idx-cli/msn/rate_limiter.go

package msn

import (
	"math/rand"
	"sync"
	"time"
)

// RateLimiter implements a token bucket rate limiter with random delay
type RateLimiter struct {
	mu           sync.Mutex
	tokens       float64
	maxTokens    float64
	refillRate   float64 // tokens per second
	lastRefill   time.Time
	minDelayMs   int // minimum delay in milliseconds
	maxDelayMs   int // maximum delay in milliseconds
	requestCount int64
}

// RateLimiterConfig holds rate limiter configuration
type RateLimiterConfig struct {
	RequestsPerSecond float64 // target RPS
	MinDelayMs        int     // minimum random delay
	MaxDelayMs        int     // maximum random delay
}

// NewRateLimiter creates a new rate limiter
func NewRateLimiter(config RateLimiterConfig) *RateLimiter {
	if config.RequestsPerSecond <= 0 {
		config.RequestsPerSecond = 10 // default 10 RPS
	}

	return &RateLimiter{
		tokens:     config.RequestsPerSecond, // start with full bucket
		maxTokens:  config.RequestsPerSecond,
		refillRate: config.RequestsPerSecond,
		lastRefill: time.Now(),
		minDelayMs: config.MinDelayMs,
		maxDelayMs: config.MaxDelayMs,
	}
}

// Wait blocks until a token is available and applies random delay
func (r *RateLimiter) Wait() {
	r.mu.Lock()
	defer r.mu.Unlock()

	// Refill tokens based on elapsed time
	now := time.Now()
	elapsed := now.Sub(r.lastRefill).Seconds()
	r.tokens += elapsed * r.refillRate
	if r.tokens > r.maxTokens {
		r.tokens = r.maxTokens
	}
	r.lastRefill = now

	// Wait if no tokens available
	if r.tokens < 1 {
		waitTime := time.Duration((1-r.tokens)/r.refillRate*1000) * time.Millisecond
		r.mu.Unlock()
		time.Sleep(waitTime)
		r.mu.Lock()
		r.tokens = 0
	} else {
		r.tokens--
	}

	r.requestCount++

	// Apply random delay if configured
	if r.maxDelayMs > 0 {
		delayRange := r.maxDelayMs - r.minDelayMs
		if delayRange <= 0 {
			delayRange = 1
		}
		delay := r.minDelayMs + rand.Intn(delayRange)
		r.mu.Unlock()
		time.Sleep(time.Duration(delay) * time.Millisecond)
		r.mu.Lock()
	}
}

// RequestCount returns the total number of requests made
func (r *RateLimiter) RequestCount() int64 {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.requestCount
}

// SetRPS dynamically adjusts the rate limit
func (r *RateLimiter) SetRPS(rps float64) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.maxTokens = rps
	r.refillRate = rps
}