fix: reduce token refresh retry amplification

This commit is contained in:
jjaw
2026-06-16 11:44:39 +08:00
committed by shaw
parent b56c501d30
commit 74199b6a6e
7 changed files with 386 additions and 14 deletions
+60 -1
View File
@@ -649,6 +649,58 @@ func (r *accountRepository) ListActive(ctx context.Context) ([]service.Account,
return r.accountsToService(ctx, accounts)
}
func (r *accountRepository) ListOAuthRefreshCandidates(ctx context.Context) ([]service.Account, error) {
if r.sql == nil {
return nil, errors.New("account repository SQL executor not configured")
}
rows, err := r.sql.QueryContext(ctx, `
SELECT id
FROM accounts
WHERE deleted_at IS NULL
AND status = 'active'
AND type = 'oauth'
AND platform IN ('anthropic', 'openai', 'gemini', 'antigravity')
AND credentials ? 'refresh_token'
AND btrim(credentials->>'refresh_token') <> ''
AND NOT (
temp_unschedulable_until > NOW()
AND temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'
)
ORDER BY priority ASC, id ASC
`)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var ids []int64
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
return nil, err
}
ids = append(ids, id)
}
if err := rows.Err(); err != nil {
return nil, err
}
if len(ids) == 0 {
return []service.Account{}, nil
}
accounts, err := r.GetByIDs(ctx, ids)
if err != nil {
return nil, err
}
out := make([]service.Account, 0, len(accounts))
for _, account := range accounts {
if account != nil {
out = append(out, *account)
}
}
return out, nil
}
func (r *accountRepository) ListByPlatform(ctx context.Context, platform string) ([]service.Account, error) {
accounts, err := r.client.Account.Query().
Where(
@@ -1148,7 +1200,7 @@ func (r *accountRepository) SetOverloaded(ctx context.Context, id int64, until t
}
func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64, until time.Time, reason string) error {
_, err := r.sql.ExecContext(ctx, `
result, err := r.sql.ExecContext(ctx, `
UPDATE accounts
SET temp_unschedulable_until = $1,
temp_unschedulable_reason = $2,
@@ -1160,6 +1212,13 @@ func (r *accountRepository) SetTempUnschedulable(ctx context.Context, id int64,
if err != nil {
return err
}
affected, err := result.RowsAffected()
if err != nil {
return err
}
if affected <= 0 {
return nil
}
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue temp unschedulable failed: account=%d err=%v", id, err)
}
@@ -0,0 +1,98 @@
package repository
import (
"context"
"database/sql"
"regexp"
"strings"
"testing"
"time"
sqlmock "github.com/DATA-DOG/go-sqlmock"
"github.com/stretchr/testify/require"
)
func TestAccountRepository_SetTempUnschedulable_NoRowsAffectedDoesNotWriteOutbox(t *testing.T) {
exec := &recordingSQLExecutor{result: rowsAffectedResult(0)}
repo := newAccountRepositoryWithSQL(nil, exec, nil)
until := time.Now().Add(10 * time.Minute)
err := repo.SetTempUnschedulable(context.Background(), 42, until, "retry")
require.NoError(t, err)
require.Len(t, exec.execQueries, 1)
require.Contains(t, exec.execQueries[0], "UPDATE accounts")
require.NotContains(t, strings.Join(exec.execQueries, "\n"), "scheduler_outbox")
}
func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherRegexp))
require.NoError(t, err)
defer func() { _ = db.Close() }()
var capturedSQL string
mock.ExpectQuery("SELECT id").
WillReturnRows(sqlmock.NewRows([]string{"id"})).
WillDelayFor(0)
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil)
accounts, err := repo.ListOAuthRefreshCandidates(context.Background())
require.NoError(t, err)
require.Empty(t, accounts)
normalized := normalizeSQLWhitespace(capturedSQL)
require.Contains(t, normalized, "deleted_at IS NULL")
require.Contains(t, normalized, "status = 'active'")
require.Contains(t, normalized, "type = 'oauth'")
require.Contains(t, normalized, "platform IN ('anthropic', 'openai', 'gemini', 'antigravity')")
require.Contains(t, normalized, "credentials ? 'refresh_token'")
require.Contains(t, normalized, "btrim(credentials->>'refresh_token') <> ''")
require.Contains(t, normalized, "temp_unschedulable_until > NOW()")
require.Contains(t, normalized, "temp_unschedulable_reason LIKE 'token refresh retry exhausted:%'")
require.Contains(t, normalized, "ORDER BY priority ASC, id ASC")
require.NotContains(t, normalized, "credentials->>'expires_at'")
require.NoError(t, mock.ExpectationsWereMet())
}
type captureQuerySQL struct {
db *sql.DB
captured *string
}
func (c captureQuerySQL) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
return c.db.ExecContext(ctx, query, args...)
}
func (c captureQuerySQL) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
if c.captured != nil {
*c.captured = query
}
return c.db.QueryContext(ctx, query, args...)
}
func normalizeSQLWhitespace(sql string) string {
return strings.Join(regexp.MustCompile(`\s+`).Split(strings.TrimSpace(sql), -1), " ")
}
type rowsAffectedResult int64
func (r rowsAffectedResult) LastInsertId() (int64, error) { return 0, nil }
func (r rowsAffectedResult) RowsAffected() (int64, error) { return int64(r), nil }
type recordingSQLExecutor struct {
result sql.Result
err error
execQueries []string
}
func (e *recordingSQLExecutor) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
e.execQueries = append(e.execQueries, query)
if e.err != nil {
return nil, e.err
}
return e.result, nil
}
func (e *recordingSQLExecutor) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
return nil, sql.ErrNoRows
}
@@ -0,0 +1,25 @@
//go:build unit
package service
import "context"
func (s *accountRepoStub) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (r *openAIAccountTestRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (m *groupAwareMockAccountRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (m *mockAccountRepoForPlatform) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected ListOAuthRefreshCandidates call")
}
func (m *mockAccountRepoForGemini) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
return m.ListActive(context.Background())
}
@@ -41,6 +41,7 @@ type AccountRepository interface {
ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, accountType, status, search string, groupID int64, privacyMode string) ([]Account, *pagination.PaginationResult, error)
ListByGroup(ctx context.Context, groupID int64) ([]Account, error)
ListActive(ctx context.Context) ([]Account, error)
ListOAuthRefreshCandidates(ctx context.Context) ([]Account, error)
ListByPlatform(ctx context.Context, platform string) ([]Account, error)
UpdateLastUsed(ctx context.Context, id int64) error
@@ -93,6 +93,9 @@ func (m *sessionWindowMockRepo) ListByGroup(context.Context, int64) ([]Account,
func (m *sessionWindowMockRepo) ListActive(context.Context) ([]Account, error) {
panic("unexpected")
}
func (m *sessionWindowMockRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
panic("unexpected")
}
func (m *sessionWindowMockRepo) ListByPlatform(context.Context, string) ([]Account, error) {
panic("unexpected")
}
@@ -12,8 +12,12 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
)
// tokenRefreshTempUnschedDuration token 刷新重试耗尽后临时不可调度的持续时间
const tokenRefreshTempUnschedDuration = 10 * time.Minute
const (
// tokenRefreshTempUnschedDuration token 刷新重试耗尽后临时不可调度的持续时间
tokenRefreshTempUnschedDuration = 10 * time.Minute
tokenRefreshRetryExhaustedReasonPrefix = "token refresh retry exhausted:"
)
// TokenRefreshService OAuth token自动刷新服务
// 定期检查并刷新即将过期的token
@@ -255,10 +259,9 @@ func (s *TokenRefreshService) processRefresh() {
}
}
// listActiveAccounts 获取所有active状态的账号
// 使用ListActive确保刷新所有活跃账号的token(包括临时禁用的)
// listActiveAccounts 获取后台 OAuth token 刷新候选账号
func (s *TokenRefreshService) listActiveAccounts(ctx context.Context) ([]Account, error) {
return s.accountRepo.ListActive(ctx)
return s.accountRepo.ListOAuthRefreshCandidates(ctx)
}
// refreshWithRetry 带重试的刷新
@@ -310,9 +313,6 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc
"error", setErr,
)
}
// 刷新失败但 access_token 可能仍有效,尝试设置隐私
s.ensureOpenAIPrivacy(ctx, account)
s.ensureAntigravityPrivacy(ctx, account)
return err
}
@@ -340,13 +340,9 @@ func (s *TokenRefreshService) refreshWithRetry(ctx context.Context, account *Acc
"error", lastErr,
)
// 刷新失败但 access_token 可能仍有效,尝试设置隐私
s.ensureOpenAIPrivacy(ctx, account)
s.ensureAntigravityPrivacy(ctx, account)
// 设置临时不可调度 10 分钟(不标记 error,保持 status=active 让下个刷新周期能继续尝试)
until := time.Now().Add(tokenRefreshTempUnschedDuration)
reason := fmt.Sprintf("token refresh retry exhausted: %v", lastErr)
reason := fmt.Sprintf("%s %v", tokenRefreshRetryExhaustedReasonPrefix, lastErr)
s.notifyAccountSchedulingBlocked(account, until, "token_refresh_retry_exhausted")
if setErr := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); setErr != nil {
slog.Warn("token_refresh.set_temp_unschedulable_failed",
@@ -0,0 +1,190 @@
package service
import (
"context"
"errors"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/imroc/req/v3"
"github.com/stretchr/testify/require"
)
type tokenRefreshCandidateRepo struct {
AccountRepository
accounts []Account
updatedCredentialIDs []int64
setErrorCalls int
setTempUnschedCalls int
lastTempUnschedReason string
listActiveCalls int
}
func (r *tokenRefreshCandidateRepo) ListActive(context.Context) ([]Account, error) {
r.listActiveCalls++
return r.accounts, nil
}
func (r *tokenRefreshCandidateRepo) ListOAuthRefreshCandidates(context.Context) ([]Account, error) {
candidates := make([]Account, 0, len(r.accounts))
now := time.Now()
for _, account := range r.accounts {
refreshToken, _ := account.Credentials["refresh_token"].(string)
inRetryCooldown := account.TempUnschedulableUntil != nil &&
account.TempUnschedulableUntil.After(now) &&
strings.HasPrefix(account.TempUnschedulableReason, tokenRefreshRetryExhaustedReasonPrefix)
if account.Status != StatusActive ||
account.Type != AccountTypeOAuth ||
!isOAuthRefreshPlatform(account.Platform) ||
strings.TrimSpace(refreshToken) == "" ||
inRetryCooldown {
continue
}
candidates = append(candidates, account)
}
return candidates, nil
}
func (r *tokenRefreshCandidateRepo) UpdateCredentials(_ context.Context, id int64, _ map[string]any) error {
r.updatedCredentialIDs = append(r.updatedCredentialIDs, id)
return nil
}
func (r *tokenRefreshCandidateRepo) SetError(context.Context, int64, string) error {
r.setErrorCalls++
return nil
}
func (r *tokenRefreshCandidateRepo) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
r.setTempUnschedCalls++
r.lastTempUnschedReason = reason
return nil
}
func isOAuthRefreshPlatform(platform string) bool {
switch platform {
case PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity:
return true
default:
return false
}
}
type tokenRefreshTestRefresher struct {
err error
}
func (r *tokenRefreshTestRefresher) CanRefresh(*Account) bool { return true }
func (r *tokenRefreshTestRefresher) NeedsRefresh(*Account, time.Duration) bool { return true }
func (r *tokenRefreshTestRefresher) Refresh(context.Context, *Account) (map[string]any, error) {
if r.err != nil {
return nil, r.err
}
return map[string]any{"access_token": "new-access-token", "refresh_token": "new-refresh-token"}, nil
}
func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing.T) {
future := time.Now().Add(10 * time.Minute)
repo := &tokenRefreshCandidateRepo{
accounts: []Account{
{
ID: 1,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-token"},
},
{
ID: 2,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{},
},
{
ID: 3,
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
Status: StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-token"},
},
{
ID: 4,
Platform: PlatformAntigravity,
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-token"},
TempUnschedulableUntil: &future,
TempUnschedulableReason: tokenRefreshRetryExhaustedReasonPrefix + " network timeout",
},
{
ID: 5,
Platform: "other",
Type: AccountTypeOAuth,
Status: StatusActive,
Credentials: map[string]any{"refresh_token": "refresh-token"},
},
},
}
svc := &TokenRefreshService{
accountRepo: repo,
refreshers: []TokenRefresher{&tokenRefreshTestRefresher{}},
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{RefreshBeforeExpiryHours: 1, MaxRetries: 1},
}
svc.processRefresh()
require.Zero(t, repo.listActiveCalls, "TokenRefreshService should not use the broad active-account query")
require.Equal(t, []int64{1}, repo.updatedCredentialIDs)
}
func TestTokenRefreshService_RefreshFailureDoesNotCallPrivacy(t *testing.T) {
tests := []struct {
name string
err error
}{
{name: "retry exhausted", err: errors.New("temporary upstream timeout")},
{name: "non retryable", err: errors.New("invalid_grant: token revoked")},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &tokenRefreshCandidateRepo{}
svc := &TokenRefreshService{
accountRepo: repo,
refreshPolicy: DefaultBackgroundRefreshPolicy(),
cfg: &config.TokenRefreshConfig{MaxRetries: 1, RetryBackoffSeconds: 0},
privacyClientFactory: func(string) (*req.Client, error) {
t.Fatalf("privacy client factory must not be called on refresh failure")
return nil, errors.New("unexpected privacy call")
},
}
account := &Account{
ID: 11,
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"access_token": "old-access-token",
"refresh_token": "refresh-token",
},
}
err := svc.refreshWithRetry(context.Background(), account, &tokenRefreshTestRefresher{err: tt.err}, nil, time.Hour)
require.Error(t, err)
if isNonRetryableRefreshError(tt.err) {
require.Equal(t, 1, repo.setErrorCalls)
require.Zero(t, repo.setTempUnschedCalls)
} else {
require.Zero(t, repo.setErrorCalls)
require.Equal(t, 1, repo.setTempUnschedCalls)
require.True(t, strings.HasPrefix(repo.lastTempUnschedReason, tokenRefreshRetryExhaustedReasonPrefix))
}
})
}
}