136 lines
3.0 KiB
Go
136 lines
3.0 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type Policy struct {
|
|
Capacity uint64
|
|
Window time.Duration
|
|
}
|
|
|
|
func (p Policy) Valid() bool { return p.Capacity > 0 && p.Window > 0 }
|
|
|
|
type Bucket struct {
|
|
Key string
|
|
Policy Policy
|
|
}
|
|
|
|
type Decision struct {
|
|
Allowed bool
|
|
RetryAt time.Time
|
|
reservation *Reservation
|
|
}
|
|
|
|
func (d Decision) Reservation() *Reservation { return d.reservation }
|
|
|
|
type windowState struct {
|
|
startedAt time.Time
|
|
expiresAt time.Time
|
|
count uint64
|
|
}
|
|
|
|
type Limiter struct {
|
|
mutex sync.Mutex
|
|
windows map[string]windowState
|
|
takes uint64
|
|
}
|
|
|
|
func New() *Limiter { return &Limiter{windows: make(map[string]windowState)} }
|
|
|
|
// Take checks all buckets atomically. Rejected checks do not consume any
|
|
// bucket, which keeps combined user/API-key limits deterministic.
|
|
func (l *Limiter) Take(now time.Time, buckets ...Bucket) Decision {
|
|
if l == nil || now.IsZero() || len(buckets) == 0 {
|
|
return Decision{}
|
|
}
|
|
seen := make(map[string]struct{}, len(buckets))
|
|
for _, bucket := range buckets {
|
|
if bucket.Key == "" || !bucket.Policy.Valid() {
|
|
return Decision{}
|
|
}
|
|
if _, exists := seen[bucket.Key]; exists {
|
|
return Decision{}
|
|
}
|
|
seen[bucket.Key] = struct{}{}
|
|
}
|
|
|
|
l.mutex.Lock()
|
|
defer l.mutex.Unlock()
|
|
l.takes++
|
|
if l.takes%256 == 0 {
|
|
for key, state := range l.windows {
|
|
if !now.Before(state.expiresAt) {
|
|
delete(l.windows, key)
|
|
}
|
|
}
|
|
}
|
|
|
|
states := make([]windowState, len(buckets))
|
|
var retryAt time.Time
|
|
for index, bucket := range buckets {
|
|
state := l.windows[bucket.Key]
|
|
if state.startedAt.IsZero() || !now.Before(state.expiresAt) {
|
|
state = windowState{startedAt: now, expiresAt: now.Add(bucket.Policy.Window)}
|
|
}
|
|
states[index] = state
|
|
if state.count >= bucket.Policy.Capacity {
|
|
candidate := state.expiresAt
|
|
if candidate.After(retryAt) {
|
|
retryAt = candidate
|
|
}
|
|
}
|
|
}
|
|
if !retryAt.IsZero() {
|
|
return Decision{RetryAt: retryAt}
|
|
}
|
|
|
|
items := make([]reservationItem, 0, len(buckets))
|
|
for index, bucket := range buckets {
|
|
state := states[index]
|
|
state.count++
|
|
l.windows[bucket.Key] = state
|
|
items = append(items, reservationItem{key: bucket.Key, startedAt: state.startedAt})
|
|
}
|
|
reservation := &Reservation{limiter: l, items: items}
|
|
return Decision{Allowed: true, reservation: reservation}
|
|
}
|
|
|
|
type reservationItem struct {
|
|
key string
|
|
startedAt time.Time
|
|
}
|
|
|
|
// Reservation can refund a pre-call Provider token when no upstream attempt
|
|
// starts. Commit is implicit; callers only cancel abandoned reservations.
|
|
type Reservation struct {
|
|
mutex sync.Mutex
|
|
limiter *Limiter
|
|
items []reservationItem
|
|
canceled bool
|
|
}
|
|
|
|
func (r *Reservation) Cancel() bool {
|
|
if r == nil || r.limiter == nil {
|
|
return false
|
|
}
|
|
r.mutex.Lock()
|
|
defer r.mutex.Unlock()
|
|
if r.canceled {
|
|
return false
|
|
}
|
|
r.limiter.mutex.Lock()
|
|
defer r.limiter.mutex.Unlock()
|
|
for _, item := range r.items {
|
|
state, ok := r.limiter.windows[item.key]
|
|
if !ok || !state.startedAt.Equal(item.startedAt) || state.count == 0 {
|
|
continue
|
|
}
|
|
state.count--
|
|
r.limiter.windows[item.key] = state
|
|
}
|
|
r.canceled = true
|
|
return true
|
|
}
|