rate-limit/core.mjs

Functions

Imports

It imports no other module of this repo. Imported by rate-limit/durable-object.mjs, rate-limit/index.mjs.

Source

1/**
2 * The counting, apart from the Durable Object so a test can run it on any SQLite. Fixed windows aligned
3 * to the epoch, one row per window length. All rules are checked first and none is counted when any is
4 * spent, so a refused request costs a read and no write. No `await` anywhere: a Durable Object sequence
5 * that spans one is not atomic.
6 *
7 * @typedef {{ limit: number, windowSeconds: number }} Rule
8 * @typedef {{ exec: (query: string, ...bindings: unknown[]) => { toArray: () => any[] } }} Sql
9 * @typedef {{ ok: boolean, used: number, limit: number, retryAfterSeconds: number, window: number }} Hit
10 */
11
12const TABLE_SQL =
13  "CREATE TABLE IF NOT EXISTS rate_window (window_seconds INTEGER PRIMARY KEY, window_index INTEGER NOT NULL, count INTEGER NOT NULL)";
14
15/** @param {Rule[]} rules */
16export function assertRules(rules) {
17  if (!Array.isArray(rules) || rules.length === 0) throw new TypeError("rules must be a non-empty array");
18  const seen = new Set();
19  for (const rule of rules) {
20    if (!Number.isInteger(rule?.limit) || rule.limit < 1) throw new TypeError(`limit must be a positive integer, got ${rule?.limit}`);
21    if (!Number.isInteger(rule?.windowSeconds) || rule.windowSeconds < 1) throw new TypeError(`windowSeconds must be a positive integer, got ${rule?.windowSeconds}`);
22    if (seen.has(rule.windowSeconds)) throw new TypeError(`two rules share the ${rule.windowSeconds} second window`);
23    seen.add(rule.windowSeconds);
24  }
25}
26
27/**
28 * Seconds left in the window containing `nowMs`, rounded up, so a client that waits that long is
29 * in the next window.
30 * @param {number} windowIndex
31 * @param {number} windowSeconds
32 * @param {number} nowMs
33 */
34export function secondsLeft(windowIndex, windowSeconds, nowMs) {
35  return Math.max(1, Math.ceil(((windowIndex + 1) * windowSeconds * 1000 - nowMs) / 1000));
36}
37
38/**
39 * Counts one use against every rule unless one is spent. When one is spent the answer is that rule's
40 * (the one that frees up last, if several), and nothing is written.
41 *
42 * @param {Sql} sql
43 * @param {Rule[]} rules
44 * @param {number} nowMs
45 * @returns {Hit}
46 */
47export function hitRules(sql, rules, nowMs) {
48  assertRules(rules);
49  sql.exec(TABLE_SQL);
50  const state = rules.map((rule) => {
51    const windowIndex = Math.floor(nowMs / 1000 / rule.windowSeconds);
52    const row = sql.exec("SELECT window_index, count FROM rate_window WHERE window_seconds = ?", rule.windowSeconds).toArray()[0];
53    const used = row && row.window_index === windowIndex ? row.count : 0;
54    return { rule, windowIndex, used };
55  });
56
57  const spent = state.filter((s) => s.used >= s.rule.limit);
58  if (spent.length > 0) {
59    const last = spent.reduce((a, b) =>
60      secondsLeft(a.windowIndex, a.rule.windowSeconds, nowMs) >= secondsLeft(b.windowIndex, b.rule.windowSeconds, nowMs) ? a : b,
61    );
62    return {
63      ok: false,
64      used: last.used,
65      limit: last.rule.limit,
66      retryAfterSeconds: secondsLeft(last.windowIndex, last.rule.windowSeconds, nowMs),
67      window: last.rule.windowSeconds,
68    };
69  }
70
71  for (const s of state) {
72    sql.exec(
73      "INSERT INTO rate_window (window_seconds, window_index, count) VALUES (?, ?, ?) ON CONFLICT(window_seconds) DO UPDATE SET window_index = excluded.window_index, count = excluded.count",
74      s.rule.windowSeconds,
75      s.windowIndex,
76      s.used + 1,
77    );
78  }
79  // Report the rule closest to its limit, so a caller can show the tightest budget.
80  const tightest = state.reduce((a, b) => (a.rule.limit - a.used <= b.rule.limit - b.used ? a : b));
81  return {
82    ok: true,
83    used: tightest.used + 1,
84    limit: tightest.rule.limit,
85    retryAfterSeconds: secondsLeft(tightest.windowIndex, tightest.rule.windowSeconds, nowMs),
86    window: tightest.rule.windowSeconds,
87  };
88}
89
90/**
91 * The moment the longest window ends, when an idle instance can forget everything.
92 * @param {Rule[]} rules
93 * @param {number} nowMs
94 */
95export function idleDeadline(rules, nowMs) {
96  return Math.max(...rules.map((r) => (Math.floor(nowMs / 1000 / r.windowSeconds) + 1) * r.windowSeconds * 1000));
97}
98
99/**
100 * Forgets the windows that have ended and says when the rest end, or null when nothing is left. An
101 * alarm that fires early, or against a clock that moved, leaves a live window alone.
102 *
103 * @param {Sql} sql
104 * @param {number} nowMs
105 * @returns {number | null}
106 */
107export function purgeExpired(sql, nowMs) {
108  sql.exec(TABLE_SQL);
109  sql.exec("DELETE FROM rate_window WHERE (window_index + 1) * window_seconds * 1000 <= ?", nowMs);
110  const rows = sql.exec("SELECT MAX((window_index + 1) * window_seconds * 1000) AS ends FROM rate_window").toArray();
111  return rows[0]?.ends ?? null;
112}