46 lines
1.3 KiB
TypeScript
46 lines
1.3 KiB
TypeScript
/**
|
|
* In-memory rate limiting with 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.
|
|
*
|
|
* @module
|
|
*/
|
|
|
|
interface Bucket {
|
|
tokens: number;
|
|
at: number;
|
|
}
|
|
const buckets = new Map<string, Bucket>();
|
|
|
|
/** A token bucket's shape. */
|
|
export interface RateLimit {
|
|
/** Burst size. */
|
|
capacity: number;
|
|
/** Tokens regained per second. */
|
|
refill: number;
|
|
}
|
|
|
|
/** Takes one token from `key`'s bucket, or returns false when it 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);
|