Refactor RateLimiter: improve request handling and cleanup logic

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