mirror of
https://github.com/lukaszraczylo/graphql-monitoring-proxy.git
synced 2026-10-22 06:38:13 +00:00
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:
@@ -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
@@ -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{
|
||||||
|
|||||||
@@ -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()
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user