| 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 | |
| 14 | /** Times `p` under `kind` and hands it back unchanged. */ |
| 15 | track<T>(kind: string, p: T): T { |
| 16 | if (!(p && typeof (p as { then?: unknown }).then === "function")) return p; |
| 17 | const t0 = Date.now(); |
| 18 | const done = () => { |
| 19 | const s = this.spans.get(kind) ?? { n: 0, ms: 0 }; |
| 20 | s.n++; |
| 21 | s.ms += Date.now() - t0; |
| 22 | this.spans.set(kind, s); |
| 23 | }; |
| 24 | (p as unknown as PromiseLike<unknown>).then(done, done); |
| 25 | return p; |
| 26 | } |
| 27 | |
| 28 | header(colo?: string): string { |
| 29 | const out = [`total;dur=${Date.now() - this.start}`]; |
| 30 | for (const [kind, s] of this.spans) out.push(`${kind};dur=${s.ms};desc="${s.n} ${kind === "d1" ? "quer" + (s.n === 1 ? "y" : "ies") : "call" + (s.n === 1 ? "" : "s")}"`); |
| 31 | if (colo && /^[A-Z]{3}$/.test(colo)) out.push(`colo;desc="${colo}"`); |
| 32 | return out.join(", "); |
| 33 | } |
| 34 | } |
| 35 | |
| 36 | const STATEMENT_RUNS = ["first", "all", "run", "raw"] as const; |
| 37 | |
| 38 | /** A D1 statement whose runs are timed. The statement itself stays the real one, so `batch` still takes it. */ |
| 39 | function timedStatement(t: Timing, s: D1PreparedStatement): D1PreparedStatement { |
| 40 | const st = s as unknown as Record<string, (...a: unknown[]) => unknown>; |
| 41 | try { |
| 42 | for (const m of STATEMENT_RUNS) { |
| 43 | const run = st[m]; |
| 44 | if (typeof run === "function") st[m] = (...a: unknown[]) => t.track("d1", run.apply(s, a)); |
| 45 | } |
| 46 | const bind = st.bind; |
| 47 | if (typeof bind === "function") st.bind = (...a: unknown[]) => timedStatement(t, bind.apply(s, a) as D1PreparedStatement); |
| 48 | } catch { |
| 49 | // A statement that can't take own methods runs untimed. |
| 50 | } |
| 51 | return s; |
| 52 | } |
| 53 | |
| 54 | function timedDb(t: Timing, db: D1Database): D1Database { |
| 55 | return new Proxy(db, { |
| 56 | get(target, prop) { |
| 57 | const v = Reflect.get(target, prop, target) as unknown; |
| 58 | if (typeof v !== "function") return v; |
| 59 | if (prop === "prepare") return (q: string) => timedStatement(t, (v as (q: string) => D1PreparedStatement).call(target, q)); |
| 60 | if (prop === "batch" || prop === "exec") return (...a: unknown[]) => t.track("d1", (v as (...a: unknown[]) => unknown).apply(target, a)); |
| 61 | return (v as (...a: unknown[]) => unknown).bind(target); |
| 62 | }, |
| 63 | }); |
| 64 | } |
| 65 | |
| 66 | /** Stub members that aren't calls to the object (or must stay as they are). */ |
| 67 | const NOT_CALLS = new Set<PropertyKey>(["then", "catch", "finally", "dup", "id", "name", "connect"]); |
| 68 | |
| 69 | function timedStub<S extends object>(t: Timing, stub: S): S { |
| 70 | return new Proxy(stub, { |
| 71 | get(target, prop) { |
| 72 | const v = Reflect.get(target, prop, target) as unknown; |
| 73 | if (typeof v !== "function" || typeof prop === "symbol" || NOT_CALLS.has(prop)) return v; |
| 74 | // Reflect.apply: the stub's members are RPC proxies, whose own `apply` would be a call to the object. |
| 75 | return (...a: unknown[]) => t.track("do", Reflect.apply(v as (...a: unknown[]) => unknown, target, a)); |
| 76 | }, |
| 77 | }); |
| 78 | } |
| 79 | |
| 80 | function timedNamespace<N extends object>(t: Timing, ns: N): N { |
| 81 | return new Proxy(ns, { |
| 82 | get(target, prop) { |
| 83 | const v = Reflect.get(target, prop, target) as unknown; |
| 84 | if (typeof v !== "function") return v; |
| 85 | const f = (v as (...a: unknown[]) => unknown).bind(target); |
| 86 | return prop === "get" || prop === "getByName" ? (...a: unknown[]) => timedStub(t, f(...a) as object) : f; |
| 87 | }, |
| 88 | }); |
| 89 | } |
| 90 | |
| 91 | /** `env` with its D1 database and the repository and agent-task Durable Objects timed into `t`. */ |
| 92 | export function timedEnv(env: Env, t: Timing): Env { |
| 93 | return { |
| 94 | ...env, |
| 95 | DB: env.DB ? timedDb(t, env.DB) : env.DB, |
| 96 | REPO: env.REPO ? timedNamespace(t, env.REPO) : env.REPO, |
| 97 | PROJECT: env.PROJECT ? timedNamespace(t, env.PROJECT) : env.PROJECT, |
| 98 | }; |
| 99 | } |