81 lines
1.5 KiB
Go
81 lines
1.5 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Memory is a fixed-window limiter held in process. Useful for tests and for a
|
|
// single-instance self-hoster who would rather not run Redis.
|
|
//
|
|
// Not suitable for more than one backend process: each would keep its own
|
|
// counters, so N processes allow N times the limit.
|
|
type Memory struct {
|
|
mu sync.Mutex
|
|
windows map[string]*window
|
|
Now func() time.Time // injectable so tests need not sleep
|
|
}
|
|
|
|
type window struct {
|
|
count int
|
|
expires time.Time
|
|
}
|
|
|
|
func NewMemory() *Memory {
|
|
return &Memory{
|
|
windows: make(map[string]*window),
|
|
Now: time.Now,
|
|
}
|
|
}
|
|
|
|
func (m *Memory) Allow(_ context.Context, key string, limit int, dur time.Duration) (Result, error) {
|
|
if limit <= 0 {
|
|
return Result{Allowed: true, Limit: 0, Remaining: 0}, nil
|
|
}
|
|
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
now := m.Now()
|
|
|
|
w, ok := m.windows[key]
|
|
if !ok || now.After(w.expires) {
|
|
w = &window{expires: now.Add(dur)}
|
|
m.windows[key] = w
|
|
}
|
|
|
|
w.count++
|
|
|
|
remaining := limit - w.count
|
|
if remaining < 0 {
|
|
remaining = 0
|
|
}
|
|
|
|
return Result{
|
|
Allowed: w.count <= limit,
|
|
Limit: limit,
|
|
Remaining: remaining,
|
|
RetryAfter: w.expires.Sub(now),
|
|
}, nil
|
|
}
|
|
|
|
// Sweep drops expired windows. Without it the map grows for every key ever
|
|
// seen, which for per-user keys is unbounded. Call it periodically.
|
|
func (m *Memory) Sweep() int {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
|
|
now := m.Now()
|
|
removed := 0
|
|
|
|
for k, w := range m.windows {
|
|
if now.After(w.expires) {
|
|
delete(m.windows, k)
|
|
removed++
|
|
}
|
|
}
|
|
|
|
return removed
|
|
}
|