| 1 | import type { Env } from "../env.js"; |
| 2 | |
| 3 | /** |
| 4 | * Where a request's time goes, for the `Server-Timing` header that browsers' developer tools |
| 5 | * show: D1 queries and Durable Object calls, each counted and summed (parallel calls overlap, so |
| 6 | * a sum can pass the total), and the whole request. Signed-in callers get it on answers that |
| 7 | * aren't errors (app.ts): an error could tell a private repository from a missing one by how |
| 8 | * many lookups it took. |
| 9 | */ |
| 10 | export class Timing { |
| 11 | private readonly start = Date.now(); |
| 12 | private readonly spans = new Map<string, { n: number; ms: number }>(); |
| 13 | /** Where D1 answered from, as its results say (only `all`, `run` and `batch` results carry it). */ |
| 14 | private d1Where: string | null = null; |
| 15 | |
| 16 | /** Notes where a D1 result came from: its region, and whether the primary or a replica answered. */ |
| 17 | d1Result(r: unknown): void { |
| 18 | const meta = (r as { meta?: { served_by_region?: string; served_by_primary?: boolean } } | null)?.meta; |
| 19 | if (this.d1Where || !meta?.served_by_region || !/^[A-Z]{2,8}$/.test(meta.served_by_region)) return; |
| 20 | this.d1Where = `${meta.served_by_region} ${meta.served_by_primary ? "primary" : "replica"}`; |
| 21 | } |
| 22 | |
| 23 | /** Times `p` under `kind` and hands it back unchanged. */ |
| 24 | track<T>(kind: string, p: T): T { |
| 25 | if (!(p && typeof (p as { then?: unknown }).then === "function")) return p; |
| 26 | const t0 = Date.now(); |
| 27 | const done = () => { |
| 28 | const s = this.spans.get(kind) ?? { n: 0, ms: 0 }; |
| 29 | s.n++; |
| 30 | s.ms += Date.now() - t0; |
| 31 | this.spans.set(kind, s); |
| 32 | }; |
| 33 | (p as unknown as PromiseLike<unknown>).then(done, done); |
| 34 | return p; |
| 35 | } |
| 36 | |
| 37 | header(colo?: string): string { |
| 38 | const out = [`total;dur=${Date.now() - this.start}`]; |
| 39 | for (const [kind, s] of this.spans) { |
| 40 | const what = kind === "d1" ? `quer${s.n === 1 ? "y" : "ies"}${this.d1Where ? `, ${this.d1Where}` : ""}` : `call${s.n === 1 ? "" : "s"}`; |
| 41 | out.push(`${kind};dur=${s.ms};desc="${s.n} ${what}"`); |
| 42 | } |
| 43 | if (colo && /^[A-Z]{3}$/.test(colo)) out.push(`colo;desc="${colo}"`); |
| 44 | return out.join(", "); |
| 45 | } |
| 46 | } |
| 47 | |
| 48 | const STATEMENT_RUNS = ["first", "all", "run", "raw"] as const; |
| 49 | |
| 50 | /** A D1 statement whose runs are timed. The statement itself stays the real one, so `batch` still takes it. */ |
| 51 | function timedStatement(t: Timing, s: D1PreparedStatement): D1PreparedStatement { |
| 52 | const st = s as unknown as Record<string, (...a: unknown[]) => unknown>; |
| 53 | try { |
| 54 | for (const m of STATEMENT_RUNS) { |
| 55 | const run = st[m]; |
| 56 | if (typeof run === "function") |
| 57 | st[m] = (...a: unknown[]) => { |
| 58 | const p = t.track("d1", run.apply(s, a)); |
| 59 | if (m === "all" || m === "run") (p as Promise<unknown>).then((r) => t.d1Result(r), () => {}); |
| 60 | return p; |
| 61 | }; |
| 62 | } |
| 63 | const bind = st.bind; |
| 64 | if (typeof bind === "function") st.bind = (...a: unknown[]) => timedStatement(t, bind.apply(s, a) as D1PreparedStatement); |
| 65 | } catch { |
| 66 | // A statement that can't take own methods runs untimed. |
| 67 | } |
| 68 | return s; |
| 69 | } |
| 70 | |
| 71 | function timedDb(t: Timing, db: D1Database): D1Database { |
| 72 | return new Proxy(db, { |
| 73 | get(target, prop) { |
| 74 | const v = Reflect.get(target, prop, target) as unknown; |
| 75 | if (typeof v !== "function") return v; |
| 76 | if (prop === "prepare") return (q: string) => timedStatement(t, (v as (q: string) => D1PreparedStatement).call(target, q)); |
| 77 | if (prop === "batch" || prop === "exec") |
| 78 | return (...a: unknown[]) => { |
| 79 | const p = t.track("d1", (v as (...a: unknown[]) => unknown).apply(target, a)); |
| 80 | if (prop === "batch") (p as Promise<unknown[]>).then((rs) => t.d1Result(rs?.[0]), () => {}); |
| 81 | return p; |
| 82 | }; |
| 83 | return (v as (...a: unknown[]) => unknown).bind(target); |
| 84 | }, |
| 85 | }); |
| 86 | } |
| 87 | |
| 88 | /** Stub members that aren't calls to the object (or must stay as they are). */ |
| 89 | const NOT_CALLS = new Set<PropertyKey>(["then", "catch", "finally", "dup", "id", "name", "connect"]); |
| 90 | |
| 91 | function timedStub<S extends object>(t: Timing, stub: S): S { |
| 92 | return new Proxy(stub, { |
| 93 | get(target, prop) { |
| 94 | const v = Reflect.get(target, prop, target) as unknown; |
| 95 | if (typeof v !== "function" || typeof prop === "symbol" || NOT_CALLS.has(prop)) return v; |
| 96 | // Reflect.apply: the stub's members are RPC proxies, whose own `apply` would be a call to the object. |
| 97 | return (...a: unknown[]) => t.track("do", Reflect.apply(v as (...a: unknown[]) => unknown, target, a)); |
| 98 | }, |
| 99 | }); |
| 100 | } |
| 101 | |
| 102 | function timedNamespace<N extends object>(t: Timing, ns: N): N { |
| 103 | return new Proxy(ns, { |
| 104 | get(target, prop) { |
| 105 | const v = Reflect.get(target, prop, target) as unknown; |
| 106 | if (typeof v !== "function") return v; |
| 107 | const f = (v as (...a: unknown[]) => unknown).bind(target); |
| 108 | return prop === "get" || prop === "getByName" ? (...a: unknown[]) => timedStub(t, f(...a) as object) : f; |
| 109 | }, |
| 110 | }); |
| 111 | } |
| 112 | |
| 113 | /** `env` with its D1 database and the repository and agent-task Durable Objects timed into `t`. */ |
| 114 | export function timedEnv(env: Env, t: Timing): Env { |
| 115 | return { |
| 116 | ...env, |
| 117 | DB: env.DB ? timedDb(t, env.DB) : env.DB, |
| 118 | REPO: env.REPO ? timedNamespace(t, env.REPO) : env.REPO, |
| 119 | PROJECT: env.PROJECT ? timedNamespace(t, env.PROJECT) : env.PROJECT, |
| 120 | }; |
| 121 | } |