fix: reduce token refresh retry amplification

This commit is contained in:
jjaw authored and shaw committed 2026-06-16 11:44:39 +08:00
1 parent b56c501d30
commit 74199b6a6e
7 files changed
+386 -14

No files matched your search

+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
}