Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/scheduler_snapshot_full_rebuild_test.go

@@ -0,0 +1,162 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+type schedulerFullRebuildTestCache struct {
+	SchedulerCache
+
+	mu        sync.Mutex
+	listErr   error
+	listCalls int
+	captures  int
+	lockCalls int
+}
+
+func (c *schedulerFullRebuildTestCache) ListBuckets(context.Context) ([]SchedulerBucket, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.listCalls++
+	return nil, c.listErr
+}
+
+func (c *schedulerFullRebuildTestCache) TryLockBucket(context.Context, SchedulerBucket, time.Duration) (bool, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.lockCalls++
+	return false, nil
+}
+
+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
+}
+
+func (c *schedulerFullRebuildTestCache) ReopenBucket(_ context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) {
+	return SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
+}
+
+func TestSchedulerSnapshotServiceFullRebuildCoalescesConcurrentRequestsIntoTrailingRun(t *testing.T) {
+	svc := &SchedulerSnapshotService{}
+	wantTrailingErr := errors.New("trailing rebuild failed")
+	firstStarted := make(chan struct{})
+	releaseFirst := make(chan struct{})
+	var releaseOnce sync.Once
+	release := func() {
+		releaseOnce.Do(func() { close(releaseFirst) })
+	}
+	defer release()
+
+	var calls atomic.Int32
+	var active atomic.Int32
+	var maxActive atomic.Int32
+	run := func() error {
+		call := calls.Add(1)
+		currentActive := active.Add(1)
+		defer active.Add(-1)
+		for {
+			previousMax := maxActive.Load()
+			if currentActive <= previousMax || maxActive.CompareAndSwap(previousMax, currentActive) {
+				break
+			}
+		}
+		if call == 1 {
+			close(firstStarted)
+			<-releaseFirst
+			return nil
+		}
+		return wantTrailingErr
+	}
+
+	firstResult := make(chan error, 1)
+	go func() {
+		firstResult <- svc.coalesceFullRebuild(run)
+	}()
+
+	select {
+	case <-firstStarted:
+	case <-time.After(time.Second):
+		t.Fatal("first rebuild did not start")
+	}
+
+	const followers = 20
+	followerResults := make(chan error, followers)
+	for range followers {
+		go func() {
+			followerResults <- svc.coalesceFullRebuild(run)
+		}()
+	}
+
+	require.Eventually(t, func() bool {
+		requested, _ := schedulerFullRebuildState(svc)
+		return requested == followers+1
+	}, time.Second, time.Millisecond)
+	release()
+
+	require.NoError(t, <-firstResult)
+	for range followers {
+		require.ErrorIs(t, <-followerResults, wantTrailingErr)
+	}
+	require.EqualValues(t, 2, calls.Load())
+	require.EqualValues(t, 1, maxActive.Load())
+	requested, completed := schedulerFullRebuildState(svc)
+	require.EqualValues(t, followers+1, requested)
+	require.Equal(t, requested, completed)
+}
+
+func TestSchedulerSnapshotServiceFullRebuildRunsAgainForSequentialRequest(t *testing.T) {
+	svc := &SchedulerSnapshotService{}
+	wantSecondErr := errors.New("second rebuild failed")
+	var calls atomic.Int32
+	run := func() error {
+		if calls.Add(1) == 2 {
+			return wantSecondErr
+		}
+		return nil
+	}
+
+	require.NoError(t, svc.coalesceFullRebuild(run))
+	require.ErrorIs(t, svc.coalesceFullRebuild(run), wantSecondErr)
+	require.EqualValues(t, 2, calls.Load())
+	requested, completed := schedulerFullRebuildState(svc)
+	require.EqualValues(t, 2, requested)
+	require.Equal(t, requested, completed)
+}
+
+func TestSchedulerSnapshotServiceInitialFullRebuildFailsClosedWhenListBucketsFails(t *testing.T) {
+	cache := &schedulerFullRebuildTestCache{listErr: errors.New("list buckets failed")}
+	svc := NewSchedulerSnapshotService(cache, nil, nil, nil, nil)
+
+	svc.runInitialRebuild()
+
+	cache.mu.Lock()
+	listCalls := cache.listCalls
+	captures := cache.captures
+	lockCalls := cache.lockCalls
+	cache.mu.Unlock()
+	require.Equal(t, 1, listCalls)
+	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) {
+	svc.fullRebuildStateMu.Lock()
+	defer svc.fullRebuildStateMu.Unlock()
+	return svc.fullRebuildRequested, svc.fullRebuildCompleted
+}