From 570b8fdb2a301c52d362dfa5901e05c5fc44b088 Mon Sep 17 00:00:00 2001 From: yetone Date: Fri, 9 Oct 2026 23:07:41 +0800 Subject: [PATCH] provider: swapping LoginUsage's or Claude /usage's stand-in waits for the readings still running through the old one (data race in CI run 37942751214) Cause: TestClaudeAccountBackOnceItsWindowRenews (a2c39053) failed the ubuntu race shard of run 37942751214. Its Cleanup's LoginUsageVia(nil) wrote the package-level stand-in while a reading Allowances had started (routing.go:419) was still reading it in loginUsageAt. Allowances reads in the background on purpose and returns without waiting (only the very first ask waits, and not for long), so a reading can outlast the caller, and the test. The window: Allowances deletes c.loading before it tells the renewed hooks, so the test's next poll can start another reading just as the rest it waits on is lifted; the test ends with that reading in flight. In production the stand-in is never set, so there is no production race; the goroutine outliving its caller is the design (it is bounded by its 30 s context). Change: the stand-in is held for reading while it runs, and LoginUsageVia takes it for writing, so it returns once no reading is running through the one it replaces and none starts through it after. That fixes every test that swaps it (gateway claude_renewed and allowance_stale, provider routing, allowance_renewed, codex_reset_deadline, window_renewed), not only the one that failed. UsageClaudeVia's claudeCLIUsage is the same shape (tests swap it, and Allowances' background reading of a Claude account reaches it through readClaudeUsage), so it gets the same treatment; claude_usage_test's restore goes through UsageClaudeVia. ProbeClaudeVia is left: a Claude account test calls it synchronously, nothing reads it in the background. Tests: login_usage_via_test.go holds a call through each stand-in and checks the swap waits for it. Run: - Red before: with a 30 ms delay injected into loginUsageAt and before Allowances tells the renewed hooks (not committed), go test -race -count=20 -run TestClaudeAccountBackOnceItsWindowRenews ./internal/gateway/ printed 7 "WARNING: DATA RACE", the write at login_usage.go:43 from the test's Cleanup and the read at login_usage.go:63 from Allowances.func1, as on CI. Same injection with this change: ok, no race. (-race -count=50 without the injection passed on the old code too; the window is too narrow on this Mac.) - The new tests with the swap made unsynchronised again: both FAIL "swapped while a call was still running through the stand-in it replaced"; with the change: PASS, -race -count=5. - gofmt -l: nothing. go vet (darwin) ./internal/provider ./internal/gateway ./internal/gui; GOOS=linux and GOOS=windows go vet -tags nogui on the same: clean. - Under a temp HOME with pinned caches: go test -tags nogui ./internal/provider ./internal/gateway ok; go test -race ./internal/provider ok; go test -race -count=20 -run '^(TestClaudeAccountBackOnceItsWindowRenews|TestAllowance.*)$' ./internal/gateway ok. - go test -race -tags nogui ./internal/gateway (whole package, temp HOME): ok in 635 s, no race. --- internal/provider/claude_usage_cli.go | 35 +++++++++--- internal/provider/claude_usage_test.go | 8 ++- internal/provider/login_usage.go | 31 ++++++++-- internal/provider/login_usage_via_test.go | 69 +++++++++++++++++++++++ 4 files changed, 127 insertions(+), 16 deletions(-) create mode 100644 internal/provider/login_usage_via_test.go diff --git a/internal/provider/claude_usage_cli.go b/internal/provider/claude_usage_cli.go index 596b6bf6..c72bae0c 100644 --- a/internal/provider/claude_usage_cli.go +++ b/internal/provider/claude_usage_cli.go @@ -11,6 +11,7 @@ import ( "slices" "strconv" "strings" + "sync" "time" // the zone /usage names its resets in: Windows has no zone database, // and without one a reset was read in the machine's own zone @@ -19,19 +20,37 @@ import ( // claudeCLIUsage runs Claude Code's /usage for the account it is signed in // to and is what it printed. The gateway, which runs Claude Code, sets it -// (UsageClaudeVia). -var claudeCLIUsage func(ctx context.Context) (string, error) +// (UsageClaudeVia). It is held for reading while it runs, as +// loginUsageFor is: Allowances reads Claude's usage in the background, and +// tests swap it. +var claudeCLIUsage struct { + sync.RWMutex + f func(ctx context.Context) (string, error) +} -// UsageClaudeVia sets how Claude Code's /usage runs. -func UsageClaudeVia(f func(ctx context.Context) (string, error)) { claudeCLIUsage = f } +// UsageClaudeVia sets how Claude Code's /usage runs. It returns once no +// run is still going through the one it replaces. +func UsageClaudeVia(f func(ctx context.Context) (string, error)) { + claudeCLIUsage.Lock() + claudeCLIUsage.f = f + claudeCLIUsage.Unlock() +} + +// runClaudeUsage is what Claude Code's /usage printed, through +// claudeCLIUsage; errClaudeCannotRun when nothing set it. +func runClaudeUsage(ctx context.Context) (string, error) { + claudeCLIUsage.RLock() + defer claudeCLIUsage.RUnlock() + if claudeCLIUsage.f == nil { + return "", errClaudeCannotRun + } + return claudeCLIUsage.f(ctx) +} // readClaudeUsage is the allowance of the account Claude Code is signed in // to, from its /usage. func readClaudeUsage(ctx context.Context) ([]QuotaWindow, error) { - if claudeCLIUsage == nil { - return []QuotaWindow{}, errClaudeCannotRun - } - text, err := claudeCLIUsage(ctx) + text, err := runClaudeUsage(ctx) if err != nil { return []QuotaWindow{}, err } diff --git a/internal/provider/claude_usage_test.go b/internal/provider/claude_usage_test.go index 6d36b114..156f5892 100644 --- a/internal/provider/claude_usage_test.go +++ b/internal/provider/claude_usage_test.go @@ -18,7 +18,10 @@ import ( func fakeClaudeUsage(t *testing.T, out *atomic.Value, fail *atomic.Bool) *atomic.Int32 { t.Helper() var runs atomic.Int32 - old, oldAsked := claudeCLIUsage, claudeAsked.Load() + claudeCLIUsage.RLock() + old := claudeCLIUsage.f + claudeCLIUsage.RUnlock() + oldAsked := claudeAsked.Load() UsageClaudeVia(func(context.Context) (string, error) { runs.Add(1) if fail != nil && fail.Load() { @@ -38,7 +41,8 @@ func fakeClaudeUsage(t *testing.T, out *atomic.Value, fail *atomic.Bool) *atomic claudeUsage.m = nil claudeUsage.Unlock() t.Cleanup(func() { - claudeCLIUsage, claudeBase, claudeUsageWait, claudeUsedSince = old, oldBase, oldWait, oldUsed + UsageClaudeVia(old) + claudeBase, claudeUsageWait, claudeUsedSince = oldBase, oldWait, oldUsed claudeAsked.Store(oldAsked) srv.Close() claudeUsage.Lock() diff --git a/internal/provider/login_usage.go b/internal/provider/login_usage.go index 31576002..5af5d2e0 100644 --- a/internal/provider/login_usage.go +++ b/internal/provider/login_usage.go @@ -35,12 +35,31 @@ type loginRead struct { } // loginUsageFor, when set, stands in for LoginUsage's readings (tests), -// as UsageClaudeVia does for Claude's own. -var loginUsageFor func(ctx context.Context, agent string) map[string]SubscriptionQuota +// as UsageClaudeVia does for Claude's own. It is held for reading while it +// runs: Allowances reads in the background and returns without waiting, +// so a reading it began can outlast the test that set the stand-in. +var loginUsageFor struct { + sync.RWMutex + f func(ctx context.Context, agent string) map[string]SubscriptionQuota +} -// LoginUsageVia has tests stand in for LoginUsage's readings. +// LoginUsageVia has tests stand in for LoginUsage's readings. It returns +// once no reading is still running through the one it replaces, and none +// starts through it after. func LoginUsageVia(f func(ctx context.Context, agent string) map[string]SubscriptionQuota) { - loginUsageFor = f + loginUsageFor.Lock() + loginUsageFor.f = f + loginUsageFor.Unlock() +} + +// loginUsageStandIn is the stand-in's reading, and false when there is none. +func loginUsageStandIn(ctx context.Context, agent string) (map[string]SubscriptionQuota, bool) { + loginUsageFor.RLock() + defer loginUsageFor.RUnlock() + if loginUsageFor.f == nil { + return nil, false + } + return loginUsageFor.f(ctx, agent), true } // LoginUsage is the allowance used by each of an agent's accounts, by @@ -57,8 +76,8 @@ func LoginUsage(ctx context.Context, agent string) map[string]SubscriptionQuota // counts its own minute from then, not from when it asked (#1295). func loginUsageAt(ctx context.Context, agent string) (map[string]SubscriptionQuota, time.Time) { asked := time.Now() - if loginUsageFor != nil { - return loginUsageFor(ctx, agent), asked + if out, ok := loginUsageStandIn(ctx, agent); ok { + return out, asked } out := map[string]SubscriptionQuota{} if agent == "grok" { diff --git a/internal/provider/login_usage_via_test.go b/internal/provider/login_usage_via_test.go new file mode 100644 index 00000000..b443f6b3 --- /dev/null +++ b/internal/provider/login_usage_via_test.go @@ -0,0 +1,69 @@ +package provider + +import ( + "context" + "testing" + "time" +) + +// swapWaits says whether swap, run while a call through the stand-in it +// replaces is held, waits for that call: Allowances reads in the +// background and returns without waiting, so a reading it began outlasts +// the test that set the stand-in, and the test's swap back raced it (the +// ubuntu race shard of run 37942751214, TestClaudeAccountBackOnceItsWindowRenews). +func swapWaits(t *testing.T, set func(block func()), call func(), swap func()) { + t.Helper() + running, release := make(chan struct{}), make(chan struct{}) + set(func() { + close(running) + <-release + }) + called := make(chan struct{}) + go func() { defer close(called); call() }() + select { + case <-running: + case <-time.After(5 * time.Second): + close(release) + t.Fatal("the stand-in wasn't called") + } + swapped := make(chan struct{}) + go func() { defer close(swapped); swap() }() + select { + case <-swapped: + close(release) + <-called + t.Fatal("swapped while a call was still running through the stand-in it replaced") + case <-time.After(200 * time.Millisecond): + } + close(release) + for _, c := range []chan struct{}{called, swapped} { + select { + case <-c: + case <-time.After(5 * time.Second): + t.Fatal("the swap didn't go on once the call was done") + } + } +} + +func TestLoginUsageViaWaitsForReadingsThroughTheOld(t *testing.T) { + t.Cleanup(func() { LoginUsageVia(nil) }) + swapWaits(t, func(block func()) { + LoginUsageVia(func(context.Context, string) map[string]SubscriptionQuota { + block() + return map[string]SubscriptionQuota{} + }) + }, func() { LoginUsage(context.Background(), "swap-waits") }, func() { LoginUsageVia(nil) }) +} + +func TestUsageClaudeViaWaitsForRunsThroughTheOld(t *testing.T) { + claudeCLIUsage.RLock() + old := claudeCLIUsage.f + claudeCLIUsage.RUnlock() + t.Cleanup(func() { UsageClaudeVia(old) }) + swapWaits(t, func(block func()) { + UsageClaudeVia(func(context.Context) (string, error) { + block() + return "", nil + }) + }, func() { _, _ = readClaudeUsage(context.Background()) }, func() { UsageClaudeVia(nil) }) +}