fix: stop scheduler work after request cancellation

This commit is contained in:
rick147
2026-08-02 21:32:31 +08:00
parent a0802f00b6
commit f970bd48c9
2 changed files with 97 additions and 0 deletions
@@ -0,0 +1,79 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
type schedulerCancellationCache struct {
SchedulerCache
cancel context.CancelFunc
tokenCaptures int
}
func (c *schedulerCancellationCache) GetSnapshot(ctx context.Context, _ SchedulerBucket) ([]*Account, bool, error) {
c.cancel()
return nil, false, ctx.Err()
}
func (c *schedulerCancellationCache) CaptureBucketWriteToken(ctx context.Context, _ SchedulerBucket) (SchedulerBucketWriteToken, error) {
c.tokenCaptures++
return SchedulerBucketWriteToken{}, ctx.Err()
}
func (c *schedulerCancellationCache) GetAccount(ctx context.Context, _ int64) (*Account, error) {
c.cancel()
return nil, ctx.Err()
}
type schedulerCancellationAccountRepo struct {
AccountRepository
listCalls int
getByIDCalls int
}
func (r *schedulerCancellationAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, _ string) ([]Account, error) {
r.listCalls++
return nil, ctx.Err()
}
func (r *schedulerCancellationAccountRepo) GetByID(ctx context.Context, _ int64) (*Account, error) {
r.getByIDCalls++
return nil, ctx.Err()
}
func TestSchedulerSnapshotListStopsAfterRequestCancellation(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
cache := &schedulerCancellationCache{cancel: cancel}
repo := &schedulerCancellationAccountRepo{}
svc := NewSchedulerSnapshotService(cache, nil, repo, nil, nil)
accounts, useMixed, err := svc.ListSchedulableAccounts(ctx, nil, PlatformOpenAI, false)
require.ErrorIs(t, err, context.Canceled)
require.Nil(t, accounts)
require.False(t, useMixed)
require.Zero(t, cache.tokenCaptures, "canceled requests must not capture a cache publish token")
require.Zero(t, repo.listCalls, "canceled requests must not fall back to the database")
}
func TestSchedulerSnapshotGetAccountStopsAfterRequestCancellation(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
cache := &schedulerCancellationCache{cancel: cancel}
repo := &schedulerCancellationAccountRepo{}
svc := NewSchedulerSnapshotService(cache, nil, repo, nil, nil)
account, err := svc.GetAccount(ctx, 42)
require.ErrorIs(t, err, context.Canceled)
require.Nil(t, account)
require.Zero(t, repo.getByIDCalls, "canceled requests must not fall back to the database")
}
@@ -213,15 +213,24 @@ func (s *SchedulerSnapshotService) ListSchedulableAccounts(ctx context.Context,
bucket := s.bucketFor(groupID, platform, mode)
var writeToken SchedulerBucketWriteToken
canPublish := false
if err := ctx.Err(); err != nil {
return nil, useMixed, err
}
if s.cache != nil {
cached, hit, err := s.cache.GetSnapshot(ctx, bucket)
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, useMixed, ctxErr
}
if err != nil {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] cache read failed: bucket=%s err=%v", bucket.String(), err)
} else if hit {
return derefAccounts(cached), useMixed, nil
}
token, err := s.cache.CaptureBucketWriteToken(ctx, bucket)
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, useMixed, ctxErr
}
if err != nil {
if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) {
slog.Debug("[Scheduler] cache publish fenced", "bucket", bucket.String())
@@ -245,6 +254,9 @@ func (s *SchedulerSnapshotService) ListSchedulableAccounts(ctx context.Context,
if err != nil {
return nil, useMixed, err
}
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, useMixed, ctxErr
}
if s.cache != nil && canPublish {
if err := s.cache.SetSnapshot(fallbackCtx, bucket, writeToken, accounts); err != nil {
@@ -263,8 +275,14 @@ func (s *SchedulerSnapshotService) GetAccount(ctx context.Context, accountID int
if accountID <= 0 {
return nil, nil
}
if err := ctx.Err(); err != nil {
return nil, err
}
if s.cache != nil {
account, err := s.cache.GetAccount(ctx, accountID)
if ctxErr := ctx.Err(); ctxErr != nil {
return nil, ctxErr
}
if err != nil {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] account cache read failed: id=%d err=%v", accountID, err)
} else if account != nil {