From 01efb83df7ffc9ec8373561a140e34516aaf8c2f Mon Sep 17 00:00:00 2001 From: tiennm99 Date: Tue, 21 Jul 2026 18:50:30 +0700 Subject: [PATCH] fix(coin): remove stale dividend cursor fields --- cmd/server/main.go | 29 ++++ cmd/server/main_test.go | 10 ++ internal/modules/coin/startup.go | 110 ++++++++++++++++ internal/modules/coin/startup_mongo_test.go | 79 +++++++++++ internal/modules/coin/startup_test.go | 138 ++++++++++++++++++++ 5 files changed, 366 insertions(+) create mode 100644 internal/modules/coin/startup.go create mode 100644 internal/modules/coin/startup_mongo_test.go create mode 100644 internal/modules/coin/startup_test.go diff --git a/cmd/server/main.go b/cmd/server/main.go index 692d175..6cd9362 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -29,6 +29,7 @@ import ( "github.com/tiennm99/miti99bot/internal/modules/wordle" "github.com/tiennm99/miti99bot/internal/server" "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" "github.com/tiennm99/miti99bot/internal/telegram" ) @@ -94,6 +95,10 @@ func factories() map[string]modules.Factory { // container; 10s leaves headroom without hiding a wedged cluster. const mongodbInitTimeout = 10 * time.Second +// portfolioCleanupTimeout bounds the temporary one-time cleanup of legacy +// coin portfolio fields before Telegram handlers begin serving requests. +const portfolioCleanupTimeout = 2 * time.Minute + func main() { rootCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer stop() @@ -131,6 +136,18 @@ func main() { if err != nil { log.Fatal("module registry build failed", "err", err) } + if moduleLoaded(reg, coin.CollectionName) { + cleanupCtx, cancel := context.WithTimeout(rootCtx, portfolioCleanupTimeout) + err = coin.InitStore( + cleanupCtx, + provider.Collection(coin.CollectionName), + provider.Collection(systemstate.CollectionName), + ) + cancel() + if err != nil { + log.Fatal("coin storage cleanup failed", "err", err) + } + } auth := modules.Auth{BotOwnerID: cfg.BotOwnerID, AdminUserIDs: cfg.AdminUserIDs} modules.Install(b, reg, auth) log.Info("modules loaded", @@ -209,6 +226,18 @@ func main() { } } +func moduleLoaded(reg *modules.Registry, name string) bool { + if reg == nil { + return false + } + for _, module := range reg.Modules { + if module.Name == name { + return true + } + } + return false +} + // buildProvider picks the storage backend. Selection order: // 1. Explicit KV_PROVIDER env (memory|mongodb) wins. // 2. Auto-detect: MONGO_URL set → mongodb; otherwise memory. diff --git a/cmd/server/main_test.go b/cmd/server/main_test.go index 6bff6fd..c8df07d 100644 --- a/cmd/server/main_test.go +++ b/cmd/server/main_test.go @@ -77,3 +77,13 @@ func TestFactoriesIncludesExpectedModules(t *testing.T) { } } } + +func TestModuleLoaded(t *testing.T) { + reg := &modules.Registry{Modules: []modules.Module{{Name: "coin"}}} + if !moduleLoaded(reg, "coin") { + t.Fatal("coin module was not detected") + } + if moduleLoaded(reg, "stock") || moduleLoaded(nil, "coin") { + t.Fatal("moduleLoaded reported an absent module") + } +} diff --git a/internal/modules/coin/startup.go b/internal/modules/coin/startup.go new file mode 100644 index 0000000..49767bd --- /dev/null +++ b/internal/modules/coin/startup.go @@ -0,0 +1,110 @@ +package coin + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/tiennm99/miti99bot/internal/log" + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" +) + +const coinDividendCursorCleanupMarker = "migration:coin-remove-dividend-checked-at-v1" + +// dividendCursorCleanupPosition exists only to detect the stale field that an +// earlier completed migration persisted on coin assets. A pointer distinguishes +// an absent field from a present zero value so clean rows are never rewritten. +type dividendCursorCleanupPosition struct { + Quantity float64 `json:"quantity" bson:"quantity"` + Base float64 `json:"base" bson:"base"` + DividendCheckedAt *int64 `json:"dividendCheckedAt,omitempty" bson:"dividendCheckedAt,omitempty"` +} + +type dividendCursorCleanupPortfolio struct { + USD float64 `json:"usd" bson:"usd"` + Assets map[string]dividendCursorCleanupPosition `json:"assets" bson:"assets"` + Meta PortfolioMeta `json:"meta" bson:"meta"` +} + +// InitStore removes the stock-only dividend cursor accidentally persisted in +// legacy coin positions. The completion marker makes the scan a one-time task; +// each affected row uses optimistic concurrency and whole-document replacement. +func InitStore(ctx context.Context, coinColl, systemColl storage.Collection) error { + system := systemstate.New(systemColl) + marker, exists, err := system.Get(ctx, coinDividendCursorCleanupMarker) + if err != nil { + return fmt.Errorf("coin dividend cursor cleanup: read marker: %w", err) + } + if exists && marker.Status == "completed" { + return nil + } + + legacyDocs := storage.Typed[dividendCursorCleanupPortfolio](coinColl) + currentDocs := storage.Typed[Portfolio](coinColl) + keys, err := legacyDocs.List(ctx, "user:") + if err != nil { + return fmt.Errorf("coin dividend cursor cleanup: list portfolios: %w", err) + } + var cleaned int64 + for _, key := range keys { + changed, cleanupErr := cleanupCoinDividendCursor(ctx, legacyDocs, currentDocs, key) + if cleanupErr != nil { + return cleanupErr + } + if changed { + cleaned++ + } + } + + now := time.Now().UnixMilli() + if !exists { + marker = systemstate.Record{Kind: "migration", Name: "coin remove dividendCheckedAt v1"} + } + marker.Status = "completed" + marker.Count += cleaned + marker.CompletedAt = now + marker.UpdatedAt = now + if err := system.Put(ctx, coinDividendCursorCleanupMarker, marker); err != nil { + return fmt.Errorf("coin dividend cursor cleanup: write marker: %w", err) + } + log.Info("coin_dividend_cursor_cleanup_completed", "portfolios", cleaned) + return nil +} + +func cleanupCoinDividendCursor( + ctx context.Context, + legacyDocs storage.DocStore[dividendCursorCleanupPortfolio], + currentDocs storage.DocStore[Portfolio], + key string, +) (bool, error) { + for attempt := 0; attempt < portfolioUpdateAttempts; attempt++ { + legacy, version, err := legacyDocs.Get(ctx, key) + if err != nil { + return false, fmt.Errorf("coin dividend cursor cleanup: read %s: %w", key, err) + } + current, stale := legacy.currentPortfolio() + if !stale { + return false, nil + } + err = currentDocs.PutVersioned(ctx, key, version, current) + if err == nil { + return true, nil + } + if !errors.Is(err, storage.ErrConflict) { + return false, fmt.Errorf("coin dividend cursor cleanup: write %s: %w", key, err) + } + } + return false, fmt.Errorf("coin dividend cursor cleanup: write %s: %w after %d attempts", key, storage.ErrConflict, portfolioUpdateAttempts) +} + +func (p dividendCursorCleanupPortfolio) currentPortfolio() (Portfolio, bool) { + assets := make(map[string]AssetPosition, len(p.Assets)) + stale := false + for symbol, position := range p.Assets { + assets[symbol] = AssetPosition{Quantity: position.Quantity, Base: position.Base} + stale = stale || position.DividendCheckedAt != nil + } + return Portfolio{USD: p.USD, Assets: assets, Meta: p.Meta}, stale +} diff --git a/internal/modules/coin/startup_mongo_test.go b/internal/modules/coin/startup_mongo_test.go new file mode 100644 index 0000000..4aeaece --- /dev/null +++ b/internal/modules/coin/startup_mongo_test.go @@ -0,0 +1,79 @@ +package coin + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "go.mongodb.org/mongo-driver/v2/bson" + "go.mongodb.org/mongo-driver/v2/mongo" + + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" + "github.com/tiennm99/miti99bot/internal/testutil/mongotest" +) + +var coinMongoTests mongotest.Manager + +func TestMain(m *testing.M) { os.Exit(coinMongoTests.Run(m)) } + +func TestInitStoreRemovesCoinDividendCursorInMongoDB(t *testing.T) { + uri := coinMongoTests.URI(t) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + client, err := storage.NewMongoClient(ctx, uri) + if err != nil { + t.Fatal(err) + } + db := client.Database(fmt.Sprintf("miti99bot_coin_cleanup_test_%d", time.Now().UnixNano())) + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + _ = db.Drop(cleanupCtx) + _ = client.Disconnect(cleanupCtx) + }) + provider := storage.NewMongoProvider(db) + coinColl := provider.Collection(CollectionName) + legacyDocs := storage.Typed[dividendCursorCleanupPortfolio](coinColl) + if err := legacyDocs.Put(ctx, "user:7", dividendCursorCleanupPortfolio{ + USD: 5, + Assets: map[string]dividendCursorCleanupPosition{ + "BTC": {Quantity: 0.25, Base: 20_000, DividendCheckedAt: int64Pointer(123)}, + }, + Meta: PortfolioMeta{Invested: 10, CreatedAt: 1}, + }); err != nil { + t.Fatal(err) + } + rawCollection, ok := storage.MongoCollection(coinColl) + if !ok { + t.Fatal("coin collection is not MongoDB") + } + if !mongoCoinDividendCursorExists(t, ctx, rawCollection) { + t.Fatal("test setup did not persist the stale dividend cursor") + } + _, versionBefore, _ := legacyDocs.Get(ctx, "user:7") + if err := InitStore(ctx, coinColl, provider.Collection(systemstate.CollectionName)); err != nil { + t.Fatal(err) + } + if err := InitStore(ctx, coinColl, provider.Collection(systemstate.CollectionName)); err != nil { + t.Fatal(err) + } + _, versionAfter, err := storage.Typed[Portfolio](coinColl).Get(ctx, "user:7") + if err != nil || versionAfter != versionBefore+1 { + t.Fatalf("version before=%d after=%d err=%v", versionBefore, versionAfter, err) + } + if mongoCoinDividendCursorExists(t, ctx, rawCollection) { + t.Fatal("stale dividend cursor remains") + } +} + +func mongoCoinDividendCursorExists(t *testing.T, ctx context.Context, collection *mongo.Collection) bool { + t.Helper() + var raw bson.Raw + if err := collection.FindOne(ctx, bson.M{"_id": "user:7"}).Decode(&raw); err != nil { + t.Fatal(err) + } + return raw.Lookup("assets", "BTC", "dividendCheckedAt").Type != 0 +} diff --git a/internal/modules/coin/startup_test.go b/internal/modules/coin/startup_test.go new file mode 100644 index 0000000..a908b51 --- /dev/null +++ b/internal/modules/coin/startup_test.go @@ -0,0 +1,138 @@ +package coin + +import ( + "context" + "testing" + + "github.com/tiennm99/miti99bot/internal/storage" + "github.com/tiennm99/miti99bot/internal/systemstate" +) + +func int64Pointer(value int64) *int64 { return &value } + +func TestInitStoreRemovesStaleDividendCursorAndIsIdempotent(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + coinColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + legacyDocs := storage.Typed[dividendCursorCleanupPortfolio](coinColl) + currentDocs := storage.Typed[Portfolio](coinColl) + + legacy := dividendCursorCleanupPortfolio{ + USD: 75, + Assets: map[string]dividendCursorCleanupPosition{ + "BTC": {Quantity: 0.5, Base: 25_000, DividendCheckedAt: int64Pointer(123)}, + }, + Meta: PortfolioMeta{Invested: 100, CreatedAt: 1}, + } + if err := legacyDocs.Put(ctx, "user:7", legacy); err != nil { + t.Fatal(err) + } + clean := Portfolio{USD: 10, Assets: map[string]AssetPosition{"ETH": {Quantity: 1, Base: 2_000}}, Meta: PortfolioMeta{Invested: 10, CreatedAt: 2}} + if err := currentDocs.Put(ctx, "user:8", clean); err != nil { + t.Fatal(err) + } + _, cleanVersionBefore, _ := currentDocs.Get(ctx, "user:8") + + if err := InitStore(ctx, coinColl, systemColl); err != nil { + t.Fatalf("InitStore: %v", err) + } + got, migratedVersion, err := currentDocs.Get(ctx, "user:7") + if err != nil { + t.Fatal(err) + } + if got.USD != 75 || got.Meta.Invested != 100 || got.Meta.CreatedAt != 1 || got.Assets["BTC"].Quantity != 0.5 || got.Assets["BTC"].Base != 25_000 { + t.Fatalf("cleanup changed business data: %+v", got) + } + rawShape, _, err := legacyDocs.Get(ctx, "user:7") + if err != nil || rawShape.Assets["BTC"].DividendCheckedAt != nil { + t.Fatalf("stale cursor remains after cleanup: %+v err=%v", rawShape, err) + } + _, cleanVersionAfter, _ := currentDocs.Get(ctx, "user:8") + if cleanVersionAfter != cleanVersionBefore { + t.Fatalf("clean portfolio version changed: before=%d after=%d", cleanVersionBefore, cleanVersionAfter) + } + marker, exists, err := systemstate.New(systemColl).Get(ctx, coinDividendCursorCleanupMarker) + if err != nil || !exists || marker.Status != "completed" || marker.Count != 1 { + t.Fatalf("marker=%+v exists=%v err=%v", marker, exists, err) + } + + if err := InitStore(ctx, coinColl, systemColl); err != nil { + t.Fatalf("InitStore second run: %v", err) + } + _, versionAfterSecondRun, _ := currentDocs.Get(ctx, "user:7") + if versionAfterSecondRun != migratedVersion { + t.Fatalf("second run rewrote portfolio: first=%d second=%d", migratedVersion, versionAfterSecondRun) + } +} + +func TestInitStoreDoesNotRewriteUnrelatedInvalidPortfolio(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + coinColl := provider.Collection(CollectionName) + systemColl := provider.Collection(systemstate.CollectionName) + legacyDocs := storage.Typed[dividendCursorCleanupPortfolio](coinColl) + if err := legacyDocs.Put(ctx, "user:7", dividendCursorCleanupPortfolio{ + Assets: map[string]dividendCursorCleanupPosition{ + "BTC": {Quantity: -1, Base: 1}, + }, + }); err != nil { + t.Fatal(err) + } + _, versionBefore, _ := legacyDocs.Get(ctx, "user:7") + if err := InitStore(ctx, coinColl, systemColl); err != nil { + t.Fatalf("unrelated invalid portfolio blocked cleanup: %v", err) + } + _, versionAfter, _ := legacyDocs.Get(ctx, "user:7") + if versionAfter != versionBefore { + t.Fatalf("unrelated portfolio was rewritten: before=%d after=%d", versionBefore, versionAfter) + } + marker, exists, err := systemstate.New(systemColl).Get(ctx, coinDividendCursorCleanupMarker) + if err != nil || !exists || marker.Status != "completed" || marker.Count != 0 { + t.Fatalf("marker=%+v exists=%v err=%v", marker, exists, err) + } +} + +type cleanupConflictOnceStore struct { + storage.DocStore[Portfolio] + legacy storage.DocStore[dividendCursorCleanupPortfolio] + conflicted bool +} + +func (s *cleanupConflictOnceStore) PutVersioned(ctx context.Context, key string, version int64, portfolio Portfolio) error { + if !s.conflicted { + s.conflicted = true + concurrent, concurrentVersion, err := s.legacy.Get(ctx, key) + if err != nil { + return err + } + concurrent.USD += 5 + if err := s.legacy.PutVersioned(ctx, key, concurrentVersion, concurrent); err != nil { + return err + } + return storage.ErrConflict + } + return s.DocStore.PutVersioned(ctx, key, version, portfolio) +} + +func TestCleanupDividendCursorRetriesConflict(t *testing.T) { + ctx := context.Background() + provider := storage.NewMemoryProvider() + coll := provider.Collection(CollectionName) + legacyDocs := storage.Typed[dividendCursorCleanupPortfolio](coll) + if err := legacyDocs.Put(ctx, "user:7", dividendCursorCleanupPortfolio{ + Assets: map[string]dividendCursorCleanupPosition{"BTC": {Quantity: 1, Base: 10, DividendCheckedAt: int64Pointer(1)}}, + Meta: PortfolioMeta{CreatedAt: 1}, + }); err != nil { + t.Fatal(err) + } + current := &cleanupConflictOnceStore{DocStore: storage.Typed[Portfolio](coll), legacy: legacyDocs} + changed, err := cleanupCoinDividendCursor(ctx, legacyDocs, current, "user:7") + if err != nil || !changed || !current.conflicted { + t.Fatalf("changed=%v conflicted=%v err=%v", changed, current.conflicted, err) + } + shape, _, err := legacyDocs.Get(ctx, "user:7") + if err != nil || shape.Assets["BTC"].DividendCheckedAt != nil || shape.USD != 5 { + t.Fatalf("cleanup lost concurrent update or left cursor: %+v err=%v", shape, err) + } +}