feat(stock): add ratio dividend commands

internal/modules/stats/startup.go

@@ -2,32 +2,253 @@ package stats
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"math"
+	"strconv"
+	"strings"
+	"time"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 
 	"github.com/tiennm99/miti99bot/internal/storage"
+	"github.com/tiennm99/miti99bot/internal/systemstate"
 )
 
 const (
 	statsCommandUsersIndexName   = "stats_cmd_n_user"
 	statsUserCommandsIndexName   = "stats_uid_n_cmd"
 	statsUsernameLookupIndexName = "stats_user_uid"
+
+	dividendStatsMigrationKey    = "migration:stock-dividend-command-stats-v1"
+	dividendStatsRowMarkerPrefix = dividendStatsMigrationKey + ":row:"
+	migrationStatusPrepared      = "prepared"
+	migrationStatusCompleted     = "completed"
 )
 
+type commandRename struct {
+	from string
+	to   string
+}
+
+var dividendCommandRenames = []commandRename{
+	{from: "stock_dividend", to: "stock_cash_dividend"},
+	{from: "stock_bonus", to: "stock_share_dividend"},
+}
+
 // InitStore performs stats collection startup maintenance. It is safe to call
-// every boot: MongoDB indexes are created idempotently.
-func InitStore(ctx context.Context, statsColl storage.Collection) error {
+// every boot: MongoDB indexes and the guarded stats migration are idempotent.
+func InitStore(ctx context.Context, statsColl, systemColl storage.Collection) error {
 	if mongoColl, ok := storage.MongoCollection(statsColl); ok {
 		if err := ensureUsageIndexes(ctx, mongoColl); err != nil {
 			return err
 		}
 	}
+	if err := migrateDividendCommandStats(
+		ctx,
+		storage.Typed[usageEntry](statsColl),
+		storage.Typed[systemstate.Record](systemColl),
+	); err != nil {
+		return fmt.Errorf("stats dividend command migration: %w", err)
+	}
 	return nil
 }
 
+func migrateDividendCommandStats(
+	ctx context.Context,
+	statsDocs storage.DocStore[usageEntry],
+	systemDocs storage.DocStore[systemstate.Record],
+) error {
+	global, exists, err := getSystemRecord(ctx, systemDocs, dividendStatsMigrationKey)
+	if err != nil {
+		return fmt.Errorf("read global marker: %w", err)
+	}
+	if exists && global.Status == migrationStatusCompleted {
+		return nil
+	}
+
+	for _, rename := range dividendCommandRenames {
+		keys, err := statsDocs.List(ctx, rename.from)
+		if err != nil {
+			return fmt.Errorf("list %s rows: %w", rename.from, err)
+		}
+		for _, key := range keys {
+			userID, ok := usageUserIDForCommandKey(key, rename.from)
+			if !ok {
+				continue
+			}
+			if err := migrateDividendStatsRow(ctx, statsDocs, systemDocs, rename, userID); err != nil {
+				return err
+			}
+		}
+	}
+
+	markerKeys, err := systemDocs.List(ctx, dividendStatsRowMarkerPrefix)
+	if err != nil {
+		return fmt.Errorf("list row checkpoints: %w", err)
+	}
+	for _, markerKey := range markerKeys {
+		rename, userID, ok := parseDividendStatsRowMarkerKey(markerKey)
+		if !ok {
+			continue
+		}
+		marker, exists, err := getSystemRecord(ctx, systemDocs, markerKey)
+		if err != nil {
+			return fmt.Errorf("read row checkpoint %s: %w", markerKey, err)
+		}
+		if !exists || marker.Status != migrationStatusPrepared {
+			continue
+		}
+		if err := migrateDividendStatsRow(ctx, statsDocs, systemDocs, rename, userID); err != nil {
+			return err
+		}
+	}
+
+	now := time.Now().UnixMilli()
+	return systemDocs.Put(ctx, dividendStatsMigrationKey, systemstate.Record{
+		Kind:        "migration",
+		Name:        "stock dividend command stats v1",
+		Status:      migrationStatusCompleted,
+		CompletedAt: now,
+		UpdatedAt:   now,
+	})
+}
+
+func migrateDividendStatsRow(
+	ctx context.Context,
+	statsDocs storage.DocStore[usageEntry],
+	systemDocs storage.DocStore[systemstate.Record],
+	rename commandRename,
+	userID int64,
+) error {
+	markerKey := dividendStatsRowMarkerKey(rename.from, userID)
+	marker, markerExists, err := getSystemRecord(ctx, systemDocs, markerKey)
+	if err != nil {
+		return fmt.Errorf("read %s checkpoint for user %d: %w", rename.from, userID, err)
+	}
+	if markerExists && marker.Status == migrationStatusCompleted {
+		return nil
+	}
+
+	sourceKey := usageKey(rename.from, userID)
+	targetKey := usageKey(rename.to, userID)
+	source, sourceExists, err := getUsageEntry(ctx, statsDocs, sourceKey)
+	if err != nil {
+		return fmt.Errorf("read source %s: %w", sourceKey, err)
+	}
+	target, targetExists, err := getUsageEntry(ctx, statsDocs, targetKey)
+	if err != nil {
+		return fmt.Errorf("read target %s: %w", targetKey, err)
+	}
+
+	if !markerExists {
+		if !sourceExists {
+			return nil
+		}
+		if source.N < 0 || target.N < 0 || source.N > math.MaxInt64-target.N {
+			return fmt.Errorf("count overflow for %s user %d", rename.from, userID)
+		}
+		now := time.Now().UnixMilli()
+		marker = systemstate.Record{
+			Kind:      "migration-row",
+			Name:      rename.from + " -> " + rename.to,
+			Status:    migrationStatusPrepared,
+			Count:     target.N + source.N,
+			UpdatedAt: now,
+		}
+		if err := systemDocs.Put(ctx, markerKey, marker); err != nil {
+			return fmt.Errorf("prepare %s user %d: %w", rename.from, userID, err)
+		}
+	}
+
+	if !targetExists {
+		target = usageEntry{UserID: userID}
+	}
+	target.Cmd = rename.to
+	target.UserID = userID
+	target.N = marker.Count
+	target.Deleted = false
+	if sourceExists && source.Username != "" {
+		target.Username = source.Username
+	}
+	if userID == 0 {
+		target.Username = ""
+	}
+	if err := statsDocs.Put(ctx, targetKey, target); err != nil {
+		return fmt.Errorf("write target %s: %w", targetKey, err)
+	}
+	if sourceExists {
+		if err := statsDocs.Delete(ctx, sourceKey); err != nil {
+			return fmt.Errorf("delete source %s: %w", sourceKey, err)
+		}
+	}
+
+	now := time.Now().UnixMilli()
+	marker.Status = migrationStatusCompleted
+	marker.CompletedAt = now
+	marker.UpdatedAt = now
+	if err := systemDocs.Put(ctx, markerKey, marker); err != nil {
+		return fmt.Errorf("complete %s user %d: %w", rename.from, userID, err)
+	}
+	return nil
+}
+
+func getUsageEntry(ctx context.Context, docs storage.DocStore[usageEntry], key string) (usageEntry, bool, error) {
+	entry, _, err := docs.Get(ctx, key)
+	if errors.Is(err, storage.ErrNotFound) {
+		return usageEntry{}, false, nil
+	}
+	return entry, err == nil, err
+}
+
+func getSystemRecord(ctx context.Context, docs storage.DocStore[systemstate.Record], key string) (systemstate.Record, bool, error) {
+	record, _, err := docs.Get(ctx, key)
+	if errors.Is(err, storage.ErrNotFound) {
+		return systemstate.Record{}, false, nil
+	}
+	return record, err == nil, err
+}
+
+func usageUserIDForCommandKey(key, command string) (int64, bool) {
+	if key == command {
+		return 0, true
+	}
+	prefix := command + ":"
+	if !strings.HasPrefix(key, prefix) {
+		return 0, false
+	}
+	userID, err := strconv.ParseInt(strings.TrimPrefix(key, prefix), 10, 64)
+	return userID, err == nil && userID != 0
+}
+
+func dividendStatsRowMarkerKey(sourceCommand string, userID int64) string {
+	return dividendStatsRowMarkerPrefix + sourceCommand + ":" + strconv.FormatInt(userID, 10)
+}
+
+func parseDividendStatsRowMarkerKey(key string) (commandRename, int64, bool) {
+	remainder := strings.TrimPrefix(key, dividendStatsRowMarkerPrefix)
+	if remainder == key {
+		return commandRename{}, 0, false
+	}
+	separator := strings.LastIndexByte(remainder, ':')
+	if separator < 1 {
+		return commandRename{}, 0, false
+	}
+	sourceCommand := remainder[:separator]
+	userID, err := strconv.ParseInt(remainder[separator+1:], 10, 64)
+	if err != nil {
+		return commandRename{}, 0, false
+	}
+	for _, rename := range dividendCommandRenames {
+		if rename.from == sourceCommand {
+			return rename, userID, true
+		}
+	}
+	return commandRename{}, 0, false
+}
+
 func ensureUsageIndexes(ctx context.Context, coll *mongo.Collection) error {
 	models := []mongo.IndexModel{
 		{

internal/modules/stats/startup_mongo_test.go

@@ -11,17 +11,17 @@ import (
 )
 
 func TestInitStore_MongoCreatesIndexes(t *testing.T) {
-	ctx, statsColl := setupMongoStatsTest(t)
+	ctx, statsColl, systemColl := setupMongoStatsTest(t)
 
 	rawStatsColl, ok := storage.MongoCollection(statsColl)
 	if !ok {
 		t.Fatal("stats collection is not Mongo-backed")
 	}
 
-	if err := InitStore(ctx, statsColl); err != nil {
+	if err := InitStore(ctx, statsColl, systemColl); err != nil {
 		t.Fatalf("InitStore: %v", err)
 	}
-	if err := InitStore(ctx, statsColl); err != nil {
+	if err := InitStore(ctx, statsColl, systemColl); err != nil {
 		t.Fatalf("InitStore second run: %v", err)
 	}
 
@@ -51,7 +51,25 @@ func TestInitStore_MongoCreatesIndexes(t *testing.T) {
 	}
 }
 
-func setupMongoStatsTest(t *testing.T) (context.Context, storage.Collection) {
+func TestInitStore_MongoMigratesDividendStatsIdempotently(t *testing.T) {
+	ctx, statsColl, systemColl := setupMongoStatsTest(t)
+	docs := storage.Typed[usageEntry](statsColl)
+	seedUsageEntries(t, docs, map[string]usageEntry{
+		usageKey("stock_bonus", 0):          {Cmd: "stock_bonus", N: 4},
+		usageKey("stock_bonus", 7):          {Cmd: "stock_bonus", UserID: 7, Username: "alice", N: 5},
+		usageKey("stock_share_dividend", 7): {Cmd: "stock_share_dividend", UserID: 7, Username: "alice", N: 2},
+	})
+	if err := InitStore(ctx, statsColl, systemColl); err != nil {
+		t.Fatalf("InitStore: %v", err)
+	}
+	if err := InitStore(ctx, statsColl, systemColl); err != nil {
+		t.Fatalf("InitStore second run: %v", err)
+	}
+	assertUsageEntry(t, docs, usageKey("stock_share_dividend", 0), 4, "")
+	assertUsageEntry(t, docs, usageKey("stock_share_dividend", 7), 7, "alice")
+}
+
+func setupMongoStatsTest(t *testing.T) (context.Context, storage.Collection, storage.Collection) {
 	t.Helper()
 
 	uri := os.Getenv("MONGODB_TEST_URL")
@@ -76,5 +94,5 @@ func setupMongoStatsTest(t *testing.T) (context.Context, storage.Collection) {
 	})
 
 	provider := storage.NewMongoProvider(db)
-	return ctx, provider.Collection("stats")
+	return ctx, provider.Collection("stats"), provider.Collection("system")
 }

internal/modules/stats/startup_test.go

@@ -0,0 +1,186 @@
+package stats
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/tiennm99/miti99bot/internal/storage"
+	"github.com/tiennm99/miti99bot/internal/systemstate"
+)
+
+func TestInitStoreMigratesDividendCommandStats(t *testing.T) {
+	ctx := context.Background()
+	provider := storage.NewMemoryProvider()
+	statsColl := provider.Collection("stats")
+	systemColl := provider.Collection(systemstate.CollectionName)
+	statsDocs := storage.Typed[usageEntry](statsColl)
+	systemDocs := storage.Typed[systemstate.Record](systemColl)
+
+	seedUsageEntries(t, statsDocs, map[string]usageEntry{
+		usageKey("stock_dividend", 0):       {Cmd: "stock_dividend", N: 10},
+		usageKey("stock_cash_dividend", 0):  {Cmd: "stock_cash_dividend", N: 3},
+		usageKey("stock_dividend", 7):       {Cmd: "stock_dividend", UserID: 7, Username: "alice", N: 4},
+		usageKey("stock_cash_dividend", 7):  {Cmd: "stock_cash_dividend", UserID: 7, Username: "old-alice", N: 2},
+		usageKey("stock_bonus", 0):          {Cmd: "stock_bonus", N: 8},
+		usageKey("stock_bonus", 9):          {Cmd: "stock_bonus", UserID: 9, Username: "bob", N: 5},
+		usageKey("stock_share_dividend", 9): {Cmd: "stock_share_dividend", UserID: 9, Username: "previous", N: 1},
+	})
+
+	if err := InitStore(ctx, statsColl, systemColl); err != nil {
+		t.Fatalf("InitStore: %v", err)
+	}
+	assertUsageEntry(t, statsDocs, usageKey("stock_cash_dividend", 0), 13, "")
+	assertUsageEntry(t, statsDocs, usageKey("stock_cash_dividend", 7), 6, "alice")
+	assertUsageEntry(t, statsDocs, usageKey("stock_share_dividend", 0), 8, "")
+	assertUsageEntry(t, statsDocs, usageKey("stock_share_dividend", 9), 6, "bob")
+	for _, key := range []string{
+		usageKey("stock_dividend", 0), usageKey("stock_dividend", 7),
+		usageKey("stock_bonus", 0), usageKey("stock_bonus", 9),
+	} {
+		if _, _, err := statsDocs.Get(ctx, key); !errors.Is(err, storage.ErrNotFound) {
+			t.Fatalf("source %s still exists or get failed: %v", key, err)
+		}
+	}
+	global, _, err := systemDocs.Get(ctx, dividendStatsMigrationKey)
+	if err != nil || global.Status != migrationStatusCompleted {
+		t.Fatalf("global marker = %+v, %v", global, err)
+	}
+	markerKeys, err := systemDocs.List(ctx, dividendStatsRowMarkerPrefix)
+	if err != nil || len(markerKeys) != 4 {
+		t.Fatalf("row markers = %v, %v", markerKeys, err)
+	}
+
+	if err := InitStore(ctx, statsColl, systemColl); err != nil {
+		t.Fatalf("InitStore second run: %v", err)
+	}
+	assertUsageEntry(t, statsDocs, usageKey("stock_cash_dividend", 0), 13, "")
+	assertUsageEntry(t, statsDocs, usageKey("stock_share_dividend", 9), 6, "bob")
+}
+
+func TestDividendStatsMigrationResumesPreparedCheckpointWithoutSource(t *testing.T) {
+	ctx := context.Background()
+	provider := storage.NewMemoryProvider()
+	statsDocs := storage.Typed[usageEntry](provider.Collection("stats"))
+	systemDocs := storage.Typed[systemstate.Record](provider.Collection(systemstate.CollectionName))
+	markerKey := dividendStatsRowMarkerKey("stock_dividend", 7)
+
+	if err := statsDocs.Put(ctx, usageKey("stock_cash_dividend", 7), usageEntry{
+		Cmd: "stock_cash_dividend", UserID: 7, Username: "alice", N: 12,
+	}); err != nil {
+		t.Fatalf("seed target: %v", err)
+	}
+	if err := systemDocs.Put(ctx, markerKey, systemstate.Record{
+		Kind: "migration-row", Name: "stock_dividend -> stock_cash_dividend", Status: migrationStatusPrepared, Count: 12,
+	}); err != nil {
+		t.Fatalf("seed marker: %v", err)
+	}
+
+	if err := migrateDividendCommandStats(ctx, statsDocs, systemDocs); err != nil {
+		t.Fatalf("migrateDividendCommandStats: %v", err)
+	}
+	assertUsageEntry(t, statsDocs, usageKey("stock_cash_dividend", 7), 12, "alice")
+	marker, _, err := systemDocs.Get(ctx, markerKey)
+	if err != nil || marker.Status != migrationStatusCompleted {
+		t.Fatalf("row marker = %+v, %v", marker, err)
+	}
+}
+
+type faultDocStore[T any] struct {
+	storage.DocStore[T]
+	fail func(op, id string, value T) error
+}
+
+func (s *faultDocStore[T]) Put(ctx context.Context, id string, value T) error {
+	if s.fail != nil {
+		if err := s.fail("put", id, value); err != nil {
+			return err
+		}
+	}
+	return s.DocStore.Put(ctx, id, value)
+}
+
+func (s *faultDocStore[T]) Delete(ctx context.Context, id string) error {
+	var zero T
+	if s.fail != nil {
+		if err := s.fail("delete", id, zero); err != nil {
+			return err
+		}
+	}
+	return s.DocStore.Delete(ctx, id)
+}
+
+func TestDividendStatsMigrationRetriesEveryWriteBoundaryWithoutDuplication(t *testing.T) {
+	for _, boundary := range []string{"prepare", "target", "delete", "complete"} {
+		t.Run(boundary, func(t *testing.T) {
+			ctx := context.Background()
+			provider := storage.NewMemoryProvider()
+			baseStats := storage.Typed[usageEntry](provider.Collection("stats"))
+			baseSystem := storage.Typed[systemstate.Record](provider.Collection(systemstate.CollectionName))
+			seedUsageEntries(t, baseStats, map[string]usageEntry{
+				usageKey("stock_dividend", 7):      {Cmd: "stock_dividend", UserID: 7, Username: "alice", N: 5},
+				usageKey("stock_cash_dividend", 7): {Cmd: "stock_cash_dividend", UserID: 7, Username: "alice", N: 7},
+			})
+
+			failed := false
+			statsDocs := &faultDocStore[usageEntry]{DocStore: baseStats}
+			systemDocs := &faultDocStore[systemstate.Record]{DocStore: baseSystem}
+			forced := errors.New("forced boundary failure")
+			statsDocs.fail = func(op, id string, _ usageEntry) error {
+				if failed {
+					return nil
+				}
+				if boundary == "target" && op == "put" && id == usageKey("stock_cash_dividend", 7) ||
+					boundary == "delete" && op == "delete" && id == usageKey("stock_dividend", 7) {
+					failed = true
+					return forced
+				}
+				return nil
+			}
+			systemDocs.fail = func(op, id string, record systemstate.Record) error {
+				if failed || op != "put" || id != dividendStatsRowMarkerKey("stock_dividend", 7) {
+					return nil
+				}
+				if boundary == "prepare" && record.Status == migrationStatusPrepared ||
+					boundary == "complete" && record.Status == migrationStatusCompleted {
+					failed = true
+					return forced
+				}
+				return nil
+			}
+
+			if err := migrateDividendCommandStats(ctx, statsDocs, systemDocs); err == nil {
+				t.Fatal("migration unexpectedly succeeded before injected failure")
+			}
+			statsDocs.fail = nil
+			systemDocs.fail = nil
+			if err := migrateDividendCommandStats(ctx, statsDocs, systemDocs); err != nil {
+				t.Fatalf("retry migration: %v", err)
+			}
+			assertUsageEntry(t, baseStats, usageKey("stock_cash_dividend", 7), 12, "alice")
+			if _, _, err := baseStats.Get(ctx, usageKey("stock_dividend", 7)); !errors.Is(err, storage.ErrNotFound) {
+				t.Fatalf("source remains after retry: %v", err)
+			}
+		})
+	}
+}
+
+func seedUsageEntries(t *testing.T, docs storage.DocStore[usageEntry], entries map[string]usageEntry) {
+	t.Helper()
+	for key, entry := range entries {
+		if err := docs.Put(context.Background(), key, entry); err != nil {
+			t.Fatalf("seed %s: %v", key, err)
+		}
+	}
+}
+
+func assertUsageEntry(t *testing.T, docs storage.DocStore[usageEntry], key string, wantN int64, wantUsername string) {
+	t.Helper()
+	entry, _, err := docs.Get(context.Background(), key)
+	if err != nil {
+		t.Fatalf("get %s: %v", key, err)
+	}
+	if entry.N != wantN || entry.Username != wantUsername || entry.Deleted {
+		t.Fatalf("entry %s = %+v, want n=%d username=%q deleted=false", key, entry, wantN, wantUsername)
+	}
+}