Download this file
| 1 | package calls |
| 2 | |
| 3 | import ( |
| 4 | "sync" |
| 5 | "time" |
| 6 | ) |
| 7 | |
| 8 | // budget is a sliding-window counter: at most limit events per key per |
| 9 | // window. It counts attempts, refused or not, because the point is to |
| 10 | // keep outsider traffic away from rastrillo's shared limiter (spec |
| 11 | // §3.2), and an attempt that reaches Begin is what consumes that. The |
| 12 | // budgets are process-local and reset on every deploy or restart. |
| 13 | type budget struct { |
| 14 | mu sync.Mutex |
| 15 | limit int |
| 16 | window time.Duration |
| 17 | hits map[string][]time.Time |
| 18 | } |
| 19 | |
| 20 | func newBudget(limit int, window time.Duration) *budget { |
| 21 | return &budget{limit: limit, window: window, hits: map[string][]time.Time{}} |
| 22 | } |
| 23 | |
| 24 | // allow records an attempt and reports whether it was within budget. |
| 25 | func (b *budget) allow(key string, now time.Time) bool { |
| 26 | b.mu.Lock() |
| 27 | defer b.mu.Unlock() |
| 28 | cut := now.Add(-b.window) |
| 29 | kept := b.hits[key][:0] |
| 30 | for _, t := range b.hits[key] { |
| 31 | if t.After(cut) { |
| 32 | kept = append(kept, t) |
| 33 | } |
| 34 | } |
| 35 | if len(kept) == 0 { |
| 36 | delete(b.hits, key) |
| 37 | } else { |
| 38 | b.hits[key] = kept |
| 39 | } |
| 40 | if len(kept) >= b.limit { |
| 41 | return false |
| 42 | } |
| 43 | b.hits[key] = append(kept, now) |
| 44 | return true |
| 45 | } |
| 46 | |
amadan