Implement sliding window rate limiter

pkg/rate/limiter.go

@@ -0,0 +1,108 @@
+package rate
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+)
+
+type Limiter struct {
+	limit  int
+	period time.Duration
+
+	mu               sync.Mutex
+	start, end, size int
+	requests         []time.Time
+}
+
+func NewLimiter(limit int, period time.Duration) *Limiter {
+	return &Limiter{
+		limit:    limit,
+		period:   period,
+		start:    0,
+		end:      0,
+		requests: make([]time.Time, limit),
+	}
+}
+
+func (l *Limiter) Wait(ctx context.Context) error {
+	return l.WaitN(ctx, 1)
+}
+
+func (l *Limiter) WaitN(ctx context.Context, n int) error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	if n <= 0 {
+		return nil
+	}
+
+	if n > l.limit {
+		return fmt.Errorf("rate: Wait (n=%d) exceed limiter %d", n, l.limit)
+	}
+
+	// Get the oldest request in queue for waiting.
+	var (
+		shouldWait bool
+		oldest     time.Time
+	)
+	if l.requestsSize()+n > l.limit {
+		shouldWait = true
+		oldest, _ = l.requestAt(l.requestsSize() + n - l.limit - 1)
+	}
+
+	// Wait if rate limit is reached.
+	if shouldWait {
+		waitDuration := l.period - time.Since(oldest)
+		if waitDuration > 0 {
+			timer := time.NewTimer(waitDuration)
+			defer timer.Stop()
+
+			select {
+			case <-timer.C:
+				// We can proceed.
+			case <-ctx.Done():
+				// Context was canceled before we could proceed.
+				return ctx.Err()
+			}
+		}
+	}
+
+	// Add new requests to queue.
+	for range n {
+		l.addRequest(time.Now())
+	}
+
+	return nil
+}
+
+func (l *Limiter) requestsSize() int {
+	return l.size
+}
+
+func (l *Limiter) requestAt(i int) (time.Time, bool) {
+	if i >= l.size {
+		return time.Now(), false
+	}
+
+	return l.requests[(l.start+i)%l.limit], true
+}
+
+func (l *Limiter) addRequest(t time.Time) {
+	if l.size == l.limit {
+		l.start++
+		if l.start >= l.limit {
+			l.start = 0
+		}
+	} else {
+		l.size++
+	}
+
+	l.requests[l.end] = t
+
+	l.end++
+	if l.end >= l.limit {
+		l.end = 0
+	}
+}