mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-21 04:04:34 +00:00
清理全量重建中的过期调度桶
This commit is contained in:
@@ -0,0 +1,666 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type fullRebuildLifecycleCache struct {
|
||||
*groupLifecycleTestCache
|
||||
|
||||
mu sync.Mutex
|
||||
captureAttempts []SchedulerBucket
|
||||
captureErrors map[string]error
|
||||
lockBusyOnce map[string]bool
|
||||
watermark int64
|
||||
watermarkWrites []int64
|
||||
}
|
||||
|
||||
func newFullRebuildLifecycleCache(buckets ...SchedulerBucket) *fullRebuildLifecycleCache {
|
||||
return &fullRebuildLifecycleCache{
|
||||
groupLifecycleTestCache: newGroupLifecycleTestCache(buckets...),
|
||||
captureErrors: make(map[string]error),
|
||||
lockBusyOnce: make(map[string]bool),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) CaptureBucketWriteToken(ctx context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) {
|
||||
c.mu.Lock()
|
||||
c.captureAttempts = append(c.captureAttempts, bucket)
|
||||
err := c.captureErrors[bucket.String()]
|
||||
c.mu.Unlock()
|
||||
if err != nil {
|
||||
return SchedulerBucketWriteToken{}, err
|
||||
}
|
||||
return c.retirementRaceCache.CaptureBucketWriteToken(ctx, bucket)
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) ListBuckets(ctx context.Context) ([]SchedulerBucket, error) {
|
||||
buckets, err := c.groupLifecycleTestCache.ListBuckets(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.retirementRaceCache.mu.Lock()
|
||||
defer c.retirementRaceCache.mu.Unlock()
|
||||
registered := make([]SchedulerBucket, 0, len(buckets))
|
||||
for _, bucket := range buckets {
|
||||
if !c.retirementRaceCache.retired[bucket.String()] {
|
||||
registered = append(registered, bucket)
|
||||
}
|
||||
}
|
||||
return registered, nil
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) TryLockBucket(ctx context.Context, bucket SchedulerBucket, ttl time.Duration) (bool, error) {
|
||||
c.mu.Lock()
|
||||
busy := c.lockBusyOnce[bucket.String()]
|
||||
if busy {
|
||||
delete(c.lockBusyOnce, bucket.String())
|
||||
}
|
||||
c.mu.Unlock()
|
||||
if busy {
|
||||
return false, nil
|
||||
}
|
||||
return c.groupLifecycleTestCache.TryLockBucket(ctx, bucket, ttl)
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) GetOutboxWatermark(context.Context) (int64, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.watermark, nil
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) SetOutboxWatermark(_ context.Context, id int64) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.watermark = id
|
||||
c.watermarkWrites = append(c.watermarkWrites, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) captureAttemptCount() int {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return len(c.captureAttempts)
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) currentWatermark() int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.watermark
|
||||
}
|
||||
|
||||
func (c *fullRebuildLifecycleCache) totalSetAttempts() int {
|
||||
c.retirementRaceCache.mu.Lock()
|
||||
defer c.retirementRaceCache.mu.Unlock()
|
||||
var total int
|
||||
for _, attempts := range c.retirementRaceCache.setAttempts {
|
||||
total += attempts
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
type fullRebuildLifecycleGroupRepo struct {
|
||||
GroupRepository
|
||||
|
||||
mu sync.Mutex
|
||||
activeIDs []int64
|
||||
activeIDsErr error
|
||||
listActiveErr error
|
||||
fresh map[int64]*Group
|
||||
freshErr map[int64]error
|
||||
activeIDCalls int
|
||||
listActiveCalls int
|
||||
freshCalls []int64
|
||||
}
|
||||
|
||||
func (r *fullRebuildLifecycleGroupRepo) ListActiveIDs(context.Context) ([]int64, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.activeIDCalls++
|
||||
return append([]int64(nil), r.activeIDs...), r.activeIDsErr
|
||||
}
|
||||
|
||||
func (r *fullRebuildLifecycleGroupRepo) ListActive(context.Context) ([]Group, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.listActiveCalls++
|
||||
return nil, r.listActiveErr
|
||||
}
|
||||
|
||||
func (r *fullRebuildLifecycleGroupRepo) GetByIDLite(_ context.Context, id int64) (*Group, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.freshCalls = append(r.freshCalls, id)
|
||||
if err := r.freshErr[id]; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
group := r.fresh[id]
|
||||
if group == nil {
|
||||
return nil, ErrGroupNotFound
|
||||
}
|
||||
copyGroup := *group
|
||||
return ©Group, nil
|
||||
}
|
||||
|
||||
func (r *fullRebuildLifecycleGroupRepo) stats() (activeIDs, listActive int, fresh []int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.activeIDCalls, r.listActiveCalls, append([]int64(nil), r.freshCalls...)
|
||||
}
|
||||
|
||||
type fullRebuildFallbackGroupRepo struct {
|
||||
GroupRepository
|
||||
|
||||
mu sync.Mutex
|
||||
groups []Group
|
||||
err error
|
||||
listCalls int
|
||||
}
|
||||
|
||||
func (r *fullRebuildFallbackGroupRepo) ListActive(context.Context) ([]Group, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.listCalls++
|
||||
return append([]Group(nil), r.groups...), r.err
|
||||
}
|
||||
|
||||
type fullRebuildAccountCall struct {
|
||||
groupID int64
|
||||
platform string
|
||||
}
|
||||
|
||||
type fullRebuildAccountRepo struct {
|
||||
AccountRepository
|
||||
|
||||
mu sync.Mutex
|
||||
calls []fullRebuildAccountCall
|
||||
beforeFirst func()
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) record(groupID int64, platform string) ([]Account, error) {
|
||||
r.mu.Lock()
|
||||
r.calls = append(r.calls, fullRebuildAccountCall{groupID: groupID, platform: platform})
|
||||
beforeFirst := r.beforeFirst
|
||||
r.mu.Unlock()
|
||||
if beforeFirst != nil {
|
||||
r.once.Do(beforeFirst)
|
||||
}
|
||||
return []Account{{ID: 1, Platform: platform, Status: StatusActive, Schedulable: true}}, nil
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) ListSchedulableByGroupIDAndPlatform(_ context.Context, groupID int64, platform string) ([]Account, error) {
|
||||
return r.record(groupID, platform)
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) ListSchedulableByGroupIDAndPlatforms(_ context.Context, groupID int64, platforms []string) ([]Account, error) {
|
||||
return r.record(groupID, firstPlatform(platforms))
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) ListSchedulableUngroupedByPlatform(_ context.Context, platform string) ([]Account, error) {
|
||||
return r.record(0, platform)
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) ListSchedulableUngroupedByPlatforms(_ context.Context, platforms []string) ([]Account, error) {
|
||||
return r.record(0, firstPlatform(platforms))
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) {
|
||||
return r.record(0, platform)
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) ListSchedulableByPlatforms(_ context.Context, platforms []string) ([]Account, error) {
|
||||
return r.record(0, firstPlatform(platforms))
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) callCount() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return len(r.calls)
|
||||
}
|
||||
|
||||
func (r *fullRebuildAccountRepo) groupCallCount(groupID int64) int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var count int
|
||||
for _, call := range r.calls {
|
||||
if call.groupID == groupID {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func firstPlatform(platforms []string) string {
|
||||
if len(platforms) == 0 {
|
||||
return ""
|
||||
}
|
||||
return platforms[0]
|
||||
}
|
||||
|
||||
func newFullRebuildLifecycleService(
|
||||
cache SchedulerCache,
|
||||
outbox SchedulerOutboxRepository,
|
||||
accounts AccountRepository,
|
||||
groups GroupRepository,
|
||||
runMode string,
|
||||
) *SchedulerSnapshotService {
|
||||
return NewSchedulerSnapshotService(cache, outbox, accounts, groups, &config.Config{RunMode: runMode})
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildActiveTombstoneDoesNotBlockFollowingGroupEvent(t *testing.T) {
|
||||
const groupID int64 = 101
|
||||
canonical := schedulerBucketsForGroup(groupID)
|
||||
cache := newFullRebuildLifecycleCache()
|
||||
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[0]))
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
activeIDs: []int64{groupID},
|
||||
fresh: map[int64]*Group{
|
||||
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
|
||||
},
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
outbox := &outboxCleanupRepo{events: []SchedulerOutboxEvent{
|
||||
{ID: 1, EventType: SchedulerOutboxEventFullRebuild},
|
||||
{ID: 2, EventType: SchedulerOutboxEventGroupChanged, GroupID: ptrInt64(groupID)},
|
||||
}}
|
||||
svc := newFullRebuildLifecycleService(cache, outbox, accounts, groups, config.RunModeStandard)
|
||||
|
||||
svc.pollOutbox()
|
||||
|
||||
require.Equal(t, int64(2), cache.currentWatermark())
|
||||
cache.mu.Lock()
|
||||
require.Equal(t, []int64{2}, cache.watermarkWrites)
|
||||
cache.mu.Unlock()
|
||||
activeCalls, fallbackCalls, freshCalls := groups.stats()
|
||||
require.Equal(t, 1, activeCalls)
|
||||
require.Zero(t, fallbackCalls)
|
||||
require.Equal(t, []int64{groupID, groupID}, freshCalls)
|
||||
require.Len(t, cache.tokens(), 24, "full rebuild and the following group event must each run fresh authority")
|
||||
_, reopenHeld := cache.lifecycleMutationLeaseStates()
|
||||
require.Len(t, reopenHeld, 24)
|
||||
for _, held := range reopenHeld {
|
||||
require.True(t, held)
|
||||
}
|
||||
require.Equal(t, 36, accounts.callCount())
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildGlobalReadErrorsFailBeforeMutationOrDB(t *testing.T) {
|
||||
t.Run("list buckets", func(t *testing.T) {
|
||||
cache := newFullRebuildLifecycleCache()
|
||||
cache.listErr = errors.New("registry failed")
|
||||
groups := &fullRebuildLifecycleGroupRepo{fresh: make(map[int64]*Group), freshErr: make(map[int64]error)}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
err := svc.rebuildFullSnapshot(context.Background(), "test")
|
||||
require.ErrorIs(t, err, cache.listErr)
|
||||
activeCalls, fallbackCalls, freshCalls := groups.stats()
|
||||
require.Zero(t, activeCalls)
|
||||
require.Zero(t, fallbackCalls)
|
||||
require.Empty(t, freshCalls)
|
||||
requireFullRebuildNoMutationOrDB(t, cache, accounts)
|
||||
})
|
||||
|
||||
t.Run("list active ids without fallback", func(t *testing.T) {
|
||||
cache := newFullRebuildLifecycleCache()
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
activeIDsErr: errors.New("active ids failed"),
|
||||
listActiveErr: errors.New("fallback must not run"),
|
||||
fresh: make(map[int64]*Group),
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
err := svc.rebuildFullSnapshot(context.Background(), "test")
|
||||
require.ErrorIs(t, err, groups.activeIDsErr)
|
||||
activeCalls, fallbackCalls, freshCalls := groups.stats()
|
||||
require.Equal(t, 1, activeCalls)
|
||||
require.Zero(t, fallbackCalls)
|
||||
require.Empty(t, freshCalls)
|
||||
requireFullRebuildNoMutationOrDB(t, cache, accounts)
|
||||
})
|
||||
|
||||
t.Run("list active fallback", func(t *testing.T) {
|
||||
cache := newFullRebuildLifecycleCache()
|
||||
groups := &fullRebuildFallbackGroupRepo{err: errors.New("active groups failed")}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
err := svc.rebuildFullSnapshot(context.Background(), "test")
|
||||
require.ErrorIs(t, err, groups.err)
|
||||
groups.mu.Lock()
|
||||
require.Equal(t, 1, groups.listCalls)
|
||||
groups.mu.Unlock()
|
||||
requireFullRebuildNoMutationOrDB(t, cache, accounts)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildFreshActivePreparesEveryTokenBeforeFirstDB(t *testing.T) {
|
||||
const groupID int64 = 102
|
||||
historical := SchedulerBucket{GroupID: groupID, Platform: "legacy", Mode: "unknown"}
|
||||
cache := newFullRebuildLifecycleCache(historical)
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
fresh: map[int64]*Group{
|
||||
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
|
||||
},
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
var capturesAtFirstDB int
|
||||
accounts.beforeFirst = func() {
|
||||
capturesAtFirstDB = cache.captureAttemptCount()
|
||||
held, reopenCount := cache.leaseHeldAndTokenCount()
|
||||
require.False(t, held)
|
||||
require.Equal(t, 12, reopenCount)
|
||||
require.Equal(t, 13, capturesAtFirstDB, "C(0) and the historical bucket must be captured before DB")
|
||||
}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
|
||||
require.Equal(t, capturesAtFirstDB, cache.captureAttemptCount())
|
||||
require.Equal(t, 25, accounts.callCount())
|
||||
require.Equal(t, 13, accounts.groupCallCount(groupID))
|
||||
_, historicalPublished := cache.counts(historical)
|
||||
require.Equal(t, 1, historicalPublished)
|
||||
activeCalls, fallbackCalls, freshCalls := groups.stats()
|
||||
require.Equal(t, 1, activeCalls)
|
||||
require.Zero(t, fallbackCalls)
|
||||
require.Equal(t, []int64{groupID}, freshCalls)
|
||||
_, _, listCalls := cache.lifecycleCounts()
|
||||
require.Equal(t, 1, listCalls, "known Rg must prevent a second registry read")
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildOrdinaryCaptureErrorReturnsBeforeFirstDB(t *testing.T) {
|
||||
first := SchedulerBucket{GroupID: 0, Platform: "legacy-a", Mode: "unknown"}
|
||||
last := SchedulerBucket{GroupID: 0, Platform: "legacy-b", Mode: "unknown"}
|
||||
cache := newFullRebuildLifecycleCache(first, last)
|
||||
wantErr := errors.New("capture failed")
|
||||
cache.captureErrors[last.String()] = wantErr
|
||||
groups := &fullRebuildLifecycleGroupRepo{fresh: make(map[int64]*Group), freshErr: make(map[int64]error)}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
err := svc.rebuildFullSnapshot(context.Background(), "test")
|
||||
require.ErrorIs(t, err, wantErr)
|
||||
require.Equal(t, 14, cache.captureAttemptCount(), "all canonical and ordinary captures must be attempted before returning")
|
||||
require.Zero(t, accounts.callCount())
|
||||
require.Zero(t, cache.totalSetAttempts())
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildPreservesGroupZeroActiveHistoricalAndInvalidRegistryBuckets(t *testing.T) {
|
||||
const groupID int64 = 103
|
||||
groupZeroHistorical := SchedulerBucket{GroupID: 0, Platform: "legacy-zero", Mode: "unknown"}
|
||||
activeHistorical := SchedulerBucket{GroupID: groupID, Platform: "legacy-active", Mode: "unknown"}
|
||||
activeForced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced}
|
||||
activeMixed := SchedulerBucket{GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeMixed}
|
||||
invalidHistorical := SchedulerBucket{GroupID: -7, Platform: "legacy-invalid", Mode: "unknown"}
|
||||
cache := newFullRebuildLifecycleCache(groupZeroHistorical, activeHistorical, activeForced, activeMixed, invalidHistorical)
|
||||
groups := &fullRebuildFallbackGroupRepo{groups: []Group{{ID: groupID, Status: StatusActive}}}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
|
||||
require.Equal(t, 27, cache.captureAttemptCount())
|
||||
require.Equal(t, 27, accounts.callCount())
|
||||
groups.mu.Lock()
|
||||
require.Equal(t, 1, groups.listCalls)
|
||||
groups.mu.Unlock()
|
||||
require.Empty(t, cache.retiredBuckets())
|
||||
require.Empty(t, cache.tokens())
|
||||
for _, bucket := range []SchedulerBucket{groupZeroHistorical, activeHistorical, activeForced, activeMixed, invalidHistorical} {
|
||||
_, published := cache.counts(bucket)
|
||||
require.Equal(t, 1, published, bucket.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildActiveTombstoneFreshInactiveOrMissingFiltersAllGroupTasks(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
group *Group
|
||||
}{
|
||||
{name: "inactive", group: &Group{ID: 104, Status: StatusDisabled, Hydrated: true}},
|
||||
{name: "missing"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
const groupID int64 = 104
|
||||
canonical := schedulerBucketsForGroup(groupID)
|
||||
historical := SchedulerBucket{GroupID: groupID, Platform: "legacy", Mode: "unknown"}
|
||||
cache := newFullRebuildLifecycleCache(historical)
|
||||
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[5]))
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
activeIDs: []int64{groupID},
|
||||
fresh: make(map[int64]*Group),
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
if tc.group != nil {
|
||||
groups.fresh[groupID] = tc.group
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
|
||||
require.Zero(t, accounts.groupCallCount(groupID))
|
||||
require.Equal(t, 12, accounts.groupCallCount(0))
|
||||
require.Empty(t, cache.tokens())
|
||||
require.Equal(t, bucketStrings(append(canonical, historical)), bucketStrings(cache.retiredBuckets()))
|
||||
for _, bucket := range append(canonical, historical) {
|
||||
attempts, published := cache.counts(bucket)
|
||||
require.Zero(t, attempts, bucket.String())
|
||||
require.Zero(t, published, bucket.String())
|
||||
}
|
||||
_, _, listCalls := cache.lifecycleCounts()
|
||||
require.Equal(t, 1, listCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildStaleCandidatesAreSortedAndNeverReopenedAcrossRounds(t *testing.T) {
|
||||
historical := []SchedulerBucket{
|
||||
{GroupID: 3, Platform: "legacy", Mode: "unknown"},
|
||||
{GroupID: 1, Platform: "legacy", Mode: "unknown"},
|
||||
{GroupID: 2, Platform: "legacy", Mode: "unknown"},
|
||||
}
|
||||
cache := newFullRebuildLifecycleCache(historical...)
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
fresh: map[int64]*Group{
|
||||
1: {ID: 1, Status: StatusDisabled, Hydrated: true},
|
||||
3: {ID: 3, Status: StatusDisabled, Hydrated: true},
|
||||
},
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "first"))
|
||||
retireCallsAfterFirst := len(cache.retiredBuckets())
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "second"))
|
||||
_, _, freshCalls := groups.stats()
|
||||
require.Equal(t, []int64{1, 2, 3}, freshCalls)
|
||||
require.Equal(t, retireCallsAfterFirst, len(cache.retiredBuckets()))
|
||||
require.Empty(t, cache.tokens())
|
||||
require.Zero(t, accounts.groupCallCount(1))
|
||||
require.Zero(t, accounts.groupCallCount(2))
|
||||
require.Zero(t, accounts.groupCallCount(3))
|
||||
_, _, listCalls := cache.lifecycleCounts()
|
||||
require.Equal(t, 2, listCalls, "each round must use only its one global registry snapshot")
|
||||
for _, bucket := range historical {
|
||||
require.Contains(t, bucketStrings(cache.retiredBuckets()), bucket.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildPartialLifecycleFailureReturnsBeforeDBAndRetries(t *testing.T) {
|
||||
historical := []SchedulerBucket{
|
||||
{GroupID: 1, Platform: "legacy", Mode: "unknown"},
|
||||
{GroupID: 2, Platform: "legacy", Mode: "unknown"},
|
||||
{GroupID: 3, Platform: "legacy", Mode: "unknown"},
|
||||
}
|
||||
wantErr := errors.New("fresh group query failed")
|
||||
cache := newFullRebuildLifecycleCache(historical...)
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
fresh: map[int64]*Group{
|
||||
1: {ID: 1, Status: StatusDisabled, Hydrated: true},
|
||||
3: {ID: 3, Status: StatusDisabled, Hydrated: true},
|
||||
},
|
||||
freshErr: map[int64]error{2: wantErr},
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
err := svc.triggerFullRebuild("first")
|
||||
require.ErrorIs(t, err, wantErr)
|
||||
_, _, freshCalls := groups.stats()
|
||||
require.Equal(t, []int64{1, 2}, freshCalls)
|
||||
require.Zero(t, accounts.callCount())
|
||||
require.Zero(t, cache.totalSetAttempts())
|
||||
require.Equal(t, 13, len(cache.retiredBuckets()))
|
||||
|
||||
groups.mu.Lock()
|
||||
delete(groups.freshErr, 2)
|
||||
groups.fresh[2] = &Group{ID: 2, Status: StatusDisabled, Hydrated: true}
|
||||
groups.mu.Unlock()
|
||||
require.NoError(t, svc.triggerFullRebuild("retry"))
|
||||
_, _, freshCalls = groups.stats()
|
||||
require.Equal(t, []int64{1, 2, 2, 3}, freshCalls)
|
||||
require.Equal(t, 39, len(cache.retiredBuckets()))
|
||||
require.Equal(t, 12, accounts.callCount())
|
||||
require.Empty(t, cache.tokens())
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildActiveTombstoneLazyRecoveryDiscardsPartialCaptureTasks(t *testing.T) {
|
||||
const groupID int64 = 105
|
||||
canonical := schedulerBucketsForGroup(groupID)
|
||||
historical := SchedulerBucket{GroupID: groupID, Platform: "legacy", Mode: "unknown"}
|
||||
cache := newFullRebuildLifecycleCache(canonical[0], canonical[4], historical)
|
||||
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[5]))
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
activeIDs: []int64{groupID},
|
||||
fresh: map[int64]*Group{
|
||||
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
|
||||
},
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
var capturesAtFirstDB int
|
||||
accounts.beforeFirst = func() {
|
||||
capturesAtFirstDB = cache.captureAttemptCount()
|
||||
held, reopenCount := cache.leaseHeldAndTokenCount()
|
||||
require.False(t, held)
|
||||
require.Equal(t, 12, reopenCount)
|
||||
require.Equal(t, 19, capturesAtFirstDB)
|
||||
}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
|
||||
require.Equal(t, capturesAtFirstDB, cache.captureAttemptCount())
|
||||
require.Equal(t, 25, accounts.callCount())
|
||||
for _, bucket := range canonical {
|
||||
attempts, published := cache.counts(bucket)
|
||||
require.Equal(t, 1, attempts, "discarded pre-recovery tokens must never publish: %s", bucket.String())
|
||||
require.Equal(t, 1, published, bucket.String())
|
||||
}
|
||||
_, published := cache.counts(historical)
|
||||
require.Equal(t, 1, published)
|
||||
_, _, listCalls := cache.lifecycleCounts()
|
||||
require.Equal(t, 1, listCalls)
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildSimpleModePreservesRegistryWithoutLifecycleAuthority(t *testing.T) {
|
||||
registered := []SchedulerBucket{
|
||||
{GroupID: 0, Platform: "legacy-zero", Mode: "unknown"},
|
||||
{GroupID: 106, Platform: "legacy-positive", Mode: "unknown"},
|
||||
{GroupID: -8, Platform: "legacy-negative", Mode: "unknown"},
|
||||
}
|
||||
cache := newFullRebuildLifecycleCache(registered...)
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
activeIDsErr: errors.New("simple mode must not query groups"),
|
||||
fresh: make(map[int64]*Group),
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeSimple)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
|
||||
activeCalls, fallbackCalls, freshCalls := groups.stats()
|
||||
require.Zero(t, activeCalls)
|
||||
require.Zero(t, fallbackCalls)
|
||||
require.Empty(t, freshCalls)
|
||||
require.Equal(t, 15, cache.captureAttemptCount())
|
||||
require.Equal(t, 15, accounts.callCount())
|
||||
require.Equal(t, 15, accounts.groupCallCount(0))
|
||||
require.Empty(t, cache.retiredBuckets())
|
||||
require.Empty(t, cache.tokens())
|
||||
for _, bucket := range registered {
|
||||
_, published := cache.counts(bucket)
|
||||
require.Equal(t, 1, published, bucket.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildFreshReopenLockBusyRetriesWithoutBlockingOrdinaryTasks(t *testing.T) {
|
||||
const groupID int64 = 107
|
||||
canonical := schedulerBucketsForGroup(groupID)
|
||||
cache := newFullRebuildLifecycleCache()
|
||||
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[0]))
|
||||
cache.lockBusyOnce[canonical[0].String()] = true
|
||||
groups := &fullRebuildLifecycleGroupRepo{
|
||||
activeIDs: []int64{groupID},
|
||||
fresh: map[int64]*Group{
|
||||
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
|
||||
},
|
||||
freshErr: make(map[int64]error),
|
||||
}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
outbox := &outboxCleanupRepo{events: []SchedulerOutboxEvent{{ID: 1, EventType: SchedulerOutboxEventFullRebuild}}}
|
||||
svc := newFullRebuildLifecycleService(cache, outbox, accounts, groups, config.RunModeStandard)
|
||||
|
||||
svc.pollOutbox()
|
||||
require.Zero(t, cache.currentWatermark())
|
||||
_, groupZeroPublished := cache.counts(schedulerCanonicalBuckets(0)[0])
|
||||
require.Equal(t, 1, groupZeroPublished, "ordinary tasks must still run when one strict Reopen task is busy")
|
||||
require.Equal(t, 23, accounts.callCount())
|
||||
|
||||
svc.pollOutbox()
|
||||
require.Equal(t, int64(1), cache.currentWatermark())
|
||||
require.Equal(t, 47, accounts.callCount())
|
||||
_, busyBucketPublished := cache.counts(canonical[0])
|
||||
require.Equal(t, 1, busyBucketPublished)
|
||||
activeCalls, fallbackCalls, freshCalls := groups.stats()
|
||||
require.Equal(t, 2, activeCalls)
|
||||
require.Zero(t, fallbackCalls)
|
||||
require.Equal(t, []int64{groupID}, freshCalls)
|
||||
}
|
||||
|
||||
func TestSchedulerFullRebuildOrdinaryLockBusyKeepsExistingSkipSemantics(t *testing.T) {
|
||||
busyBucket := schedulerCanonicalBuckets(0)[0]
|
||||
cache := newFullRebuildLifecycleCache()
|
||||
cache.lockBusyOnce[busyBucket.String()] = true
|
||||
groups := &fullRebuildLifecycleGroupRepo{fresh: make(map[int64]*Group), freshErr: make(map[int64]error)}
|
||||
accounts := &fullRebuildAccountRepo{}
|
||||
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
|
||||
|
||||
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
|
||||
require.Equal(t, 11, accounts.callCount())
|
||||
attempts, published := cache.counts(busyBucket)
|
||||
require.Zero(t, attempts)
|
||||
require.Zero(t, published)
|
||||
}
|
||||
|
||||
func requireFullRebuildNoMutationOrDB(t *testing.T, cache *fullRebuildLifecycleCache, accounts *fullRebuildAccountRepo) {
|
||||
t.Helper()
|
||||
require.Zero(t, cache.captureAttemptCount())
|
||||
require.Empty(t, cache.retiredBuckets())
|
||||
require.Empty(t, cache.tokens())
|
||||
require.Zero(t, accounts.callCount())
|
||||
require.Zero(t, cache.totalSetAttempts())
|
||||
}
|
||||
@@ -17,6 +17,7 @@ type schedulerFullRebuildTestCache struct {
|
||||
mu sync.Mutex
|
||||
listErr error
|
||||
listCalls int
|
||||
captures int
|
||||
lockCalls int
|
||||
}
|
||||
|
||||
@@ -35,6 +36,9 @@ func (c *schedulerFullRebuildTestCache) TryLockBucket(context.Context, Scheduler
|
||||
}
|
||||
|
||||
func (c *schedulerFullRebuildTestCache) CaptureBucketWriteToken(_ context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) {
|
||||
c.mu.Lock()
|
||||
c.captures++
|
||||
c.mu.Unlock()
|
||||
return SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
|
||||
}
|
||||
|
||||
@@ -129,7 +133,7 @@ func TestSchedulerSnapshotServiceFullRebuildRunsAgainForSequentialRequest(t *tes
|
||||
require.Equal(t, requested, completed)
|
||||
}
|
||||
|
||||
func TestSchedulerSnapshotServiceInitialFullRebuildFallsBackWhenListBucketsFails(t *testing.T) {
|
||||
func TestSchedulerSnapshotServiceInitialFullRebuildFailsClosedWhenListBucketsFails(t *testing.T) {
|
||||
cache := &schedulerFullRebuildTestCache{listErr: errors.New("list buckets failed")}
|
||||
svc := NewSchedulerSnapshotService(cache, nil, nil, nil, nil)
|
||||
|
||||
@@ -137,13 +141,18 @@ func TestSchedulerSnapshotServiceInitialFullRebuildFallsBackWhenListBucketsFails
|
||||
|
||||
cache.mu.Lock()
|
||||
listCalls := cache.listCalls
|
||||
captures := cache.captures
|
||||
lockCalls := cache.lockCalls
|
||||
cache.mu.Unlock()
|
||||
require.Equal(t, 1, listCalls)
|
||||
require.Positive(t, lockCalls, "startup should rebuild default buckets after ListBuckets fails")
|
||||
require.Zero(t, captures)
|
||||
require.Zero(t, lockCalls)
|
||||
requested, completed := schedulerFullRebuildState(svc)
|
||||
require.EqualValues(t, 1, requested)
|
||||
require.Equal(t, requested, completed)
|
||||
svc.fullRebuildStateMu.Lock()
|
||||
require.ErrorIs(t, svc.fullRebuildLastErr, cache.listErr)
|
||||
svc.fullRebuildStateMu.Unlock()
|
||||
}
|
||||
|
||||
func schedulerFullRebuildState(svc *SchedulerSnapshotService) (requested uint64, completed uint64) {
|
||||
|
||||
@@ -4,7 +4,6 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -161,7 +160,7 @@ func TestSchedulerFullRebuildCapturesAllRegistryTokensBeforeDBLoad(t *testing.T)
|
||||
return []Account{{ID: 6101, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true}}, nil
|
||||
},
|
||||
}
|
||||
svc := NewSchedulerSnapshotService(cache, nil, repo, nil, &config.Config{
|
||||
svc := NewSchedulerSnapshotService(cache, nil, repo, &retirementGroupRepo{groups: []Group{{ID: 61, Status: StatusActive}}}, &config.Config{
|
||||
RunMode: config.RunModeStandard,
|
||||
Gateway: config.GatewayConfig{Scheduling: config.GatewaySchedulingConfig{
|
||||
DbFallbackEnabled: true,
|
||||
@@ -177,7 +176,7 @@ func TestSchedulerFullRebuildCapturesAllRegistryTokensBeforeDBLoad(t *testing.T)
|
||||
}
|
||||
|
||||
captures, reopens := cache.captureAndReopenCounts()
|
||||
require.Equal(t, 2, captures, "all registry tokens must be captured before the first DB load")
|
||||
require.Equal(t, 24, captures, "group0 and active-group canonical tokens must be captured before the first DB load")
|
||||
require.Zero(t, reopens)
|
||||
require.NoError(t, cache.RetireBucket(context.Background(), queued))
|
||||
_, err := cache.ReopenBucket(context.Background(), queued)
|
||||
@@ -262,28 +261,3 @@ func TestSchedulerFallbackReturnsDBAccountsWhenBucketRetired(t *testing.T) {
|
||||
require.Zero(t, setAttempts)
|
||||
require.Zero(t, published)
|
||||
}
|
||||
|
||||
func TestSchedulerDefaultBucketsUseCaptureAndListActiveFailureKeepsGroupZero(t *testing.T) {
|
||||
cache := newRetirementRaceCache()
|
||||
svc := NewSchedulerSnapshotService(
|
||||
cache,
|
||||
nil,
|
||||
nil,
|
||||
&retirementGroupRepo{err: errors.New("list active failed")},
|
||||
testConfig(),
|
||||
)
|
||||
|
||||
buckets, err := svc.defaultBuckets(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, buckets)
|
||||
for _, bucket := range buckets {
|
||||
require.Zero(t, bucket.GroupID)
|
||||
}
|
||||
|
||||
tasks, err := svc.prepareBucketWriteTasks(context.Background(), buckets)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, tasks, len(buckets))
|
||||
captures, reopens := cache.captureAndReopenCounts()
|
||||
require.Equal(t, len(buckets), captures)
|
||||
require.Zero(t, reopens)
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sort"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -47,6 +48,10 @@ type schedulerGroupLifecyclePlan struct {
|
||||
tasks []schedulerBucketWriteTask
|
||||
}
|
||||
|
||||
type schedulerActiveGroupIDLister interface {
|
||||
ListActiveIDs(ctx context.Context) ([]int64, error)
|
||||
}
|
||||
|
||||
type SchedulerSnapshotService struct {
|
||||
cache SchedulerCache
|
||||
outboxRepo SchedulerOutboxRepository
|
||||
@@ -225,18 +230,7 @@ func (s *SchedulerSnapshotService) runInitialRebuild() {
|
||||
_ = s.coalesceFullRebuild(func() error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
buckets, err := s.cache.ListBuckets(ctx)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err)
|
||||
}
|
||||
if len(buckets) == 0 {
|
||||
buckets, err = s.defaultBuckets(ctx)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := s.rebuildBuckets(ctx, buckets, "startup"); err != nil {
|
||||
if err := s.rebuildFullSnapshot(ctx, "startup"); err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup failed: %v", err)
|
||||
return err
|
||||
}
|
||||
@@ -665,6 +659,10 @@ func schedulerBucketsForGroup(groupID int64) []SchedulerBucket {
|
||||
if groupID <= 0 {
|
||||
return nil
|
||||
}
|
||||
return schedulerCanonicalBuckets(groupID)
|
||||
}
|
||||
|
||||
func schedulerCanonicalBuckets(groupID int64) []SchedulerBucket {
|
||||
buckets := make([]SchedulerBucket, 0, 12)
|
||||
for _, platform := range schedulerSnapshotPlatforms() {
|
||||
buckets = append(buckets,
|
||||
@@ -718,10 +716,8 @@ func (s *SchedulerSnapshotService) bucketsForPlatform(platform string, groupIDs
|
||||
|
||||
func (s *SchedulerSnapshotService) rebuildBuckets(ctx context.Context, buckets []SchedulerBucket, reason string) error {
|
||||
tasks, firstErr := s.prepareBucketWriteTasks(ctx, buckets)
|
||||
for _, task := range tasks {
|
||||
if err := s.rebuildBucketWithToken(ctx, task, reason); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
if err := s.rebuildPreparedBucketTasks(ctx, tasks, reason, false); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
@@ -748,8 +744,14 @@ func (s *SchedulerSnapshotService) prepareBucketWriteTasks(ctx context.Context,
|
||||
return tasks, firstErr
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) rebuildBucketWithToken(ctx context.Context, task schedulerBucketWriteTask, reason string) error {
|
||||
return s.rebuildBucketWithTokenPolicy(ctx, task, reason, false)
|
||||
func (s *SchedulerSnapshotService) rebuildPreparedBucketTasks(ctx context.Context, tasks []schedulerBucketWriteTask, reason string, strict bool) error {
|
||||
var firstErr error
|
||||
for _, task := range tasks {
|
||||
if err := s.rebuildBucketWithTokenPolicy(ctx, task, reason, strict); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicy(ctx context.Context, task schedulerBucketWriteTask, reason string, strict bool) error {
|
||||
@@ -801,21 +803,213 @@ func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error {
|
||||
return s.coalesceFullRebuild(func() error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
return s.rebuildFullSnapshot(ctx, reason)
|
||||
})
|
||||
}
|
||||
|
||||
buckets, err := s.cache.ListBuckets(ctx)
|
||||
func (s *SchedulerSnapshotService) rebuildFullSnapshot(ctx context.Context, reason string) error {
|
||||
if s.cache == nil {
|
||||
return ErrSchedulerCacheNotReady
|
||||
}
|
||||
|
||||
registered, err := s.cache.ListBuckets(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
registered = dedupeBuckets(registered)
|
||||
|
||||
if s.isRunModeSimple() {
|
||||
canonical := schedulerCanonicalBuckets(0)
|
||||
captured, err := s.captureFullRebuildCanonicalTasks(ctx, canonical)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err)
|
||||
return err
|
||||
}
|
||||
if len(buckets) == 0 {
|
||||
buckets, err = s.defaultBuckets(ctx)
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err)
|
||||
return err
|
||||
}
|
||||
ordinary := appendBucketsExcept(nil, registered, canonical)
|
||||
return s.prepareAndRebuildFullSnapshot(ctx, captured, nil, ordinary, reason)
|
||||
}
|
||||
|
||||
activeGroupIDs, err := s.listActiveSchedulerGroupIDs(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
activeGroups := make(map[int64]struct{}, len(activeGroupIDs))
|
||||
for _, groupID := range activeGroupIDs {
|
||||
activeGroups[groupID] = struct{}{}
|
||||
}
|
||||
|
||||
registeredByGroup := make(map[int64][]SchedulerBucket)
|
||||
for _, bucket := range registered {
|
||||
registeredByGroup[bucket.GroupID] = append(registeredByGroup[bucket.GroupID], bucket)
|
||||
}
|
||||
|
||||
groupZeroCanonical := schedulerCanonicalBuckets(0)
|
||||
capturedTasks, err := s.captureFullRebuildCanonicalTasks(ctx, groupZeroCanonical)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ordinaryBuckets := appendBucketsExcept(nil, registeredByGroup[0], groupZeroCanonical)
|
||||
for groupID, buckets := range registeredByGroup {
|
||||
if groupID < 0 {
|
||||
ordinaryBuckets = append(ordinaryBuckets, buckets...)
|
||||
}
|
||||
return s.rebuildBuckets(ctx, buckets, reason)
|
||||
})
|
||||
}
|
||||
|
||||
reopenedTasks := make([]schedulerBucketWriteTask, 0)
|
||||
for _, groupID := range activeGroupIDs {
|
||||
canonical := schedulerBucketsForGroup(groupID)
|
||||
canonicalTasks, captureErr := s.captureFullRebuildCanonicalTasks(ctx, canonical)
|
||||
if captureErr == nil {
|
||||
capturedTasks = append(capturedTasks, canonicalTasks...)
|
||||
ordinaryBuckets = appendBucketsExcept(ordinaryBuckets, registeredByGroup[groupID], canonical)
|
||||
continue
|
||||
}
|
||||
if !errors.Is(captureErr, ErrSchedulerBucketRetired) && !errors.Is(captureErr, ErrSchedulerBucketWriteFenced) {
|
||||
return captureErr
|
||||
}
|
||||
|
||||
// A prior full_rebuild event can observe the active state committed for a
|
||||
// later group_changed event. Recover here under fresh authority so the
|
||||
// earlier event cannot block the outbox watermark before that event runs.
|
||||
knownHistorical := registeredByGroup[groupID]
|
||||
if knownHistorical == nil {
|
||||
knownHistorical = []SchedulerBucket{}
|
||||
}
|
||||
plan, err := s.prepareGroupLifecycle(ctx, groupID, knownHistorical)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if plan.active {
|
||||
reopenedTasks = append(reopenedTasks, plan.tasks...)
|
||||
ordinaryBuckets = appendBucketsExcept(ordinaryBuckets, registeredByGroup[groupID], canonical)
|
||||
}
|
||||
}
|
||||
|
||||
staleGroupIDs := make([]int64, 0)
|
||||
for groupID := range registeredByGroup {
|
||||
if groupID <= 0 {
|
||||
continue
|
||||
}
|
||||
if _, active := activeGroups[groupID]; !active {
|
||||
staleGroupIDs = append(staleGroupIDs, groupID)
|
||||
}
|
||||
}
|
||||
sort.Slice(staleGroupIDs, func(i, j int) bool { return staleGroupIDs[i] < staleGroupIDs[j] })
|
||||
|
||||
for _, groupID := range staleGroupIDs {
|
||||
plan, err := s.prepareGroupLifecycle(ctx, groupID, registeredByGroup[groupID])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if plan.active {
|
||||
reopenedTasks = append(reopenedTasks, plan.tasks...)
|
||||
ordinaryBuckets = appendBucketsExcept(ordinaryBuckets, registeredByGroup[groupID], schedulerBucketsForGroup(groupID))
|
||||
}
|
||||
}
|
||||
|
||||
return s.prepareAndRebuildFullSnapshot(ctx, capturedTasks, reopenedTasks, ordinaryBuckets, reason)
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) listActiveSchedulerGroupIDs(ctx context.Context) ([]int64, error) {
|
||||
if s.groupRepo == nil {
|
||||
return nil, ErrSchedulerCacheNotReady
|
||||
}
|
||||
|
||||
var groupIDs []int64
|
||||
if lister, ok := s.groupRepo.(schedulerActiveGroupIDLister); ok {
|
||||
ids, err := lister.ListActiveIDs(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groupIDs = ids
|
||||
} else {
|
||||
groups, err := s.groupRepo.ListActive(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
groupIDs = make([]int64, 0, len(groups))
|
||||
for _, group := range groups {
|
||||
groupIDs = append(groupIDs, group.ID)
|
||||
}
|
||||
}
|
||||
|
||||
seen := make(map[int64]struct{}, len(groupIDs))
|
||||
normalized := make([]int64, 0, len(groupIDs))
|
||||
for _, groupID := range groupIDs {
|
||||
if groupID <= 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[groupID]; ok {
|
||||
continue
|
||||
}
|
||||
seen[groupID] = struct{}{}
|
||||
normalized = append(normalized, groupID)
|
||||
}
|
||||
sort.Slice(normalized, func(i, j int) bool { return normalized[i] < normalized[j] })
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) prepareAndRebuildFullSnapshot(
|
||||
ctx context.Context,
|
||||
captured []schedulerBucketWriteTask,
|
||||
reopened []schedulerBucketWriteTask,
|
||||
ordinaryBuckets []SchedulerBucket,
|
||||
reason string,
|
||||
) error {
|
||||
preparedBuckets := make(map[SchedulerBucket]struct{}, len(captured)+len(reopened))
|
||||
for _, task := range captured {
|
||||
preparedBuckets[task.bucket] = struct{}{}
|
||||
}
|
||||
for _, task := range reopened {
|
||||
preparedBuckets[task.bucket] = struct{}{}
|
||||
}
|
||||
|
||||
ordinaryBuckets = dedupeBuckets(ordinaryBuckets)
|
||||
toCapture := make([]SchedulerBucket, 0, len(ordinaryBuckets))
|
||||
for _, bucket := range ordinaryBuckets {
|
||||
if _, ok := preparedBuckets[bucket]; !ok {
|
||||
toCapture = append(toCapture, bucket)
|
||||
}
|
||||
}
|
||||
ordinary, firstErr := s.prepareBucketWriteTasks(ctx, toCapture)
|
||||
if firstErr != nil {
|
||||
return firstErr
|
||||
}
|
||||
captured = append(captured, ordinary...)
|
||||
if err := s.rebuildPreparedBucketTasks(ctx, reopened, reason, true); err != nil {
|
||||
firstErr = err
|
||||
}
|
||||
if err := s.rebuildPreparedBucketTasks(ctx, captured, reason, false); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) captureFullRebuildCanonicalTasks(ctx context.Context, buckets []SchedulerBucket) ([]schedulerBucketWriteTask, error) {
|
||||
if s.cache == nil {
|
||||
return nil, ErrSchedulerCacheNotReady
|
||||
}
|
||||
tasks := make([]schedulerBucketWriteTask, 0, len(buckets))
|
||||
for _, bucket := range buckets {
|
||||
token, err := s.cache.CaptureBucketWriteToken(ctx, bucket)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tasks = append(tasks, schedulerBucketWriteTask{bucket: bucket, token: token})
|
||||
}
|
||||
return tasks, nil
|
||||
}
|
||||
|
||||
func appendBucketsExcept(dst, buckets, excluded []SchedulerBucket) []SchedulerBucket {
|
||||
excludedKeys := make(map[SchedulerBucket]struct{}, len(excluded))
|
||||
for _, bucket := range excluded {
|
||||
excludedKeys[bucket] = struct{}{}
|
||||
}
|
||||
for _, bucket := range buckets {
|
||||
if _, ok := excludedKeys[bucket]; !ok {
|
||||
dst = append(dst, bucket)
|
||||
}
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) coalesceFullRebuild(run func() error) error {
|
||||
@@ -1053,38 +1247,6 @@ func (s *SchedulerSnapshotService) fullRebuildInterval() time.Duration {
|
||||
return time.Duration(sec) * time.Second
|
||||
}
|
||||
|
||||
func (s *SchedulerSnapshotService) defaultBuckets(ctx context.Context) ([]SchedulerBucket, error) {
|
||||
buckets := make([]SchedulerBucket, 0)
|
||||
platforms := []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok}
|
||||
for _, platform := range platforms {
|
||||
buckets = append(buckets, SchedulerBucket{GroupID: 0, Platform: platform, Mode: SchedulerModeSingle})
|
||||
buckets = append(buckets, SchedulerBucket{GroupID: 0, Platform: platform, Mode: SchedulerModeForced})
|
||||
if platform == PlatformAnthropic || platform == PlatformGemini {
|
||||
buckets = append(buckets, SchedulerBucket{GroupID: 0, Platform: platform, Mode: SchedulerModeMixed})
|
||||
}
|
||||
}
|
||||
|
||||
if s.isRunModeSimple() || s.groupRepo == nil {
|
||||
return dedupeBuckets(buckets), nil
|
||||
}
|
||||
|
||||
groups, err := s.groupRepo.ListActive(ctx)
|
||||
if err != nil {
|
||||
return dedupeBuckets(buckets), nil
|
||||
}
|
||||
for _, group := range groups {
|
||||
if group.Platform == "" {
|
||||
continue
|
||||
}
|
||||
buckets = append(buckets, SchedulerBucket{GroupID: group.ID, Platform: group.Platform, Mode: SchedulerModeSingle})
|
||||
buckets = append(buckets, SchedulerBucket{GroupID: group.ID, Platform: group.Platform, Mode: SchedulerModeForced})
|
||||
if group.Platform == PlatformAnthropic || group.Platform == PlatformGemini {
|
||||
buckets = append(buckets, SchedulerBucket{GroupID: group.ID, Platform: group.Platform, Mode: SchedulerModeMixed})
|
||||
}
|
||||
}
|
||||
return dedupeBuckets(buckets), nil
|
||||
}
|
||||
|
||||
func dedupeBuckets(in []SchedulerBucket) []SchedulerBucket {
|
||||
seen := make(map[string]struct{}, len(in))
|
||||
out := make([]SchedulerBucket, 0, len(in))
|
||||
|
||||
Reference in New Issue
Block a user