修复调度全量重建并发合并

2026-09-24 16:05:44 +08:00 · 2026-07-13 21:59:55 +08:00
parent 7d239d62e8
commit 831862b924
2 changed files with 211 additions and 25 deletions
@@ -0,0 +1,145 @@
+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
+	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 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 TestSchedulerSnapshotServiceInitialFullRebuildFallsBackWhenListBucketsFails(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
+	lockCalls := cache.lockCalls
+	cache.mu.Unlock()
+	require.Equal(t, 1, listCalls)
+	require.Positive(t, lockCalls, "startup should rebuild default buckets after ListBuckets fails")
+	requested, completed := schedulerFullRebuildState(svc)
+	require.EqualValues(t, 1, requested)
+	require.Equal(t, requested, completed)
+}
+
+func schedulerFullRebuildState(svc *SchedulerSnapshotService) (requested uint64, completed uint64) {
+	svc.fullRebuildStateMu.Lock()
+	defer svc.fullRebuildStateMu.Unlock()
+	return svc.fullRebuildRequested, svc.fullRebuildCompleted
+}
@@ -43,6 +43,12 @@ type SchedulerSnapshotService struct {
 	fallbackLimit *fallbackLimiter
 	lagMu         sync.Mutex
 	lagFailures   int
+
+	fullRebuildRunMu     sync.Mutex
+	fullRebuildStateMu   sync.Mutex
+	fullRebuildRequested uint64
+	fullRebuildCompleted uint64
+	fullRebuildLastErr   error
 }

 func NewSchedulerSnapshotService(
@@ -183,22 +189,26 @@ func (s *SchedulerSnapshotService) runInitialRebuild() {
 	if s.cache == nil {
 		return
 	}
-	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)
+	_ = 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] default buckets failed: %v", err)
-			return
+			logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err)
 		}
-	}
-	if err := s.rebuildBuckets(ctx, buckets, "startup"); err != nil {
-		logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup 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 {
+			logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup failed: %v", err)
+			return err
+		}
+		return nil
+	})
 }

 func (s *SchedulerSnapshotService) runOutboxWorker(interval time.Duration) {
@@ -602,22 +612,53 @@ func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error {
 	if s.cache == nil {
 		return ErrSchedulerCacheNotReady
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
-	defer cancel()
+	return 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)
-		return err
-	}
-	if len(buckets) == 0 {
-		buckets, err = s.defaultBuckets(ctx)
+		buckets, err := s.cache.ListBuckets(ctx)
 		if err != nil {
-			logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err)
+			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
+			}
+		}
+		return s.rebuildBuckets(ctx, buckets, reason)
+	})
+}
+
+func (s *SchedulerSnapshotService) coalesceFullRebuild(run func() error) error {
+	s.fullRebuildStateMu.Lock()
+	s.fullRebuildRequested++
+	requestID := s.fullRebuildRequested
+	s.fullRebuildStateMu.Unlock()
+
+	s.fullRebuildRunMu.Lock()
+	defer s.fullRebuildRunMu.Unlock()
+
+	s.fullRebuildStateMu.Lock()
+	if s.fullRebuildCompleted >= requestID {
+		err := s.fullRebuildLastErr
+		s.fullRebuildStateMu.Unlock()
+		return err
 	}
-	return s.rebuildBuckets(ctx, buckets, reason)
+	// 当前轮重建可能早于新 outbox 事件对应事务的提交，不能让后到请求直接复用当前轮。
+	// 每轮开始前记录可覆盖的请求代次，执行期间登记的请求统一合并到下一轮。
+	coveredThrough := s.fullRebuildRequested
+	s.fullRebuildStateMu.Unlock()
+
+	err := run()
+
+	s.fullRebuildStateMu.Lock()
+	s.fullRebuildCompleted = coveredThrough
+	s.fullRebuildLastErr = err
+	s.fullRebuildStateMu.Unlock()
+	return err
 }

 func (s *SchedulerSnapshotService) checkOutboxLag(ctx context.Context, oldest SchedulerOutboxEvent, watermark int64) {