diff options
Diffstat (limited to 'utils')
| -rw-r--r-- | utils/ratelimit/limiter.go | 39 |
1 files changed, 20 insertions, 19 deletions
diff --git a/utils/ratelimit/limiter.go b/utils/ratelimit/limiter.go index 5ef1792..ecbcc56 100644 --- a/utils/ratelimit/limiter.go +++ b/utils/ratelimit/limiter.go @@ -5,32 +5,42 @@ import ( "time" ) -// RateLimiter manages request rate limiting for APIs type RateLimiter struct { mu sync.Mutex + lastRequest time.Time lastRequests []time.Time - maxRequests int // Maximum requests per time window - window time.Duration // Time window duration + maxRequests int + window time.Duration + minDelay time.Duration } -// NewRateLimiter creates a new rate limiter -// maxRequests is the maximum number of requests allowed in the specified time window func NewRateLimiter(maxRequests int, window time.Duration) *RateLimiter { + minDelay := window / time.Duration(maxRequests) return &RateLimiter{ lastRequests: make([]time.Time, 0, maxRequests), maxRequests: maxRequests, window: window, + minDelay: minDelay, } } -// Wait blocks until a request can be made according to rate limiting rules func (r *RateLimiter) Wait() { r.mu.Lock() defer r.mu.Unlock() now := time.Now() - // Clean up old requests + if !r.lastRequest.IsZero() { + elapsed := now.Sub(r.lastRequest) + if elapsed < r.minDelay { + waitTime := r.minDelay - elapsed + r.mu.Unlock() + time.Sleep(waitTime) + r.mu.Lock() + now = time.Now() + } + } + cutoff := now.Add(-r.window) i := 0 for i < len(r.lastRequests) && r.lastRequests[i].Before(cutoff) { @@ -40,18 +50,14 @@ func (r *RateLimiter) Wait() { r.lastRequests = r.lastRequests[i:] } - // If we've reached max requests in the window, wait until we can make another if len(r.lastRequests) >= r.maxRequests { - // Calculate wait time based on the oldest request in the window oldestInWindow := r.lastRequests[0] waitDuration := r.window - now.Sub(oldestInWindow) - // Release lock while waiting r.mu.Unlock() - time.Sleep(waitDuration + time.Millisecond) // Add 1ms to be safe - r.mu.Lock() // Re-acquire lock + time.Sleep(waitDuration + time.Millisecond) + r.mu.Lock() - // Refresh current time and clean up again after waiting now = time.Now() cutoff = now.Add(-r.window) i = 0 @@ -63,11 +69,10 @@ func (r *RateLimiter) Wait() { } } - // Add current request timestamp + r.lastRequest = now r.lastRequests = append(r.lastRequests, now) } -// RemainingRequests returns the number of requests that can still be made in the current window func (r *RateLimiter) RemainingRequests() int { r.mu.Lock() defer r.mu.Unlock() @@ -75,7 +80,6 @@ func (r *RateLimiter) RemainingRequests() int { now := time.Now() cutoff := now.Add(-r.window) - // Clean up old requests i := 0 for i < len(r.lastRequests) && r.lastRequests[i].Before(cutoff) { i++ @@ -87,19 +91,16 @@ func (r *RateLimiter) RemainingRequests() int { return r.maxRequests - len(r.lastRequests) } -// MultiLimiter combines multiple rate limiters type MultiLimiter struct { limiters []*RateLimiter } -// NewMultiLimiter creates a new multi-limiter from the given limiters func NewMultiLimiter(limiters ...*RateLimiter) *MultiLimiter { return &MultiLimiter{ limiters: limiters, } } -// Wait waits for all underlying limiters func (m *MultiLimiter) Wait() { for _, limiter := range m.limiters { limiter.Wait() |
