workers/ask-budget.ts

Functions

Imports

It imports no other module of this repo. Imported by workers/app.ts.

Source

1import { RateLimiter } from "@dustinedwards/site-runtime/rate-limit/durable-object";
2
3/**
4 * The per-key rate limits are the shared class's `hit` (@dustinedwards/site-runtime/rate-limit), which this class
5 * extends so the binding and the stored class keep their name and need no migration; `consume` is the
6 * site-wide daily budget, which stays here. Not the `ratelimit` binding or KV: the binding is documented
7 * as eventually consistent and not for accounting, and KV read-modify-writes race. Synchronous SQL,
8 * because a DO sequence spanning `await` is not atomic.
9 */
10export class AskBudget extends RateLimiter<Env> {
11  constructor(ctx: DurableObjectState, env: Env) {
12    super(ctx, env);
13    // `day` is the UTC date of the daily budget. Rows an earlier per-IP limit left here are never read again.
14    this.ctx.storage.sql.exec(
15      `CREATE TABLE IF NOT EXISTS budget (day TEXT PRIMARY KEY, count INTEGER NOT NULL)`,
16    );
17  }
18
19  /**
20   * Counts one use in the window `key` unless `limit` is already reached, and forgets every other
21   * window. No `await` anywhere in the body, or the count can be raced.
22   */
23  private increment(key: string, limit: number): { ok: boolean; count: number } {
24    const rows = this.ctx.storage.sql
25      .exec<{ count: number }>(`SELECT count FROM budget WHERE day = ?`, key)
26      .toArray();
27    const current = rows[0]?.count ?? 0;
28
29    if (current >= limit) {
30      return { ok: false, count: current };
31    }
32
33    const next = current + 1;
34    this.ctx.storage.sql.exec(
35      `INSERT INTO budget (day, count) VALUES (?, ?)
36         ON CONFLICT(day) DO UPDATE SET count = excluded.count`,
37      key,
38      next,
39    );
40    this.ctx.storage.sql.exec(`DELETE FROM budget WHERE day <> ?`, key);
41    return { ok: true, count: next };
42  }
43
44  /** The site-wide daily budget, keyed by the UTC date. */
45  consume(limit: number): { ok: boolean; spent: number } {
46    const { ok, count } = this.increment(new Date().toISOString().slice(0, 10), limit);
47    return { ok, spent: count };
48  }
49
50  peek(): { day: string; count: number } {
51    const day = new Date().toISOString().slice(0, 10);
52    const rows = this.ctx.storage.sql
53      .exec<{ count: number }>(`SELECT count FROM budget WHERE day = ?`, day)
54      .toArray();
55    return { day, count: rows[0]?.count ?? 0 };
56  }
57
58  reset(): void {
59    this.ctx.storage.sql.exec(`DELETE FROM budget`);
60  }
61}