40 lines
1.2 KiB
TypeScript
40 lines
1.2 KiB
TypeScript
// Token buckets keyed by the caller's choice (e.g. `${kind}:${userId}`) — per
|
|
// user rather than per socket or request, so more tabs or parallel requests
|
|
// don't buy more throughput.
|
|
interface Bucket {
|
|
tokens: number;
|
|
at: number;
|
|
}
|
|
const buckets = new Map<string, Bucket>();
|
|
|
|
export interface RateLimit {
|
|
/** Burst size. */
|
|
capacity: number;
|
|
/** Tokens regained per second. */
|
|
refill: number;
|
|
}
|
|
|
|
/** Takes one token for `key`; false when the bucket is empty. */
|
|
export function takeToken(key: string, limit: RateLimit): boolean {
|
|
const now = Date.now();
|
|
const bucket = buckets.get(key) ?? { tokens: limit.capacity, at: now };
|
|
bucket.tokens = Math.min(
|
|
limit.capacity,
|
|
bucket.tokens + ((now - bucket.at) / 1000) * limit.refill,
|
|
);
|
|
bucket.at = now;
|
|
buckets.set(key, bucket);
|
|
if (bucket.tokens < 1) return false;
|
|
bucket.tokens -= 1;
|
|
return true;
|
|
}
|
|
|
|
// Idle buckets have refilled by now for every limit in use; drop them so the
|
|
// map tracks only keys that have been active lately.
|
|
setInterval(() => {
|
|
const cutoff = Date.now() - 60 * 60_000;
|
|
for (const [key, bucket] of buckets) {
|
|
if (bucket.at < cutoff) buckets.delete(key);
|
|
}
|
|
}, 60_000);
|