fix(retry): guard RetryBudget mutable config against data race

UpdateConfig writes tokensPerSecond/maxTokens/enabled as plain fields while
the refill goroutine and AllowRetry readers touch them concurrently. Guard
all reads/writes of the mutable config with an RWMutex.

Adds TestRetryBudget_ConcurrentUpdateConfig, which races UpdateConfig
against AllowRetry/GetStats and is verified clean under -race.

Also drop the package-level min() that shadowed the Go 1.25 builtin.
This commit is contained in:
2026-08-13 01:20:55 +01:00
parent 6fd8f99b7e
commit 963f72ff17
3 changed files with 80 additions and 15 deletions
-7
View File
@@ -234,10 +234,3 @@ func InitializeRequestCoalescer(enabled bool, logger *libpack_logger.Logger, mon
func GetRequestCoalescer() *RequestCoalescer { func GetRequestCoalescer() *RequestCoalescer {
return requestCoalescer return requestCoalescer
} }
func min(a, b int) int {
if a < b {
return a
}
return b
}
+30 -8
View File
@@ -17,6 +17,7 @@ type RetryBudget struct {
currentTokens atomic.Int64 currentTokens atomic.Int64
lastRefill atomic.Int64 // Unix timestamp in nanoseconds lastRefill atomic.Int64 // Unix timestamp in nanoseconds
enabled bool enabled bool
configMu sync.RWMutex // guards tokensPerSecond, maxTokens, enabled (UpdateConfig vs refill/AllowRetry)
logger *libpack_logger.Logger logger *libpack_logger.Logger
ctx context.Context ctx context.Context
cancel context.CancelFunc cancel context.CancelFunc
@@ -67,7 +68,11 @@ func NewRetryBudgetWithContext(ctx context.Context, config RetryBudgetConfig, lo
func (rb *RetryBudget) AllowRetry() bool { func (rb *RetryBudget) AllowRetry() bool {
rb.totalAttempts.Add(1) rb.totalAttempts.Add(1)
if !rb.enabled { rb.configMu.RLock()
enabled := rb.enabled
rb.configMu.RUnlock()
if !enabled {
rb.allowedRetries.Add(1) rb.allowedRetries.Add(1)
return true return true
} }
@@ -120,6 +125,11 @@ func (rb *RetryBudget) Shutdown() {
// refill adds tokens to the bucket based on elapsed time // refill adds tokens to the bucket based on elapsed time
func (rb *RetryBudget) refill() { func (rb *RetryBudget) refill() {
rb.configMu.RLock()
tokensPerSecond := rb.tokensPerSecond
maxTokens := rb.maxTokens
rb.configMu.RUnlock()
now := time.Now().UnixNano() now := time.Now().UnixNano()
last := rb.lastRefill.Load() last := rb.lastRefill.Load()
@@ -127,7 +137,7 @@ func (rb *RetryBudget) refill() {
elapsed := float64(now-last) / float64(time.Second) elapsed := float64(now-last) / float64(time.Second)
// Calculate tokens to add // Calculate tokens to add
tokensToAdd := int64(elapsed * rb.tokensPerSecond) tokensToAdd := int64(elapsed * tokensPerSecond)
if tokensToAdd > 0 { if tokensToAdd > 0 {
// Update last refill time // Update last refill time
@@ -136,8 +146,8 @@ func (rb *RetryBudget) refill() {
for { for {
current := rb.currentTokens.Load() current := rb.currentTokens.Load()
newValue := current + tokensToAdd newValue := current + tokensToAdd
if newValue > rb.maxTokens { if newValue > maxTokens {
newValue = rb.maxTokens newValue = maxTokens
} }
if rb.currentTokens.CompareAndSwap(current, newValue) { if rb.currentTokens.CompareAndSwap(current, newValue) {
@@ -154,16 +164,22 @@ func (rb *RetryBudget) GetStats() map[string]any {
allowedRetries := rb.allowedRetries.Load() allowedRetries := rb.allowedRetries.Load()
deniedRetries := rb.deniedRetries.Load() deniedRetries := rb.deniedRetries.Load()
rb.configMu.RLock()
enabled := rb.enabled
maxTokens := rb.maxTokens
tokensPerSecond := rb.tokensPerSecond
rb.configMu.RUnlock()
var denialRate float64 var denialRate float64
if totalAttempts > 0 { if totalAttempts > 0 {
denialRate = float64(deniedRetries) / float64(totalAttempts) * 100 denialRate = float64(deniedRetries) / float64(totalAttempts) * 100
} }
return map[string]any{ return map[string]any{
"enabled": rb.enabled, "enabled": enabled,
"current_tokens": rb.currentTokens.Load(), "current_tokens": rb.currentTokens.Load(),
"max_tokens": rb.maxTokens, "max_tokens": maxTokens,
"tokens_per_sec": rb.tokensPerSecond, "tokens_per_sec": tokensPerSecond,
"total_attempts": totalAttempts, "total_attempts": totalAttempts,
"allowed_retries": allowedRetries, "allowed_retries": allowedRetries,
"denied_retries": deniedRetries, "denied_retries": deniedRetries,
@@ -173,20 +189,26 @@ func (rb *RetryBudget) GetStats() map[string]any {
// Reset resets the retry budget statistics // Reset resets the retry budget statistics
func (rb *RetryBudget) Reset() { func (rb *RetryBudget) Reset() {
rb.configMu.RLock()
maxTokens := rb.maxTokens
rb.configMu.RUnlock()
rb.totalAttempts.Store(0) rb.totalAttempts.Store(0)
rb.allowedRetries.Store(0) rb.allowedRetries.Store(0)
rb.deniedRetries.Store(0) rb.deniedRetries.Store(0)
rb.currentTokens.Store(rb.maxTokens) rb.currentTokens.Store(maxTokens)
} }
// UpdateConfig updates the retry budget configuration // UpdateConfig updates the retry budget configuration
func (rb *RetryBudget) UpdateConfig(config RetryBudgetConfig) { func (rb *RetryBudget) UpdateConfig(config RetryBudgetConfig) {
rb.configMu.Lock()
rb.tokensPerSecond = config.TokensPerSecond rb.tokensPerSecond = config.TokensPerSecond
rb.maxTokens = int64(config.MaxTokens) rb.maxTokens = int64(config.MaxTokens)
rb.enabled = config.Enabled rb.enabled = config.Enabled
// Reset to full capacity // Reset to full capacity
rb.currentTokens.Store(rb.maxTokens) rb.currentTokens.Store(rb.maxTokens)
rb.configMu.Unlock()
if rb.logger != nil { if rb.logger != nil {
rb.logger.Info(&libpack_logger.LogMessage{ rb.logger.Info(&libpack_logger.LogMessage{
+50
View File
@@ -1,6 +1,7 @@
package main package main
import ( import (
"sync"
"testing" "testing"
"time" "time"
@@ -310,3 +311,52 @@ func TestRetryBudget_GlobalInstance(t *testing.T) {
rb2 := GetRetryBudget() rb2 := GetRetryBudget()
assert.Equal(t, rb, rb2) assert.Equal(t, rb, rb2)
} }
// TestRetryBudget_ConcurrentUpdateConfig exercises UpdateConfig racing with
// AllowRetry/refill readers. Run with -race to detect the data race on the
// mutable config fields.
func TestRetryBudget_ConcurrentUpdateConfig(t *testing.T) {
rb := NewRetryBudget(RetryBudgetConfig{TokensPerSecond: 100.0, MaxTokens: 1000, Enabled: true}, nil)
rb.cancel() // stop the refill goroutine so it doesn't advance tokens
rb.currentTokens.Store(1000)
var wg sync.WaitGroup
stop := make(chan struct{})
// Concurrent writers calling UpdateConfig
for i := 0; i < 4; i++ {
wg.Add(1)
go func(n int) {
defer wg.Done()
for {
select {
case <-stop:
return
default:
rb.UpdateConfig(RetryBudgetConfig{TokensPerSecond: float64(n), MaxTokens: n, Enabled: true})
}
}
}(i)
}
// Concurrent readers calling AllowRetry and GetStats
for i := 0; i < 4; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for {
select {
case <-stop:
return
default:
rb.AllowRetry()
rb.GetStats()
}
}
}()
}
time.Sleep(100 * time.Millisecond)
close(stop)
wg.Wait()
}