memory.go RawHistoryFolder

1// Package ratelimit provides rate limiter implementations.
2package ratelimit
3
4import (
5	"context"
6	"sync"
7	"time"
8
9	"github.com/bnema/zerowrap"
10	"golang.org/x/time/rate"
11
12	"github.com/bnema/gordon/internal/boundaries/out"
13)
14
15// Ensure MemoryStore implements out.RateLimiter.
16var _ out.RateLimiter = (*MemoryStore)(nil)
17
18// MemoryStore is an in-memory rate limiter implementation using golang.org/x/time/rate.
19// Each unique key gets its own independent rate limiter.
20type MemoryStore struct {
21	limiters map[string]*rate.Limiter
22	mu       sync.RWMutex
23	rps      float64
24	burst    int
25	log      zerowrap.Logger
26}
27
28// NewMemoryStore creates a new in-memory rate limiter store.
29func NewMemoryStore(rps float64, burst int, log zerowrap.Logger) *MemoryStore {
30	return &MemoryStore{
31		limiters: make(map[string]*rate.Limiter),
32		rps:      rps,
33		burst:    burst,
34		log:      log,
35	}
36}
37
38// Allow checks if a request identified by key is allowed.
39// Returns true if allowed, false if rate limited.
40func (s *MemoryStore) Allow(_ context.Context, key string) bool {
41	return s.getLimiter(key).Allow()
42}
43
44// AllowN checks if n requests identified by key are allowed.
45func (s *MemoryStore) AllowN(_ context.Context, key string, n int) bool {
46	return s.getLimiter(key).AllowN(time.Now(), n)
47}
48
49// getLimiter returns the rate limiter for the given key, creating one if it doesn't exist.
50func (s *MemoryStore) getLimiter(key string) *rate.Limiter {
51	s.mu.RLock()
52	limiter, exists := s.limiters[key]
53	s.mu.RUnlock()
54
55	if exists {
56		return limiter
57	}
58
59	s.mu.Lock()
60	defer s.mu.Unlock()
61
62	// Double-check after acquiring write lock
63	if limiter, exists = s.limiters[key]; exists {
64		return limiter
65	}
66
67	limiter = rate.NewLimiter(rate.Limit(s.rps), s.burst)
68	s.limiters[key] = limiter
69	return limiter
70}