rate-limit.ts · Source

rate-limit.tsapps/realtime/src/rate-limit.tsraw ← src4,254 B← Back to the app
1/**
2 * Rate limiters (TASK-0091 R4/R6).
3 *
4 * - {@link TokenBucket} — per-socket SUSTAINED message rate (100 msg/s;
5 *   overflow → close 4008 + audit). Continuous refill: capacity tokens,
6 *   `rate` tokens/second. Deterministic via an injectable clock.
7 * - {@link ConnectLimiter} — connect-time admission (30 conns/min/IP):
8 *   a sliding window of connect timestamps per key.
9 *
10 * Both are pure (no timers) so the heartbeat/limit integration tests can
11 * drive them with a fake clock.
12 */
13 
14export interface TokenBucketOpts {
15  /** Bucket capacity (= burst allowance). */
16  readonly capacity: number;
17  /** Sustained tokens per second. */
18  readonly rate: number;
19  /** Injectable clock (ms) — defaults to Date.now. */
20  readonly now?: () => number;
21}
22 
23export class TokenBucket {
24  private readonly capacity: number;
25  private readonly rate: number;
26  private readonly now: () => number;
27  private readonly buckets = new Map<string, { tokens: number; last: number }>();
28 
29  constructor(opts: TokenBucketOpts) {
30    this.capacity = opts.capacity;
31    this.rate = opts.rate;
32    this.now = opts.now ?? Date.now;
33  }
34 
35  private refill(key: string, t: number): { tokens: number; last: number } {
36    const b = this.buckets.get(key);
37    if (b === undefined) {
38      const bucket = { tokens: this.capacity, last: t };
39      this.buckets.set(key, bucket);
40      return bucket;
41    }
42    const elapsed = Math.max(0, t - b.last);
43    b.tokens = Math.min(this.capacity, b.tokens + (elapsed / 1000) * this.rate);
44    b.last = t;
45    return b;
46  }
47 
48  /** Consume one token. Returns `{ok}` or `{ok:false, retryAfterMs}`. */
49  consume(key: string): { ok: true } | { ok: false; retryAfterMs: number } {
50    const b = this.refill(key, this.now());
51    if (b.tokens >= 1) {
52      b.tokens -= 1;
53      return { ok: true };
54    }
55    const deficit = 1 - b.tokens;
56    const retryAfterMs = Math.ceil((deficit / this.rate) * 1000);
57    return { ok: false, retryAfterMs };
58  }
59 
60  /** Peek without consuming. */
61  available(key: string): number {
62    return this.refill(key, this.now()).tokens;
63  }
64 
65  /** Drop a bucket (socket close / prune). */
66  drop(key: string): void {
67    this.buckets.delete(key);
68  }
69 
70  /** Bucket count (metrics). */
71  size(): number {
72    return this.buckets.size;
73  }
74 
75  /** Drop buckets idle longer than `staleMs` (housekeeping). */
76  sweep(staleMs: number): number {
77    const t = this.now();
78    let dropped = 0;
79    for (const [key, b] of this.buckets) {
80      if (t - b.last > staleMs) {
81        this.buckets.delete(key);
82        dropped += 1;
83      }
84    }
85    return dropped;
86  }
87}
88 
89/**
90 * Connect-time admission over a sliding window (30 conns/min/IP): counts
91 * distinct connect events per key inside the window.
92 */
93export class ConnectLimiter {
94  private readonly windowMs: number;
95  private readonly max: number;
96  private readonly now: () => number;
97  private readonly hits = new Map<string, number[]>();
98 
99  constructor(opts: { windowMs: number; max: number; now?: () => number }) {
100    this.windowMs = opts.windowMs;
101    this.max = opts.max;
102    this.now = opts.now ?? Date.now;
103  }
104 
105  /** Record a connect attempt; `ok:false` when the window is already full. */
106  tryAdmit(key: string): { ok: true } | { ok: false } {
107    const t = this.now();
108    const cutoff = t - this.windowMs;
109    const prev = (this.hits.get(key) ?? []).filter((x) => x > cutoff);
110    if (prev.length >= this.max) {
111      this.hits.set(key, prev);
112      return { ok: false };
113    }
114    prev.push(t);
115    this.hits.set(key, prev);
116    return { ok: true };
117  }
118 
119  /** Current in-window count for a key (tests/metrics). */
120  count(key: string): number {
121    const t = this.now();
122    const cutoff = t - this.windowMs;
123    return (this.hits.get(key) ?? []).filter((x) => x > cutoff).length;
124  }
125 
126  drop(key: string): void {
127    this.hits.delete(key);
128  }
129 
130  size(): number {
131    return this.hits.size;
132  }
133 
134  /** Drop keys idle longer than `staleMs` (housekeeping). */
135  sweep(staleMs: number): number {
136    const t = this.now();
137    let dropped = 0;
138    for (const [key, arr] of this.hits) {
139      if (arr.length === 0 || t - arr[arr.length - 1]! > staleMs) {
140        this.hits.delete(key);
141        dropped += 1;
142      }
143    }
144    return dropped;
145  }
146}
147