SpotiFLAC-Mobile/go_backend/ratelimit.go

package gobackend

import (
	"context"
	"sync"
	"time"
)

type RateLimiter struct {
	mu          sync.Mutex
	maxRequests int
	window      time.Duration
	timestamps  []time.Time
}

func NewRateLimiter(maxRequests int, window time.Duration) *RateLimiter {
	return &RateLimiter{
		maxRequests: maxRequests,
		window:      window,
		timestamps:  make([]time.Time, 0, maxRequests),
	}
}

func (r *RateLimiter) WaitForSlot() {
	_ = r.WaitForSlotContext(context.Background())

}

// WaitForSlotContext reserves exactly one slot, rechecking the window after
// every wake-up. Multiple waiters may wake together, but only the first one
// that reacquires the mutex can consume the newly available slot.
func (r *RateLimiter) WaitForSlotContext(ctx context.Context) error {
	if ctx == nil {
		ctx = context.Background()
	}
	for {
		r.mu.Lock()
		now := time.Now()
		r.cleanOldTimestamps(now)
		if len(r.timestamps) < r.maxRequests {
			r.timestamps = append(r.timestamps, now)
			r.mu.Unlock()
			return nil
		}

		waitDuration := r.timestamps[0].Add(r.window).Sub(now)
		r.mu.Unlock()
		if waitDuration <= 0 {
			continue
		}

		timer := time.NewTimer(waitDuration)
		select {
		case <-ctx.Done():
			if !timer.Stop() {
				<-timer.C
			}
			return ctx.Err()
		case <-timer.C:
		}
	}
}

func (r *RateLimiter) cleanOldTimestamps(now time.Time) {
	cutoff := now.Add(-r.window)
	validStart := 0

	for i, ts := range r.timestamps {
		if ts.After(cutoff) {
			validStart = i
			break
		}
		validStart = i + 1
	}

	if validStart > 0 {
		r.timestamps = r.timestamps[validStart:]
	}
}

func (r *RateLimiter) TryAcquire() bool {
	r.mu.Lock()
	defer r.mu.Unlock()

	now := time.Now()
	r.cleanOldTimestamps(now)

	if len(r.timestamps) < r.maxRequests {
		r.timestamps = append(r.timestamps, now)
		return true
	}

	return false
}

func (r *RateLimiter) Available() int {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.cleanOldTimestamps(time.Now())
	return r.maxRequests - len(r.timestamps)
}