Ilia Denisov
137 lines
3.1 KiB
Go
137 lines
3.1 KiB
Go
// Package ratelimit provides small process-local rate-limit primitives used by
|
|
// the gateway edge policy layers.
|
|
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
// Policy describes one token-bucket budget enforced for a concrete key.
|
|
type Policy struct {
|
|
// Requests is the number of accepted requests replenished per Window.
|
|
Requests int
|
|
|
|
// Window is the interval over which Requests are replenished.
|
|
Window time.Duration
|
|
|
|
// Burst is the maximum number of immediately available tokens.
|
|
Burst int
|
|
}
|
|
|
|
// Decision describes the result of one limiter reservation attempt.
|
|
type Decision struct {
|
|
// Allowed reports whether the request may proceed immediately.
|
|
Allowed bool
|
|
|
|
// RetryAfter is the minimum delay the caller should wait before retrying
|
|
// when Allowed is false.
|
|
RetryAfter time.Duration
|
|
}
|
|
|
|
// Limiter applies a policy to one concrete key.
|
|
type Limiter interface {
|
|
// Reserve evaluates key under policy and reports whether the request may
|
|
// proceed immediately.
|
|
Reserve(key string, policy Policy) Decision
|
|
}
|
|
|
|
// InMemory is a process-local Limiter backed by x/time/rate token buckets.
|
|
type InMemory struct {
|
|
now func() time.Time
|
|
cleanupInterval time.Duration
|
|
|
|
mu sync.Mutex
|
|
entries map[string]*entry
|
|
nextCleanup time.Time
|
|
}
|
|
|
|
type entry struct {
|
|
limiter *rate.Limiter
|
|
limit rate.Limit
|
|
burst int
|
|
expiresAt time.Time
|
|
}
|
|
|
|
// NewInMemory constructs a process-local limiter suitable for one gateway
|
|
// process instance.
|
|
func NewInMemory() *InMemory {
|
|
return &InMemory{
|
|
now: time.Now,
|
|
cleanupInterval: time.Minute,
|
|
entries: make(map[string]*entry),
|
|
}
|
|
}
|
|
|
|
// Reserve evaluates key against policy and reports whether the request may
|
|
// proceed immediately.
|
|
func (l *InMemory) Reserve(key string, policy Policy) Decision {
|
|
if policy.Requests <= 0 || policy.Window <= 0 || policy.Burst <= 0 {
|
|
return Decision{}
|
|
}
|
|
|
|
now := l.now()
|
|
limit := rate.Limit(float64(policy.Requests) / policy.Window.Seconds())
|
|
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
l.cleanupExpiredBucketsLocked(now)
|
|
|
|
current, ok := l.entries[key]
|
|
if !ok || current.limit != limit || current.burst != policy.Burst {
|
|
current = &entry{
|
|
limiter: rate.NewLimiter(limit, policy.Burst),
|
|
limit: limit,
|
|
burst: policy.Burst,
|
|
}
|
|
l.entries[key] = current
|
|
}
|
|
|
|
current.expiresAt = now.Add(entryTTL(policy.Window))
|
|
|
|
reservation := current.limiter.ReserveN(now, 1)
|
|
if !reservation.OK() {
|
|
return Decision{
|
|
Allowed: false,
|
|
RetryAfter: policy.Window,
|
|
}
|
|
}
|
|
|
|
retryAfter := reservation.DelayFrom(now)
|
|
if retryAfter > 0 {
|
|
return Decision{
|
|
Allowed: false,
|
|
RetryAfter: retryAfter,
|
|
}
|
|
}
|
|
|
|
return Decision{Allowed: true}
|
|
}
|
|
|
|
func (l *InMemory) cleanupExpiredBucketsLocked(now time.Time) {
|
|
if !l.nextCleanup.IsZero() && now.Before(l.nextCleanup) {
|
|
return
|
|
}
|
|
|
|
for key, current := range l.entries {
|
|
if !current.expiresAt.After(now) {
|
|
delete(l.entries, key)
|
|
}
|
|
}
|
|
|
|
l.nextCleanup = now.Add(l.cleanupInterval)
|
|
}
|
|
|
|
func entryTTL(window time.Duration) time.Duration {
|
|
if window < time.Minute {
|
|
return time.Minute
|
|
}
|
|
|
|
return 2 * window
|
|
}
|
|
|
|
var _ Limiter = (*InMemory)(nil)
|