memory.go RawHistoryFolder

1// Package ratelimit provides rate limiter implementations.
2package ratelimit
3
4import (
5	"container/list"
6	"context"
7	"sync"
8	"time"
9
10	"github.com/bnema/zerowrap"
11	"golang.org/x/time/rate"
12
13	"github.com/bnema/gordon/internal/boundaries/out"
14)
15
16// Ensure MemoryStore implements out.RateLimiter.
17var _ out.RateLimiter = (*MemoryStore)(nil)
18
19// MemoryStore is an in-memory rate limiter implementation using golang.org/x/time/rate.
20// Each unique key gets its own independent rate limiter.
21type limiterEntry struct {
22	key      string
23	limiter  *rate.Limiter
24	lastSeen time.Time
25}
26
27const (
28	maxLimiterEntries = 10000
29	limiterIdleTTL    = 10 * time.Minute
30)
31
32type MemoryStore struct {
33	limiters map[string]*list.Element
34	lru      *list.List
35	mu       sync.Mutex
36	rps      float64
37	burst    int
38	log      zerowrap.Logger
39}
40
41// NewMemoryStore creates a new in-memory rate limiter store.
42func NewMemoryStore(rps float64, burst int, log zerowrap.Logger) *MemoryStore {
43	return &MemoryStore{
44		limiters: make(map[string]*list.Element),
45		lru:      list.New(),
46		rps:      rps,
47		burst:    burst,
48		log:      log,
49	}
50}
51
52// Allow checks if a request identified by key is allowed.
53// Returns true if allowed, false if rate limited.
54func (s *MemoryStore) Allow(_ context.Context, key string) bool {
55	return s.getLimiter(key).Allow()
56}
57
58// AllowN checks if n requests identified by key are allowed.
59func (s *MemoryStore) AllowN(_ context.Context, key string, n int) bool {
60	return s.getLimiter(key).AllowN(time.Now(), n)
61}
62
63// getLimiter returns the rate limiter for the given key, creating one if it doesn't exist.
64func (s *MemoryStore) getLimiter(key string) *rate.Limiter {
65	now := time.Now()
66	s.mu.Lock()
67	defer s.mu.Unlock()
68
69	if element, exists := s.limiters[key]; exists {
70		entry := element.Value.(*limiterEntry)
71		entry.lastSeen = now
72		s.lru.MoveToFront(element)
73		return entry.limiter
74	}
75
76	s.evictExpiredOrOldest(now)
77	limiter := rate.NewLimiter(rate.Limit(s.rps), s.burst)
78	entry := &limiterEntry{key: key, limiter: limiter, lastSeen: now}
79	s.limiters[key] = s.lru.PushFront(entry)
80	return limiter
81}
82
83func (s *MemoryStore) evictExpiredOrOldest(now time.Time) {
84	for element := s.lru.Back(); element != nil; {
85		entry := element.Value.(*limiterEntry)
86		if now.Sub(entry.lastSeen) <= limiterIdleTTL {
87			break
88		}
89		previous := element.Prev()
90		s.removeElement(element)
91		element = previous
92	}
93	if len(s.limiters) >= maxLimiterEntries {
94		s.removeElement(s.lru.Back())
95	}
96}
97
98func (s *MemoryStore) removeElement(element *list.Element) {
99	if element == nil {
100		return
101	}
102	entry := element.Value.(*limiterEntry)
103	delete(s.limiters, entry.key)
104	s.lru.Remove(element)
105}