internal/control/ratelimit.go

 1package control
 2
 3import (
 4	"sync"
 5	"time"
 6)
 7
 8// writeLimiter bounds mutating commands per account. It lives in the
 9// dispatcher rather than in a surface because every surface reaches the
10// same registry: SSH, the JSON API and the web draw on one budget, so a
11// caller cannot get a fresh allowance by changing how it connects (#148).
12//
13// The API's own limiter stays in front of it, keyed by token or IP: that
14// one bounds a network source, this one bounds an account.
15//
16// A command is one token whatever it writes. `account import-bundle`
17// replays a whole bundle in a single dispatch, so bulk work costs one
18// write and only a loop of separate commands spends the budget.
19type writeLimiter struct {
20	mu      sync.Mutex
21	buckets map[int64]*writeBucket
22}
23
24type writeBucket struct {
25	tokens float64
26	last   time.Time
27}
28
29var writes = &writeLimiter{buckets: map[int64]*writeBucket{}}
30
31// allow reports whether this account may write now, and how long until it
32// can if not. perMinute is both the sustained rate and the burst, so an
33// idle account gets a full minute's worth at once.
34func (l *writeLimiter) allow(user int64, perMinute int) (bool, time.Duration) {
35	l.mu.Lock()
36	defer l.mu.Unlock()
37	now := time.Now()
38
39	// Bounded by the number of accounts that wrote in the last ten
40	// minutes; a busy instance never grows this without also using it.
41	if len(l.buckets) > 4096 {
42		for k, b := range l.buckets {
43			if now.Sub(b.last) > 10*time.Minute {
44				delete(l.buckets, k)
45			}
46		}
47	}
48
49	burst := float64(perMinute)
50	rate := burst / 60
51	b := l.buckets[user]
52	if b == nil {
53		b = &writeBucket{tokens: burst, last: now}
54		l.buckets[user] = b
55	}
56	elapsed := now.Sub(b.last).Seconds()
57	b.last = now
58	if b.tokens += elapsed * rate; b.tokens > burst {
59		b.tokens = burst
60	}
61	if b.tokens < 1 {
62		return false, time.Duration((1-b.tokens)/rate*float64(time.Second)) + time.Second
63	}
64	b.tokens--
65	return true, 0
66}