优化批量账户事件的调度重建范围

2026-09-21 14:19:18 +08:00 · 2026-07-16 00:11:17 +08:00
parent eb2b8632de
commit a8a8a58f65
2 changed files with 279 additions and 6 deletions
@@ -0,0 +1,211 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/stretchr/testify/require"
+)
+
+type bulkEventAccountRepo struct {
+	*batchAccountQueryRepo
+	accounts []*Account
+}
+
+func newBulkEventAccountRepo(accounts ...*Account) *bulkEventAccountRepo {
+	return &bulkEventAccountRepo{
+		batchAccountQueryRepo: newBatchAccountQueryRepo(),
+		accounts:              accounts,
+	}
+}
+
+func (r *bulkEventAccountRepo) GetByIDs(context.Context, []int64) ([]*Account, error) {
+	return append([]*Account(nil), r.accounts...), nil
+}
+
+type bulkEventSnapshotCache struct {
+	*batchSnapshotCache
+
+	accountMu        sync.Mutex
+	setAccountIDs    []int64
+	deleteAccountIDs []int64
+}
+
+func newBulkEventSnapshotCache() *bulkEventSnapshotCache {
+	return &bulkEventSnapshotCache{batchSnapshotCache: newBatchSnapshotCache()}
+}
+
+func (c *bulkEventSnapshotCache) SetAccount(_ context.Context, account *Account) error {
+	c.accountMu.Lock()
+	defer c.accountMu.Unlock()
+	c.setAccountIDs = append(c.setAccountIDs, account.ID)
+	return nil
+}
+
+func (c *bulkEventSnapshotCache) DeleteAccount(_ context.Context, accountID int64) error {
+	c.accountMu.Lock()
+	defer c.accountMu.Unlock()
+	c.deleteAccountIDs = append(c.deleteAccountIDs, accountID)
+	return nil
+}
+
+func (c *bulkEventSnapshotCache) accountWrites() (set []int64, deleted []int64) {
+	c.accountMu.Lock()
+	defer c.accountMu.Unlock()
+	return append([]int64(nil), c.setAccountIDs...), append([]int64(nil), c.deleteAccountIDs...)
+}
+
+func (c *bulkEventSnapshotCache) capturedBuckets() []SchedulerBucket {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return append([]SchedulerBucket(nil), c.captures...)
+}
+
+func newBulkEventTestService(cache SchedulerCache, accounts AccountRepository) *SchedulerSnapshotService {
+	return NewSchedulerSnapshotService(cache, nil, accounts, nil, &config.Config{RunMode: config.RunModeStandard})
+}
+
+func bulkEventPayload(accountIDs []int64, groupIDs []int64) map[string]any {
+	accountValues := make([]any, 0, len(accountIDs))
+	for _, id := range accountIDs {
+		accountValues = append(accountValues, id)
+	}
+	groupValues := make([]any, 0, len(groupIDs))
+	for _, id := range groupIDs {
+		groupValues = append(groupValues, id)
+	}
+	return map[string]any{
+		"account_ids": accountValues,
+		"group_ids":   groupValues,
+	}
+}
+
+func schedulerBucketsForTest(groupIDs []int64, platforms ...string) []SchedulerBucket {
+	buckets := make([]SchedulerBucket, 0, len(groupIDs)*len(platforms)*3)
+	for _, platform := range platforms {
+		for _, groupID := range groupIDs {
+			buckets = append(buckets,
+				SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeSingle},
+				SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeForced},
+			)
+			if platform == PlatformAnthropic || platform == PlatformGemini {
+				buckets = append(buckets, SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeMixed})
+			}
+		}
+	}
+	return buckets
+}
+
+func TestSchedulerBulkAccountEventScopesOpenAIRebuildToFreshPlatform(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(&Account{ID: 1, Platform: PlatformOpenAI, GroupIDs: []int64{12}})
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{1}, []int64{11}), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	require.ElementsMatch(t, schedulerBucketsForTest([]int64{11, 12}, PlatformOpenAI), cache.capturedBuckets())
+	set, deleted := cache.accountWrites()
+	require.Equal(t, []int64{1}, set)
+	require.Empty(t, deleted)
+}
+
+func TestSchedulerBulkAccountEventRebuildsOpenAIUngroupedBucket(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(&Account{ID: 6, Platform: PlatformOpenAI})
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{6}, nil), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	require.ElementsMatch(t, schedulerBucketsForTest([]int64{0}, PlatformOpenAI), cache.capturedBuckets())
+}
+
+func TestSchedulerBulkAccountEventKeepsGroupedAndUngroupedBuckets(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(
+		&Account{ID: 7, Platform: PlatformOpenAI, GroupIDs: []int64{51}},
+		&Account{ID: 8, Platform: PlatformOpenAI},
+	)
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{7, 8}, nil), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	require.ElementsMatch(t, schedulerBucketsForTest([]int64{0, 51}, PlatformOpenAI), cache.capturedBuckets())
+}
+
+func TestSchedulerBulkAccountEventDoesNotCrossCurrentGroupsBetweenPlatforms(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(
+		&Account{ID: 9, Platform: PlatformOpenAI, GroupIDs: []int64{61}},
+		&Account{ID: 10, Platform: PlatformGrok, GroupIDs: []int64{62}},
+	)
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{9, 10}, []int64{63}), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	want := append(
+		schedulerBucketsForTest([]int64{61, 63}, PlatformOpenAI),
+		schedulerBucketsForTest([]int64{62, 63}, PlatformGrok)...,
+	)
+	require.ElementsMatch(t, want, cache.capturedBuckets())
+}
+
+func TestSchedulerBulkAccountEventUsesGroupZeroInSimpleMode(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(&Account{ID: 11, Platform: PlatformOpenAI, GroupIDs: []int64{71}})
+	svc := NewSchedulerSnapshotService(cache, nil, repo, nil, &config.Config{RunMode: config.RunModeSimple})
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{11}, []int64{72}), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	require.ElementsMatch(t, schedulerBucketsForTest([]int64{0}, PlatformOpenAI), cache.capturedBuckets())
+}
+
+func TestSchedulerBulkAccountEventConservativelyExpandsAntigravityPlatforms(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	// fresh 值可能已经关闭 mixed_scheduling，兼容平台仍要重建以清理旧快照。
+	repo := newBulkEventAccountRepo(&Account{ID: 2, Platform: PlatformAntigravity, GroupIDs: []int64{22}})
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{2}, []int64{21}), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	require.ElementsMatch(t,
+		schedulerBucketsForTest([]int64{21, 22}, PlatformAnthropic, PlatformGemini, PlatformAntigravity),
+		cache.capturedBuckets(),
+	)
+}
+
+func TestSchedulerBulkAccountEventMissingAccountFallsBackToAllPlatforms(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(&Account{ID: 3, Platform: PlatformOpenAI, GroupIDs: []int64{32}})
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{3, 4}, []int64{31}), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	platforms := schedulerSnapshotPlatforms()
+	require.ElementsMatch(t, schedulerBucketsForTest([]int64{31, 32}, platforms[:]...), cache.capturedBuckets())
+	set, deleted := cache.accountWrites()
+	require.Equal(t, []int64{3}, set)
+	require.Equal(t, []int64{4}, deleted)
+}
+
+func TestSchedulerBulkAccountEventUnknownPlatformFallsBackToAllPlatforms(t *testing.T) {
+	cache := newBulkEventSnapshotCache()
+	repo := newBulkEventAccountRepo(&Account{ID: 5, GroupIDs: []int64{42}})
+	svc := newBulkEventTestService(cache, repo)
+
+	err := svc.handleBulkAccountEvent(context.Background(), bulkEventPayload([]int64{5}, []int64{41}), make(map[batchSeenKey]struct{}))
+
+	require.NoError(t, err)
+	platforms := schedulerSnapshotPlatforms()
+	require.ElementsMatch(t, schedulerBucketsForTest([]int64{41, 42}, platforms[:]...), cache.capturedBuckets())
+}
@@ -529,11 +529,13 @@ func (s *SchedulerSnapshotService) handleBulkAccountEvent(ctx context.Context, p
 		}
 	}

-	if s.cache != nil {
-		for _, id := range ids {
-			if _, ok := found[id]; ok {
-				continue
-			}
+	allAccountsFound := true
+	for _, id := range ids {
+		if _, ok := found[id]; ok {
+			continue
+		}
+		allAccountsFound = false
+		if s.cache != nil {
 			if err := s.cache.DeleteAccount(ctx, id); err != nil {
 				return err
 			}
@@ -544,7 +546,67 @@ func (s *SchedulerSnapshotService) handleBulkAccountEvent(ctx context.Context, p
 	for gid := range rebuildGroupSet {
 		rebuildGroupIDs = append(rebuildGroupIDs, gid)
 	}
-	return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen)
+
+	// 缺失账户无法确定原平台，保留五平台重建以避免遗留旧快照。
+	if !allAccountsFound {
+		return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen)
+	}
+
+	platformGroupSets := make(map[string]map[int64]struct{}, len(accounts))
+	addPlatformGroups := func(platform string, groupIDs []int64) {
+		groupSet := platformGroupSets[platform]
+		if groupSet == nil {
+			groupSet = make(map[int64]struct{}, len(groupIDs))
+			platformGroupSets[platform] = groupSet
+		}
+		for _, groupID := range groupIDs {
+			groupSet[groupID] = struct{}{}
+		}
+	}
+	for _, account := range accounts {
+		if account == nil || account.ID <= 0 {
+			continue
+		}
+		accountGroupIDs := s.normalizeGroupIDs(account.GroupIDs)
+		switch account.Platform {
+		case PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformGrok:
+			addPlatformGroups(account.Platform, accountGroupIDs)
+		case PlatformAntigravity:
+			// 批量更新可能刚关闭 mixed_scheduling，仍需清理两个兼容平台的旧快照。
+			addPlatformGroups(PlatformAntigravity, accountGroupIDs)
+			addPlatformGroups(PlatformAnthropic, accountGroupIDs)
+			addPlatformGroups(PlatformGemini, accountGroupIDs)
+		default:
+			return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen)
+		}
+	}
+
+	// payload 携带更新前的组；只扩散到本事件实际涉及的平台，避免平台间交叉重建。
+	if len(preloadGroupIDs) > 0 {
+		preloadGroupIDs = s.normalizeGroupIDs(preloadGroupIDs)
+		for platform := range platformGroupSets {
+			addPlatformGroups(platform, preloadGroupIDs)
+		}
+	}
+
+	bucketCapacity := 0
+	for _, groupSet := range platformGroupSets {
+		bucketCapacity += len(groupSet) * 3
+	}
+	buckets := make([]SchedulerBucket, 0, bucketCapacity)
+	for _, platform := range schedulerSnapshotPlatforms() {
+		groupSet, ok := platformGroupSets[platform]
+		if !ok {
+			continue
+		}
+		platformGroupIDs := make([]int64, 0, len(groupSet))
+		for groupID := range groupSet {
+			platformGroupIDs = append(platformGroupIDs, groupID)
+		}
+		sort.Slice(platformGroupIDs, func(i, j int) bool { return platformGroupIDs[i] < platformGroupIDs[j] })
+		buckets = append(buckets, s.bucketsForPlatform(platform, platformGroupIDs, seen)...)
+	}
+	return s.rebuildBuckets(ctx, buckets, "account_bulk_change")
 }

 func (s *SchedulerSnapshotService) handleAccountEvent(ctx context.Context, accountID *int64, payload map[string]any, seen map[batchSeenKey]struct{}) error {