diff --git a/.gitignore b/.gitignore index 8d43501..889ed57 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,6 @@ dist/ coverage.out CLAUDE.md graphql-monitoring-proxy + +# Lean proofs: local only, never committed +theorems/ diff --git a/retry_budget.go b/retry_budget.go index 5cf6d6a..a15f88e 100644 --- a/retry_budget.go +++ b/retry_budget.go @@ -148,31 +148,39 @@ func (rb *RetryBudget) refill() { now := time.Now().UnixNano() last := rb.lastRefill.Load() - // Calculate elapsed time in seconds - elapsed := float64(now-last) / float64(time.Second) + tokensToAdd, newLast := refillStep(tokensPerSecond, last, now) + if tokensToAdd <= 0 { + return + } + if rb.lastRefill.CompareAndSwap(last, newLast) { + for { + current := rb.currentTokens.Load() + newValue := current + tokensToAdd + if newValue > maxTokens { + newValue = maxTokens + } - // Calculate tokens to add - tokensToAdd := int64(elapsed * tokensPerSecond) - - if tokensToAdd > 0 { - // Update last refill time - if rb.lastRefill.CompareAndSwap(last, now) { - // Add tokens, capped at maxTokens - for { - current := rb.currentTokens.Load() - newValue := current + tokensToAdd - if newValue > maxTokens { - newValue = maxTokens - } - - if rb.currentTokens.CompareAndSwap(current, newValue) { - break - } + if rb.currentTokens.CompareAndSwap(current, newValue) { + break } } } } +// refillStep returns the whole tokens earned between last and now (ns) and the +// clock to store. The clock advances only by the time those tokens pay for, so +// the fractional remainder carries into the next tick. +func refillStep(tokensPerSecond float64, last, now int64) (int64, int64) { + if tokensPerSecond <= 0 || now <= last { + return 0, last + } + tokens := int64(float64(now-last) / float64(time.Second) * tokensPerSecond) + if tokens <= 0 { + return 0, last + } + return tokens, last + int64(float64(tokens)/tokensPerSecond*float64(time.Second)) +} + // GetStats returns current statistics func (rb *RetryBudget) GetStats() map[string]any { totalAttempts := rb.totalAttempts.Load() diff --git a/retry_budget_refill_test.go b/retry_budget_refill_test.go new file mode 100644 index 0000000..4bcf064 --- /dev/null +++ b/retry_budget_refill_test.go @@ -0,0 +1,47 @@ +package main + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +// Pins theorems/Theorems/RetryRefill.lean: total_added(15 tps, 100ms ticks, 10s) = 150. +func TestRefillStep_KeepsFractionalTokens(t *testing.T) { + tests := []struct { + name string + tps float64 + want int64 + }{ + {"15 tps", 15, 150}, + {"25 tps", 25, 250}, + {"5 tps (sub-token per tick)", 5, 50}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tick := 100 * time.Millisecond + var last, total int64 + now := int64(tick) + for i := 0; i < 100; i++ { + added, newLast := refillStep(tt.tps, last, now) + total += added + last = newLast + now += int64(tick) + } + assert.Equal(t, tt.want, total) + }) + } +} + +func TestRefillStep_NoElapsedNoChange(t *testing.T) { + added, last := refillStep(10, 1000, 1000) + assert.Zero(t, added) + assert.Equal(t, int64(1000), last) +} + +func TestRefillStep_ZeroRateNeverAdds(t *testing.T) { + added, last := refillStep(0, 0, int64(time.Hour)) + assert.Zero(t, added) + assert.Zero(t, last) +}