graphql-monitoring-proxy/retry_budget.go

package main

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	libpack_logger "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
)

// RetryBudget implements a token bucket algorithm to limit the rate of retries
// This prevents retry storms and cascading failures
type RetryBudget struct {
	tokensPerSecond float64
	maxTokens       int64
	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
	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
	allowedRetries atomic.Int64
	deniedRetries  atomic.Int64
}

// RetryBudgetConfig holds configuration for retry budget
type RetryBudgetConfig struct {
	TokensPerSecond float64 // Rate at which tokens are refilled
	MaxTokens       int     // Maximum number of tokens (burst capacity)
	Enabled         bool    // Whether retry budget is enabled
}

// NewRetryBudget creates a new retry budget (deprecated, use NewRetryBudgetWithContext)
func NewRetryBudget(config RetryBudgetConfig, logger *libpack_logger.Logger) *RetryBudget {
	return NewRetryBudgetWithContext(context.Background(), config, logger)
}

// NewRetryBudgetWithContext creates a new retry budget with context for graceful shutdown
func NewRetryBudgetWithContext(ctx context.Context, config RetryBudgetConfig, logger *libpack_logger.Logger) *RetryBudget {
	budgetCtx, cancel := context.WithCancel(ctx)
	rb := &RetryBudget{
		tokensPerSecond: config.TokensPerSecond,
		maxTokens:       int64(config.MaxTokens),
		enabled:         config.Enabled,
		logger:          logger,
		ctx:             budgetCtx,
		cancel:          cancel,
	}

	// Initialize with full bucket
	rb.currentTokens.Store(rb.maxTokens)
	rb.lastRefill.Store(time.Now().UnixNano())

	// Start refill goroutine
	if rb.enabled {
		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)

	rb.configMu.RLock()
	enabled := rb.enabled
	rb.configMu.RUnlock()

	if !enabled {
		rb.allowedRetries.Add(1)
		return true
	}

	// Try to consume a token
	for {
		current := rb.currentTokens.Load()
		if current <= 0 {
			rb.deniedRetries.Add(1)
			if rb.logger != nil {
				rb.logger.Debug(&libpack_logger.LogMessage{
					Message: "Retry denied: budget exhausted",
					Pairs: map[string]any{
						"current_tokens": current,
						"denied_count":   rb.deniedRetries.Load(),
					},
				})
			}
			return false
		}

		if rb.currentTokens.CompareAndSwap(current, current-1) {
			rb.allowedRetries.Add(1)
			return true
		}
	}
}

// refillLoop periodically refills tokens
func (rb *RetryBudget) refillLoop() {
	ticker := time.NewTicker(100 * time.Millisecond) // Refill every 100ms
	defer ticker.Stop()

	for {
		select {
		case <-rb.ctx.Done():
			return
		case <-ticker.C:
			rb.refill()
		}
	}
}

// Shutdown stops the retry budget goroutine
func (rb *RetryBudget) Shutdown() {
	if rb.cancel != nil {
		rb.cancel()
	}
}

// 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()

	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
			}

			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()
	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":         enabled,
		"current_tokens":  rb.currentTokens.Load(),
		"max_tokens":      maxTokens,
		"tokens_per_sec":  tokensPerSecond,
		"total_attempts":  totalAttempts,
		"allowed_retries": allowedRetries,
		"denied_retries":  deniedRetries,
		"denial_rate_pct": denialRate,
	}
}

// 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(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()

	// 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",
			Pairs: map[string]any{
				"tokens_per_sec": config.TokensPerSecond,
				"max_tokens":     config.MaxTokens,
				"enabled":        config.Enabled,
			},
		})
	}
}

// Global retry budget instance
var (
	retryBudget     *RetryBudget
	retryBudgetOnce sync.Once
)

// InitializeRetryBudget initializes the global retry budget (deprecated, use InitializeRetryBudgetWithContext)
func InitializeRetryBudget(config RetryBudgetConfig, logger *libpack_logger.Logger) *RetryBudget {
	return InitializeRetryBudgetWithContext(context.Background(), config, logger)
}

// InitializeRetryBudgetWithContext initializes the global retry budget with context for graceful shutdown
func InitializeRetryBudgetWithContext(ctx context.Context, config RetryBudgetConfig, logger *libpack_logger.Logger) *RetryBudget {
	retryBudgetOnce.Do(func() {
		retryBudget = NewRetryBudgetWithContext(ctx, config, logger)
		if logger != nil && config.Enabled {
			logger.Info(&libpack_logger.LogMessage{
				Message: "Retry budget initialized",
				Pairs: map[string]any{
					"tokens_per_sec": config.TokensPerSecond,
					"max_tokens":     config.MaxTokens,
				},
			})
		}
	})
	return retryBudget
}

// GetRetryBudget returns the global retry budget instance
func GetRetryBudget() *RetryBudget {
	return retryBudget
}