mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-21 14:19:18 +08:00
fix: reduce token refresh retry amplification
This commit is contained in:
@@ -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))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user