A budget spent on the server's block no longer blames the client limit

internal/resilience/gate.go

@@ -28,8 +28,10 @@ func (e *GateError) Error() string { return e.Message }
 // Unwrap returns the SDK sentinel the rejection stands for.
 func (e *GateError) Unwrap() error { return e.sentinel }
 
-// pause sleeps for d or until ctx is done, whichever comes first.
-func pause(ctx context.Context, d time.Duration) error {
+// pause sleeps for d or until ctx is done, whichever comes first. A
+// variable, like jitter, so that a test can wake a gate on its own terms
+// rather than waiting on a real timer and hoping about the wake-up.
+var pause = func(ctx context.Context, d time.Duration) error {
 	timer := time.NewTimer(d)
 	defer timer.Stop()
 	select {

internal/resilience/gate_test.go

@@ -367,6 +367,113 @@ func TestRateLimiterWaitRoundsTheRetryAfterUp(t *testing.T) {
 	assert.Equal(t, "Wait 41s, then re-run.", gateErr.Hint)
 }
 
+// heldClock is a clock a test moves, and the gate's own sleeps move it:
+// pause does not wait, it advances the clock by the sleep it was asked for
+// plus the lateness every real wake-up has. A gate only ever reaches a spent
+// budget by waking past its deadline, and that is the one thing a test
+// cannot ask a real timer for — the 20ms margin measured on an idle box and
+// spent by the scheduler on a loaded one is the defect this card came from.
+// Here the overshoot is a fact of the test.
+type heldClock struct {
+	mu sync.Mutex
+	t  time.Time
+}
+
+// holdClock freezes time and hands every sleep the given lateness. The
+// jitter goes to zero with it, so a sleep for a wait is exactly that wait.
+func holdClock(t *testing.T, lateness time.Duration) *heldClock {
+	t.Helper()
+	clock := &heldClock{t: time.Now()}
+
+	previousJitter, previousPause := jitter, pause
+	jitter = func(time.Duration) time.Duration { return 0 }
+	pause = func(ctx context.Context, d time.Duration) error {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		clock.advance(d + lateness)
+		return nil
+	}
+	t.Cleanup(func() { jitter, pause = previousJitter, previousPause })
+
+	return clock
+}
+
+func (c *heldClock) Now() time.Time {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.t
+}
+
+func (c *heldClock) advance(d time.Duration) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.t = c.t.Add(d)
+}
+
+// A budget spent on the server's block says so. Blaming the client limit
+// here reads as "lower your parallelism", which is advice about a knob that
+// had nothing to do with a wait the server asked of everybody.
+//
+// The gate sleeps a 10s block out against a 10s budget and wakes a
+// millisecond late, which is the shape the CI failure in #763 had and the
+// only way a spent budget is ever reached on a block.
+func TestRateLimiterBudgetSpentOnTheServersBlockNamesTheServer(t *testing.T) {
+	clock := holdClock(t, time.Millisecond)
+	rl := NewRateLimiter(NewStore(t.TempDir()), RateLimiterConfig{})
+	rl.clock = clock.Now
+	start := clock.Now()
+	require.NoError(t, rl.SetRetryAfter(start.Add(DefaultMaxWait)))
+
+	err := rl.waitSince(context.Background(), start, start.Add(DefaultMaxWait))
+
+	var gateErr *GateError
+	require.ErrorAs(t, err, &gateErr)
+	assert.ErrorIs(t, err, basecamp.ErrRateLimited)
+	assert.Equal(t, "Rate limited by the server; waited 10s", gateErr.Message)
+	assert.Equal(t, "Re-run.", gateErr.Hint)
+	assert.Equal(t, start.Add(DefaultMaxWait+time.Millisecond), clock.Now(), "slept the block out and woke late")
+}
+
+// The store is shared, so the Retry-After it holds when a gate gives up is
+// no evidence that this gate waited on one. Here another invocation's 429
+// lands while this gate sleeps for a refill of its own — the write happens
+// inside the sleep, so the order is the test's and not the scheduler's — and
+// the wait our own bucket imposed keeps its own name, and the advice that
+// goes with it.
+func TestRateLimiterBudgetSpentOnOurOwnRefillNamesTheClientLimit(t *testing.T) {
+	clock := holdClock(t, time.Millisecond)
+	store := NewStore(t.TempDir())
+	rl := NewRateLimiter(store, RateLimiterConfig{MaxTokens: 1, RefillRate: 1, TokensPerRequest: 1})
+	rl.clock = clock.Now
+	start := clock.Now()
+
+	allowed, err := rl.Allow()
+	require.NoError(t, err)
+	require.True(t, allowed, "the bucket's one token")
+
+	sleep := pause
+	pause = func(ctx context.Context, d time.Duration) error {
+		//nolint:contextcheck // lock acquisition is context-independent by design
+		require.NoError(t, rl.SetRetryAfterDuration(30*time.Second), "someone else's 429, mid-sleep")
+		return sleep(ctx, d)
+	}
+
+	// The budget ends with the refill of the token this gate is waiting for.
+	err = rl.waitSince(context.Background(), start, start.Add(time.Second))
+
+	var gateErr *GateError
+	require.ErrorAs(t, err, &gateErr)
+	assert.ErrorIs(t, err, basecamp.ErrRateLimited)
+	assert.Equal(t, "Too many requests (client limit 1/s); waited 1s", gateErr.Message)
+	assert.Equal(t, "Re-run, or lower parallelism.", gateErr.Hint)
+
+	state, err := store.Load()
+	require.NoError(t, err)
+	require.True(t, state.RateLimiter.RetryAfterUntil.After(start),
+		"the block really did land, and a gate reading the store back would have called this the server's")
+}
+
 func TestCeilSeconds(t *testing.T) {
 	assert.Equal(t, 41*time.Second, ceilSeconds(40*time.Second+time.Millisecond))
 	assert.Equal(t, 40*time.Second, ceilSeconds(40*time.Second))

internal/resilience/rate_limiter.go

@@ -12,6 +12,9 @@ import (
 type RateLimiter struct {
 	config RateLimiterConfig
 	store  *Store
+	// clock is the time the limiter reads. A field so that a test can put
+	// the deadline boundary where it wants it instead of racing a real one.
+	clock func() time.Time
 }
 
 // NewRateLimiter creates a new rate limiter with the given config.
@@ -30,12 +33,13 @@ func NewRateLimiter(store *Store, config RateLimiterConfig) *RateLimiter {
 	return &RateLimiter{
 		config: config,
 		store:  store,
+		clock:  time.Now,
 	}
 }
 
 // now returns the current time.
 func (rl *RateLimiter) now() time.Time {
-	return time.Now()
+	return rl.clock()
 }
 
 // refill adds tokens based on elapsed time since last refill.
@@ -78,9 +82,11 @@ func (rl *RateLimiter) take() (allowed bool, wait time.Duration, blocked bool) {
 		rlState := &state.RateLimiter
 		now := rl.now()
 
-		// Check Retry-After block
-		if rlState.IsBlocked() {
-			allowed, blocked, wait = false, true, rlState.BlockedFor()
+		// Check Retry-After block, against the transaction's own now the
+		// way the refill below is, and not against a second reading of the
+		// clock inside the state.
+		if blockEnds := rlState.RetryAfterUntil; blockEnds.After(now) {
+			allowed, blocked, wait = false, true, blockEnds.Sub(now)
 			return nil
 		}
 
@@ -128,23 +134,28 @@ func sleepWithin(wait, remaining time.Duration) time.Duration {
 // Retry-After block to lift, until deadline. It returns nil when the request
 // may proceed, a *GateError when the deadline would pass first, or ctx.Err().
 // A Retry-After block that outlasts the deadline is reported immediately
-// rather than waited on, with the remaining time in the message.
+// rather than waited on, with the remaining time in the message, and a
+// budget spent queueing names the limit that held the gate: the server's
+// block, or our own bucket.
 // Cancellation and the deadline are checked before every attempt, so an
 // expired or canceled gate consumes nothing, and cancellation outranks the
 // deadline.
 func (rl *RateLimiter) Wait(ctx context.Context, deadline time.Time) error {
 	return rl.waitSince(ctx, rl.now(), deadline)
 }
 
-// waitSince is Wait for a gate that started queueing at start.
+// waitSince is Wait for a gate that started queueing at start. It carries
+// the cause of the sleep it is in, so that a budget which runs out can name
+// what it was spent on.
 func (rl *RateLimiter) waitSince(ctx context.Context, start, deadline time.Time) error {
+	sleptOnServerBlock := false
 	for {
 		if err := ctx.Err(); err != nil {
 			return err
 		}
 		remaining := deadline.Sub(rl.now())
 		if remaining <= 0 {
-			return rl.gateError(false, 0, rl.now().Sub(start))
+			return rl.budgetError(sleptOnServerBlock, rl.now().Sub(start))
 		}
 		allowed, wait, blocked := rl.take() //nolint:contextcheck // lock acquisition is context-independent by design
 		if allowed {
@@ -153,6 +164,7 @@ func (rl *RateLimiter) waitSince(ctx context.Context, start, deadline time.Time)
 		if wait > remaining {
 			return rl.gateError(blocked, wait, rl.now().Sub(start))
 		}
+		sleptOnServerBlock = blocked
 		if err := pause(ctx, sleepWithin(wait, remaining)); err != nil {
 			return err
 		}
@@ -170,6 +182,33 @@ func (rl *RateLimiter) gateError(blocked bool, wait, waited time.Duration) *Gate
 			sentinel: basecamp.ErrRateLimited,
 		}
 	}
+	return rl.clientLimitError(waited)
+}
+
+// budgetError is the rejection for a gate whose budget ran out while it was
+// still sleeping, and sleptOnServerBlock is what that last sleep was for, as
+// take saw it. A server Retry-After is the server's doing and nobody's
+// parallelism; anything else is our own bucket.
+//
+// The cause is carried out of the sleep rather than read back from the store
+// here, because the store is shared and by now says something else: another
+// invocation's 429 can land while this gate sleeps for a refill of its own,
+// and a block this gate never waited on can expire between the two. Either
+// one would hand a client-side timeout the server's name.
+func (rl *RateLimiter) budgetError(sleptOnServerBlock bool, waited time.Duration) *GateError {
+	if !sleptOnServerBlock {
+		return rl.clientLimitError(waited)
+	}
+	return &GateError{
+		Message:  fmt.Sprintf("Rate limited by the server; waited %s", waited.Round(time.Second)),
+		Hint:     "Re-run.",
+		sentinel: basecamp.ErrRateLimited,
+	}
+}
+
+// clientLimitError is the rejection for a wait our own token bucket imposed,
+// reported as the request rate the bucket allows.
+func (rl *RateLimiter) clientLimitError(waited time.Duration) *GateError {
 	requestsPerSecond := rl.config.RefillRate / rl.config.TokensPerRequest
 	return &GateError{
 		Message:  fmt.Sprintf("Too many requests (client limit %g/s); waited %s", requestsPerSecond, waited.Round(time.Second)),