1package ratelimit
  2
  3import (
  4	"container/list"
  5	"errors"
  6	"net/netip"
  7	"sync"
  8	"time"
  9)
 10
 11type Rate struct {
 12	Burst int
 13	Every time.Duration
 14}
 15
 16type entry[K comparable] struct {
 17	key    K
 18	tokens float64
 19	at     time.Time
 20}
 21
 22var errNoEntries = errors.New("ratelimit: a table needs at least one entry")
 23
 24type Table[K comparable] struct {
 25	mu      sync.Mutex
 26	now     func() time.Time
 27	max     int
 28	entries map[K]*list.Element
 29	order   *list.List
 30	evicted uint64
 31}
 32
 33func New[K comparable](max int, now func() time.Time) *Table[K] {
 34	if max < 1 {
 35		panic(errNoEntries)
 36	}
 37	return &Table[K]{now: now, max: max, entries: map[K]*list.Element{}, order: list.New()}
 38}
 39
 40func (t *Table[K]) bucket(k K, rate Rate, now time.Time) *entry[K] {
 41	el, ok := t.entries[k]
 42	if !ok {
 43		el = t.order.PushFront(&entry[K]{key: k, tokens: float64(rate.Burst), at: now})
 44		t.entries[k] = el
 45		for t.order.Len() > t.max {
 46			oldest := t.order.Back()
 47			t.order.Remove(oldest)
 48			delete(t.entries, oldest.Value.(*entry[K]).key)
 49			t.evicted++
 50		}
 51	}
 52	t.order.MoveToFront(el)
 53	e := el.Value.(*entry[K])
 54	if elapsed := now.Sub(e.at); elapsed > 0 {
 55		e.tokens = min(float64(rate.Burst), e.tokens+float64(elapsed)/float64(rate.Every))
 56	}
 57	e.at = now
 58	return e
 59}
 60
 61func wait(tokens float64, rate Rate) time.Duration {
 62	return time.Duration((1 - tokens) * float64(rate.Every))
 63}
 64
 65func (t *Table[K]) Take(k K, rate Rate) (time.Duration, bool) {
 66	if rate.Burst <= 0 || rate.Every <= 0 {
 67		return 0, true
 68	}
 69	t.mu.Lock()
 70	defer t.mu.Unlock()
 71	e := t.bucket(k, rate, t.now())
 72	if e.tokens >= 1 {
 73		e.tokens--
 74		return 0, true
 75	}
 76	return wait(e.tokens, rate), false
 77}
 78
 79func (t *Table[K]) Allowed(k K, rate Rate) (time.Duration, bool) {
 80	if rate.Burst <= 0 || rate.Every <= 0 {
 81		return 0, true
 82	}
 83	t.mu.Lock()
 84	defer t.mu.Unlock()
 85	if _, ok := t.entries[k]; !ok {
 86		return 0, true
 87	}
 88	e := t.bucket(k, rate, t.now())
 89	if e.tokens >= 1 {
 90		return 0, true
 91	}
 92	return wait(e.tokens, rate), false
 93}
 94
 95func (t *Table[K]) Len() int {
 96	t.mu.Lock()
 97	defer t.mu.Unlock()
 98	return t.order.Len()
 99}
100
101func Prefix(a netip.Addr) netip.Prefix {
102	if !a.IsValid() {
103		return netip.Prefix{}
104	}
105	a = a.WithZone("").Unmap()
106	bits := 32
107	if a.Is6() {
108		bits = 64
109	}
110	p, _ := a.Prefix(bits)
111	return p
112}