mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-21 04:10:58 +00:00
fix: stop scheduler work after request cancellation
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user