fix(retry): start refill loop when budget enabled via UpdateConfig

The refill goroutine started only when the budget was enabled at
construction. Enabling it later through UpdateConfig left a bucket that
never refilled.

Start the refill loop from UpdateConfig, guarded by sync.Once so it
starts at most once, sharing the existing shutdown lifecycle.
This commit is contained in:
2026-08-14 01:42:04 +01:00
parent 6226717aad
commit 77447bd353
2 changed files with 94 additions and 1 deletions
+24 -1
View File
@@ -21,6 +21,8 @@ type RetryBudget struct {
logger *libpack_logger.Logger
ctx context.Context
cancel context.CancelFunc
refillOnce sync.Once // ensures refillLoop starts at most once, whether from construction or a later UpdateConfig
refillStarts atomic.Int32 // counts startRefillLoop executions; used by tests to assert no duplicate refiller
// Statistics
totalAttempts atomic.Int64
@@ -58,12 +60,25 @@ func NewRetryBudgetWithContext(ctx context.Context, config RetryBudgetConfig, lo
// Start refill goroutine
if rb.enabled {
go rb.refillLoop()
rb.startRefillLoop()
}
return rb
}
// startRefillLoop starts the background refill goroutine, at most once,
// regardless of how many times it is called from construction or UpdateConfig.
// Shutdown is terminal: once rb.ctx is cancelled, a later call here still
// starts refillLoop, but it exits immediately on the cancelled context
// without ever refilling, so enabling the budget again after Shutdown does
// not bring back a working refiller.
func (rb *RetryBudget) startRefillLoop() {
rb.refillOnce.Do(func() {
rb.refillStarts.Add(1)
go rb.refillLoop()
})
}
// AllowRetry checks if a retry is allowed based on the current budget
func (rb *RetryBudget) AllowRetry() bool {
rb.totalAttempts.Add(1)
@@ -210,6 +225,14 @@ func (rb *RetryBudget) UpdateConfig(config RetryBudgetConfig) {
rb.currentTokens.Store(rb.maxTokens)
rb.configMu.Unlock()
// Enabling via UpdateConfig must (re)start the refill loop, otherwise a
// budget constructed disabled and later enabled would never refill.
// startRefillLoop is idempotent, so re-enabling an already-running budget
// is a no-op here.
if config.Enabled {
rb.startRefillLoop()
}
if rb.logger != nil {
rb.logger.Info(&libpack_logger.LogMessage{
Message: "Retry budget configuration updated",