5c8c15cdb1
Two anti-thundering-herd improvements: 1. OAuthRefreshAPI.RefreshIfNeeded Wrap the existing distributed-lock + DB-reread + executor.Refresh pipeline in a per-process singleflight keyed by cacheKey+window. Without this, N concurrent goroutines on the same account each pay one Redis lock RTT and one DB reread; with it, only the leader pays and the rest share the result. The refreshWindow is part of the key so a long background-refresh window cannot starve a short foreground-refresh window. 2. accountRepository.SetTempUnschedulable Wrap the same path (UPDATE + scheduler outbox enqueue + scheduler cache sync) in a per-process singleflight keyed by id+until+reason. The SQL guard (existing < new) already makes the UPDATE idempotent, but N callers still cost N round-trips and N outbox inserts. With singleflight, an upstream 401 burst that hits the same account collapses to one execution. Tests cover dedup behavior, key separation by account / refresh window, and that the SQL exec count drops from N to <=2 (UPDATE + outbox).
120 lines
3.6 KiB
Go
120 lines
3.6 KiB
Go
package repository
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"sync"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// blockingExecutor 是一个最小化的 sqlExecutor 实现，用于精确控制并发时序。
|
||
// ExecContext 会等待 release 信号才返回，便于让多个 goroutine 集中堆积在
|
||
// singleflight 的同一窗口内。
|
||
type blockingExecutor struct {
|
||
mu sync.Mutex
|
||
execCalls int32
|
||
queryCalls int32
|
||
release chan struct{}
|
||
concurrent int32
|
||
maxObserved int32
|
||
}
|
||
|
||
func newBlockingExecutor() *blockingExecutor {
|
||
return &blockingExecutor{release: make(chan struct{})}
|
||
}
|
||
|
||
func (e *blockingExecutor) Release() { close(e.release) }
|
||
|
||
func (e *blockingExecutor) ExecContext(_ context.Context, _ string, _ ...any) (sql.Result, error) {
|
||
atomic.AddInt32(&e.execCalls, 1)
|
||
c := atomic.AddInt32(&e.concurrent, 1)
|
||
for {
|
||
old := atomic.LoadInt32(&e.maxObserved)
|
||
if c <= old || atomic.CompareAndSwapInt32(&e.maxObserved, old, c) {
|
||
break
|
||
}
|
||
}
|
||
defer atomic.AddInt32(&e.concurrent, -1)
|
||
<-e.release
|
||
return driverResult{}, nil
|
||
}
|
||
|
||
func (e *blockingExecutor) QueryContext(_ context.Context, _ string, _ ...any) (*sql.Rows, error) {
|
||
atomic.AddInt32(&e.queryCalls, 1)
|
||
return nil, sql.ErrNoRows
|
||
}
|
||
|
||
// driverResult 是一个零值 sql.Result，用于测试。
|
||
type driverResult struct{}
|
||
|
||
func (driverResult) LastInsertId() (int64, error) { return 0, nil }
|
||
func (driverResult) RowsAffected() (int64, error) { return 1, nil }
|
||
|
||
func TestSetTempUnschedulable_SingleflightDedupesConcurrentCallers(t *testing.T) {
|
||
// 同一账号 + 同一 until + 同一 reason 的 N 个并发调用，应只触发一次实际
|
||
// SQL 路径（UPDATE + outbox INSERT = 2 次 ExecContext）。
|
||
exec := newBlockingExecutor()
|
||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||
|
||
const callers = 30
|
||
until := time.Now().Add(10 * time.Minute)
|
||
const reason = "OAuth 401: invalid_grant"
|
||
|
||
var wg sync.WaitGroup
|
||
wg.Add(callers)
|
||
for i := 0; i < callers; i++ {
|
||
go func() {
|
||
defer wg.Done()
|
||
_ = repo.SetTempUnschedulable(context.Background(), 42, until, reason)
|
||
}()
|
||
}
|
||
|
||
// 等首个 ExecContext 进入阻塞，确认 sf 已聚拢调用。
|
||
deadline := time.Now().Add(2 * time.Second)
|
||
for atomic.LoadInt32(&exec.concurrent) == 0 && time.Now().Before(deadline) {
|
||
time.Sleep(5 * time.Millisecond)
|
||
}
|
||
require.Equal(t, int32(1), atomic.LoadInt32(&exec.concurrent),
|
||
"singleflight should serialize the SQL call to exactly one in-flight execution")
|
||
|
||
exec.Release()
|
||
wg.Wait()
|
||
|
||
// 1 次 UPDATE + 1 次 outbox INSERT = 2 次 exec；其余 29 个 caller 共享结果。
|
||
require.LessOrEqual(t, atomic.LoadInt32(&exec.execCalls), int32(2),
|
||
"expected at most 2 ExecContext calls (UPDATE + outbox), got %d", exec.execCalls)
|
||
require.Equal(t, int32(1), atomic.LoadInt32(&exec.maxObserved),
|
||
"no two SQL execs should run concurrently for the same singleflight key")
|
||
}
|
||
|
||
func TestSetTempUnschedulable_DifferentAccountsRunInParallel(t *testing.T) {
|
||
// 不同 account 应分属不同 sf key，能并行写库。
|
||
exec := newBlockingExecutor()
|
||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||
|
||
until := time.Now().Add(10 * time.Minute)
|
||
var wg sync.WaitGroup
|
||
for i := int64(1); i <= 3; i++ {
|
||
i := i
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
_ = repo.SetTempUnschedulable(context.Background(), i, until, "different reason")
|
||
}()
|
||
}
|
||
|
||
deadline := time.Now().Add(2 * time.Second)
|
||
for atomic.LoadInt32(&exec.concurrent) < 3 && time.Now().Before(deadline) {
|
||
time.Sleep(5 * time.Millisecond)
|
||
}
|
||
require.Equal(t, int32(3), atomic.LoadInt32(&exec.maxObserved),
|
||
"different accounts should be able to write in parallel")
|
||
|
||
exec.Release()
|
||
wg.Wait()
|
||
}
|