amadan

carlos / platform Public

relay: admission, rolling budgets, token buckets (spec §3 limits)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Paul Campbell pushed by paul@keymail.dev 0853fd9dbd1b0c6fc0e9124589498f61b26787b6 parent f1593c0
2 files changed, +310 −0
  • internal/relay/limits.go +216 −0
  • internal/relay/limits_test.go +94 −0
diff --git a/internal/relay/limits.go b/internal/relay/limits.go
new file mode 100644
index 000000000..e0e208132
--- /dev/null
+++ b/internal/relay/limits.go
@@ -0,0 +1,216 @@
+package relay
+
+import (
+ "strconv"
+ "sync"
+ "time"
+)
+
+type Limits struct {
+ AppLive, AppPerMinute int
+ AppBytesPerHour int64
+ AccountLive, AccountPerMinute int
+ AccountBytesPerHour int64
+ BoxLive, BoxPerMinute int
+ AllocBitsPerSec, BoxBitsPerSec int64
+}
+
+func DefaultLimits() Limits {
+ return Limits{
+ AppLive: 64, AppPerMinute: 30, AppBytesPerHour: 20 << 30,
+ AccountLive: 256, AccountPerMinute: 60, AccountBytesPerHour: 60 << 30,
+ BoxLive: 2048, BoxPerMinute: 500,
+ AllocBitsPerSec: 8_000_000, BoxBitsPerSec: 200_000_000,
+ }
+}
+
+func LimitsFromEnv(getenv func(string) string) Limits {
+ l := DefaultLimits()
+ i := func(name string, dst *int) {
+ if v, err := strconv.Atoi(getenv("CARLOS_RELAY_LIMIT_" + name)); err == nil && v > 0 {
+ *dst = v
+ }
+ }
+ i64 := func(name string, dst *int64) {
+ if v, err := strconv.ParseInt(getenv("CARLOS_RELAY_LIMIT_"+name), 10, 64); err == nil && v > 0 {
+ *dst = v
+ }
+ }
+ i("APP_LIVE", &l.AppLive)
+ i("APP_PER_MINUTE", &l.AppPerMinute)
+ i64("APP_BYTES_PER_HOUR", &l.AppBytesPerHour)
+ i("ACCOUNT_LIVE", &l.AccountLive)
+ i("ACCOUNT_PER_MINUTE", &l.AccountPerMinute)
+ i64("ACCOUNT_BYTES_PER_HOUR", &l.AccountBytesPerHour)
+ i("BOX_LIVE", &l.BoxLive)
+ i("BOX_PER_MINUTE", &l.BoxPerMinute)
+ i64("ALLOC_BPS", &l.AllocBitsPerSec)
+ i64("BOX_BPS", &l.BoxBitsPerSec)
+ return l
+}
+
+type Identity struct{ Account, App string }
+
+type window struct { // sliding counters bucketed per minute, 60 buckets for an hour
+ at [60]time.Time
+ n [60]int64
+}
+
+func (w *window) add(now time.Time, n int64) {
+ m := now.Truncate(time.Minute)
+ i := int(m.Unix()/60) % 60
+ if !w.at[i].Equal(m) {
+ w.at[i], w.n[i] = m, 0
+ }
+ w.n[i] += n
+}
+
+func (w *window) sum(now time.Time, span time.Duration) int64 {
+ var s int64
+ cut := now.Add(-span)
+ for i := range w.at {
+ if !w.at[i].IsZero() && w.at[i].After(cut) {
+ s += w.n[i]
+ }
+ }
+ return s
+}
+
+type counters struct {
+ live int
+ starts window // attempts per minute
+ bytes window // bytes per rolling hour
+}
+
+type Admission struct {
+ mu sync.Mutex
+ l Limits
+ now func() time.Time
+ app map[Identity]*counters
+ acct map[string]*counters
+ box counters
+}
+
+func NewAdmission(l Limits, now func() time.Time) *Admission {
+ return &Admission{l: l, now: now, app: map[Identity]*counters{}, acct: map[string]*counters{}}
+}
+
+func (a *Admission) get(id Identity) (*counters, *counters) {
+ c, ok := a.app[id]
+ if !ok {
+ c = &counters{}
+ a.app[id] = c
+ }
+ ac, ok := a.acct[id.Account]
+ if !ok {
+ ac = &counters{}
+ a.acct[id.Account] = ac
+ }
+ return c, ac
+}
+
+// Admit counts an attempt for the app, the account and the box whether or
+// not it is admitted, then checks live caps, rates and budgets in that order.
+func (a *Admission) Admit(id Identity) (bool, string) {
+ a.mu.Lock()
+ defer a.mu.Unlock()
+ now := a.now()
+ c, ac := a.get(id)
+ c.starts.add(now, 1)
+ ac.starts.add(now, 1)
+ a.box.starts.add(now, 1)
+ switch {
+ case c.live >= a.l.AppLive:
+ return false, "app live"
+ case ac.live >= a.l.AccountLive:
+ return false, "account live"
+ case a.box.live >= a.l.BoxLive:
+ return false, "box live"
+ case c.starts.sum(now, time.Minute) > int64(a.l.AppPerMinute):
+ return false, "app rate"
+ case ac.starts.sum(now, time.Minute) > int64(a.l.AccountPerMinute):
+ return false, "account rate"
+ case a.box.starts.sum(now, time.Minute) > int64(a.l.BoxPerMinute):
+ return false, "box rate"
+ case c.bytes.sum(now, time.Hour) > a.l.AppBytesPerHour:
+ return false, "app budget"
+ case ac.bytes.sum(now, time.Hour) > a.l.AccountBytesPerHour:
+ return false, "account budget"
+ }
+ c.live++
+ ac.live++
+ a.box.live++
+ return true, ""
+}
+
+func (a *Admission) Release(id Identity) {
+ a.mu.Lock()
+ defer a.mu.Unlock()
+ c, ac := a.get(id)
+ if c.live > 0 {
+ c.live--
+ }
+ if ac.live > 0 {
+ ac.live--
+ }
+ if a.box.live > 0 {
+ a.box.live--
+ }
+}
+
+// AddBytes records forwarded bytes and reports whether the app or account
+// rolling-hour budget is now exceeded.
+func (a *Admission) AddBytes(id Identity, n int64) bool {
+ a.mu.Lock()
+ defer a.mu.Unlock()
+ now := a.now()
+ c, ac := a.get(id)
+ c.bytes.add(now, n)
+ ac.bytes.add(now, n)
+ return c.bytes.sum(now, time.Hour) > a.l.AppBytesPerHour || ac.bytes.sum(now, time.Hour) > a.l.AccountBytesPerHour
+}
+
+func (a *Admission) Seed(id Identity, at time.Time, n int64) {
+ a.mu.Lock()
+ defer a.mu.Unlock()
+ c, ac := a.get(id)
+ c.bytes.add(at, n)
+ ac.bytes.add(at, n)
+}
+
+func (a *Admission) OverBudget(id Identity, now time.Time) bool {
+ a.mu.Lock()
+ defer a.mu.Unlock()
+ c, ac := a.get(id)
+ return c.bytes.sum(now, time.Hour) > a.l.AppBytesPerHour || ac.bytes.sum(now, time.Hour) > a.l.AccountBytesPerHour
+}
+
+// Bucket is a token bucket in bytes with a one-second burst.
+type Bucket struct {
+ mu sync.Mutex
+ rate float64 // bytes per second
+ tokens float64
+ last time.Time
+ now func() time.Time
+}
+
+func NewBucket(bitsPerSec int64, now func() time.Time) *Bucket {
+ r := float64(bitsPerSec) / 8
+ return &Bucket{rate: r, tokens: r, last: now(), now: now}
+}
+
+func (b *Bucket) Allow(nbytes int) bool {
+ b.mu.Lock()
+ defer b.mu.Unlock()
+ t := b.now()
+ b.tokens += t.Sub(b.last).Seconds() * b.rate
+ if b.tokens > b.rate {
+ b.tokens = b.rate
+ }
+ b.last = t
+ if float64(nbytes) > b.tokens {
+ return false
+ }
+ b.tokens -= float64(nbytes)
+ return true
+}
diff --git a/internal/relay/limits_test.go b/internal/relay/limits_test.go
new file mode 100644
index 000000000..1c01f78f2
--- /dev/null
+++ b/internal/relay/limits_test.go
@@ -0,0 +1,94 @@
+package relay
+
+import (
+ "testing"
+ "time"
+)
+
+func TestAdmissionCaps(t *testing.T) {
+ now := time.Unix(1_800_000_000, 0)
+ clock := func() time.Time { return now }
+ l := DefaultLimits()
+ l.AppLive, l.AppPerMinute, l.AccountLive, l.BoxLive = 2, 100, 3, 100
+ a := NewAdmission(l, clock)
+ meet := Identity{"abc", "meet"}
+ docs := Identity{"abc", "docs"}
+ for i := 0; i < 2; i++ {
+ if ok, r := a.Admit(meet); !ok {
+ t.Fatalf("admit %d refused: %s", i, r)
+ }
+ }
+ if ok, r := a.Admit(meet); ok || r != "app live" {
+ t.Fatalf("third meet = %v %q", ok, r)
+ }
+ if ok, _ := a.Admit(docs); !ok {
+ t.Fatal("docs first should pass")
+ }
+ if ok, r := a.Admit(docs); ok || r != "account live" {
+ t.Fatalf("account cap = %v %q", ok, r)
+ }
+ a.Release(meet)
+ if ok, _ := a.Admit(docs); !ok {
+ t.Fatal("after release, account has room")
+ }
+}
+
+func TestAdmissionRate(t *testing.T) {
+ now := time.Unix(1_800_000_000, 0)
+ clock := func() time.Time { return now }
+ l := DefaultLimits()
+ l.AppPerMinute = 2
+ a := NewAdmission(l, clock)
+ id := Identity{"abc", "meet"}
+ a.Admit(id)
+ a.Release(id)
+ a.Admit(id)
+ a.Release(id)
+ if ok, r := a.Admit(id); ok || r != "app rate" {
+ t.Fatalf("rate = %v %q", ok, r)
+ }
+ now = now.Add(61 * time.Second)
+ if ok, _ := a.Admit(id); !ok {
+ t.Fatal("rate window should have rolled")
+ }
+}
+
+func TestAdmissionBudgetRollsAndSeeds(t *testing.T) {
+ now := time.Unix(1_800_000_000, 0)
+ clock := func() time.Time { return now }
+ l := DefaultLimits()
+ l.AppBytesPerHour = 100
+ a := NewAdmission(l, clock)
+ id := Identity{"abc", "meet"}
+ if a.AddBytes(id, 60) {
+ t.Fatal("60 under 100")
+ }
+ if !a.AddBytes(id, 50) {
+ t.Fatal("110 over 100")
+ }
+ now = now.Add(61 * time.Minute)
+ if a.OverBudget(id, now) {
+ t.Fatal("window rolled")
+ }
+ b := NewAdmission(l, clock)
+ b.Seed(id, now.Add(-30*time.Minute), 90)
+ if !b.AddBytes(id, 20) {
+ t.Fatal("seeded 90 + 20 over 100")
+ }
+}
+
+func TestBucket(t *testing.T) {
+ now := time.Unix(1_800_000_000, 0)
+ clock := func() time.Time { return now }
+ b := NewBucket(8000, clock) // 1000 bytes/s
+ if !b.Allow(1000) {
+ t.Fatal("first second's worth")
+ }
+ if b.Allow(1) {
+ t.Fatal("empty")
+ }
+ now = now.Add(500 * time.Millisecond)
+ if !b.Allow(500) || b.Allow(1) {
+ t.Fatal("refilled 500")
+ }
+}