diff --git a/request_coalescing.go b/request_coalescing.go index 0d91346..411a503 100644 --- a/request_coalescing.go +++ b/request_coalescing.go @@ -234,10 +234,3 @@ func InitializeRequestCoalescer(enabled bool, logger *libpack_logger.Logger, mon func GetRequestCoalescer() *RequestCoalescer { return requestCoalescer } - -func min(a, b int) int { - if a < b { - return a - } - return b -} diff --git a/retry_budget.go b/retry_budget.go index e04cd2e..240a3a5 100644 --- a/retry_budget.go +++ b/retry_budget.go @@ -17,6 +17,7 @@ type RetryBudget struct { currentTokens atomic.Int64 lastRefill atomic.Int64 // Unix timestamp in nanoseconds enabled bool + configMu sync.RWMutex // guards tokensPerSecond, maxTokens, enabled (UpdateConfig vs refill/AllowRetry) logger *libpack_logger.Logger ctx context.Context cancel context.CancelFunc @@ -67,7 +68,11 @@ func NewRetryBudgetWithContext(ctx context.Context, config RetryBudgetConfig, lo func (rb *RetryBudget) AllowRetry() bool { rb.totalAttempts.Add(1) - if !rb.enabled { + rb.configMu.RLock() + enabled := rb.enabled + rb.configMu.RUnlock() + + if !enabled { rb.allowedRetries.Add(1) return true } @@ -120,6 +125,11 @@ func (rb *RetryBudget) Shutdown() { // refill adds tokens to the bucket based on elapsed time func (rb *RetryBudget) refill() { + rb.configMu.RLock() + tokensPerSecond := rb.tokensPerSecond + maxTokens := rb.maxTokens + rb.configMu.RUnlock() + now := time.Now().UnixNano() last := rb.lastRefill.Load() @@ -127,7 +137,7 @@ func (rb *RetryBudget) refill() { elapsed := float64(now-last) / float64(time.Second) // Calculate tokens to add - tokensToAdd := int64(elapsed * rb.tokensPerSecond) + tokensToAdd := int64(elapsed * tokensPerSecond) if tokensToAdd > 0 { // Update last refill time @@ -136,8 +146,8 @@ func (rb *RetryBudget) refill() { for { current := rb.currentTokens.Load() newValue := current + tokensToAdd - if newValue > rb.maxTokens { - newValue = rb.maxTokens + if newValue > maxTokens { + newValue = maxTokens } if rb.currentTokens.CompareAndSwap(current, newValue) { @@ -154,16 +164,22 @@ func (rb *RetryBudget) GetStats() map[string]any { allowedRetries := rb.allowedRetries.Load() deniedRetries := rb.deniedRetries.Load() + rb.configMu.RLock() + enabled := rb.enabled + maxTokens := rb.maxTokens + tokensPerSecond := rb.tokensPerSecond + rb.configMu.RUnlock() + var denialRate float64 if totalAttempts > 0 { denialRate = float64(deniedRetries) / float64(totalAttempts) * 100 } return map[string]any{ - "enabled": rb.enabled, + "enabled": enabled, "current_tokens": rb.currentTokens.Load(), - "max_tokens": rb.maxTokens, - "tokens_per_sec": rb.tokensPerSecond, + "max_tokens": maxTokens, + "tokens_per_sec": tokensPerSecond, "total_attempts": totalAttempts, "allowed_retries": allowedRetries, "denied_retries": deniedRetries, @@ -173,20 +189,26 @@ func (rb *RetryBudget) GetStats() map[string]any { // Reset resets the retry budget statistics func (rb *RetryBudget) Reset() { + rb.configMu.RLock() + maxTokens := rb.maxTokens + rb.configMu.RUnlock() + rb.totalAttempts.Store(0) rb.allowedRetries.Store(0) rb.deniedRetries.Store(0) - rb.currentTokens.Store(rb.maxTokens) + rb.currentTokens.Store(maxTokens) } // UpdateConfig updates the retry budget configuration func (rb *RetryBudget) UpdateConfig(config RetryBudgetConfig) { + rb.configMu.Lock() rb.tokensPerSecond = config.TokensPerSecond rb.maxTokens = int64(config.MaxTokens) rb.enabled = config.Enabled // Reset to full capacity rb.currentTokens.Store(rb.maxTokens) + rb.configMu.Unlock() if rb.logger != nil { rb.logger.Info(&libpack_logger.LogMessage{ diff --git a/retry_budget_test.go b/retry_budget_test.go index b925767..3e80ceb 100644 --- a/retry_budget_test.go +++ b/retry_budget_test.go @@ -1,6 +1,7 @@ package main import ( + "sync" "testing" "time" @@ -310,3 +311,52 @@ func TestRetryBudget_GlobalInstance(t *testing.T) { rb2 := GetRetryBudget() 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() +}