diff --git a/server/BENCHMARKING.md b/server/BENCHMARKING.md index ea235783b..7ef0e8561 100644 --- a/server/BENCHMARKING.md +++ b/server/BENCHMARKING.md @@ -36,6 +36,10 @@ bun run bench:backend --compare .test-results/before.json --output .test-results # Explore sensitivity to additional DB transport latency on one path. bun run bench:backend --scenario send.dm.open --rtt 0,3,5 +# Compare safe default preparation with qualified broad preparation. +bun run bench:backend --preparation access --scenario send.dm.open --rtt 0,3,5 --output .test-results/access.json +bun run bench:backend --preparation all-custom --scenario send.dm.open --rtt 0,3,5 --compare .test-results/access.json --output .test-results/all-custom.json + # Check the entire benchmark workflow quickly; this is not a tail-latency study. bun run bench:backend --all --samples 1 --warmup 1 ``` @@ -47,6 +51,20 @@ never overwritten. Unknown selections, invalid arguments, failed assertions, background failures, failed cleanup and incomplete runs exit unsuccessfully. Interrupted worker artifacts are candidates, not accepted comparison baselines. +`--preparation access` is the default; shell rollout flags are deliberately not +inherited. `--preparation all-custom` sets `force_custom_plan` only as a database +default on the lifecycle's disposable local database. It then recreates the +application pools and validates startup before measuring. The provisioning +database and role defaults are never changed. Use at least three warmups when +studying retained broad statements, which require two observations for admission. + +Reports record the requested mode, the actual `plan_cache_mode`, and statement/ +execution counters from one sampled backend after the matrix. These counters +include driver-owned statements and are not an admission count or pool total. +Access-to-all-custom comparisons display both treatments and planning settings. +Legacy reports without preparation metadata remain readable but cannot compare +to new qualified reports; rerun those baselines with an explicit mode. + Run timing experiments alone, on the same otherwise idle machine, database version and power configuration. Avoid concurrent builds and tests. Repeat A/B measurements before claiming a small improvement. Keep reports with the change; @@ -77,6 +95,9 @@ Each iteration follows this sequence: | `commandsBeforeReturn` | Commands submitted by return time; background work can overlap, so this is **not** a foreground-only count | | `wire.exchangeBoundaries` | Frontend `Query`, `Flush`, and `Sync` frames observed at the proxy | | `calibrationMs` | Ten real `SELECT 1` timings at the requested delay, outside the operation samples | +| `preparation.mode` / `planCacheMode` | Explicit benchmark treatment and actual backend planning policy | +| `preparation.sampledBackendStatements` | Retained statements visible on one backend after the matrix; includes driver-owned queries | +| `preparation.sampledBackendCustomPlans` / `sampledBackendGenericPlans` | Execution counters on that sampled backend, rather than pool-wide totals | One driver submission can contain several SQL statements, so `sql.commands` is not a SQL parser's statement count. A transaction does not make its statements @@ -123,10 +144,12 @@ Do not disable domain jobs to make a path pass its budget. `catalog.ts` holds reviewed maximum command counts. They represent current costs, including costly linear fanout; they are not desired targets. Lower them with a proven optimization. Do not automatically regenerate budgets or loosen one to -hide a regression. The 1/10/100 cases expose scaling: current public sends cost -`32 + recipients`, reply sends `66 + 7 × recipients`, and user update batches -`2 + 4 × updates` for these fixtures. Chat-list queries remain constant at 11 -while result volume grows. These are local fixture observations, not production +hide a regression. The 1/10/100 cases expose scaling: the reviewed public-send budget is now constant at +`32` after batching sidebar owner locks, reply sends retain a conservative +`66 + 7 × recipients` bound, and user update batches +`2 + 4 × updates` for these fixtures. Unchanged dialog batches cost 5 commands at 1/10/100 owners. Chat-list +command budgets include their snapshot and admission overhead while remaining +constant as result volume grows. These are local fixture observations, not production traffic statistics. `test:backend` also selects the existing send, chat-list, history, message integrity, @@ -191,3 +214,47 @@ matrix, runtime/driver/database/platform/CPU/pool settings, warmups and sample counts. They report saved commands, exchanges and milliseconds separately, along with before/after RTT calibration. A comparison is evidence to inspect, not an automatic correctness proof or a promise of production latency savings. + +## Prepared query policy + +The default `INLINE_DATABASE_PREPARED_QUERIES=access` prepares only the existing +`sync.replay` and `sync.sidecars` access family. Its 100 batch lengths per source +fit a 200-signature budget. Other Drizzle queries retain unnamed execution and +custom planning. Rows and authorization decisions are never cached. + +`INLINE_DATABASE_PREPARED_QUERIES=all-custom` opts into broader scalar/ordinary-array +preparation, including transactions, savepoints and reserved clients. Startup +must first observe `plan_cache_mode=force_custom_plan` on the actual query path; +readiness checks it again. This requires an operator-qualified application role +default on every backend, with existing pooled sessions reestablished after the +role change. PgBouncer does not accept this as an ordinary startup parameter. +The application never changes roles or session settings. A sampled startup check +cannot establish that every existing backend has inherited the role default. + +Broad admission requires a second observation, with at most 200 pending candidates. +This avoids retaining one-off SQL shapes, such as transient `inArray` lengths. +It admits at most 100 broad signatures and leaves 100 slots for access queries, +within the same total cap of 200. Retained signatures live until pool shutdown; +they are not an LRU. Once a partition or total budget fills, new shapes use the +original unnamed path. Changing hot lists to a single array binding avoids a +new signature for each list length. The default access mode keeps its full 200 +slots. Explicit per-query options and the pool-level preparation switch remain +authoritative. Queries longer than 32 KiB, more than 301 parameters, and custom +typed parameters retain their prior mode. + +Warm admitted calls skip Describe/Flush while PostgreSQL builds custom plans +from current bindings. The wire fixture opts into `all-custom` using a dedicated +direct connection with the setting established before startup. The benchmark +launcher provides its own disposable-database qualification via `--preparation`. Database tests +assert partial-index selection and zero generic plans after more than five +executions, plus rollback, changed array bindings and admission pressure. + +Measurements from commit `5dc49a0f2` used unconditional broad preparation under +the database's prior planning policy. They do not describe this safe default or +qualify `all-custom` latency. New comparisons must record the preparation mode, +role setting, admission limits and enough warmups (at least three) for repeated +shapes; inspect `pg_prepared_statements` counters to verify prepared reuse and planning policy. + +Role-wide custom planning also applies to the previously tuned access queries. +Qualify representative replay latency as well as simple point reads before +enabling broad mode; faster protocol exchanges alone do not prove a net gain. diff --git a/server/bench/backend.test.ts b/server/bench/backend.test.ts index 597034fc8..2ed4aa730 100644 --- a/server/bench/backend.test.ts +++ b/server/bench/backend.test.ts @@ -4,18 +4,20 @@ import { createHash } from "node:crypto" import { cpus } from "node:os" import { resolve } from "node:path" import { sql } from "drizzle-orm" -import { closeDb, db, initDb } from "@in/server/db" +import { closeDb, db, initDb, validateDatabaseStartup } from "@in/server/db" import { setupTestDatabase, teardownTestDatabase } from "@in/server/__tests__/database" import { scenarios, scenarioVersion, selectScenarios } from "@in/server/__tests__/performance/catalog" import { prepareScenario } from "@in/server/__tests__/performance/scenarios" import { observeScenarioWork } from "@in/server/__tests__/performance/observe" import { assertQueryBudget, measureOperation, type Sample } from "@in/server/__tests__/performance/measure" import { createWireProxy } from "@in/server/__tests__/performance/wire" -import { summarize, type BenchmarkReport } from "@in/server/__tests__/performance/report" +import { summarize, type BenchmarkReport, type BenchmarkPreparation } from "@in/server/__tests__/performance/report" const encoded = process.env["INLINE_BACKEND_BENCH_CONFIG"] if (!encoded) throw new Error("Use bun run bench:backend; this worker needs the isolated benchmark launcher") -const config = JSON.parse(encoded) as { ids: string[]; samples: number; warmup: number; rtts: number[]; output: string } +const config = JSON.parse(encoded) as { ids: string[]; samples: number; warmup: number; rtts: number[]; preparation: BenchmarkPreparation["mode"]; output: string } +if (config.preparation !== "access" && config.preparation !== "all-custom") throw new Error("Invalid benchmark preparation mode") +const previousPreparation = process.env["INLINE_DATABASE_PREPARED_QUERIES"] const selected = selectScenarios(config.ids) const background = observeScenarioWork() let proxy: Awaited> | undefined @@ -23,15 +25,36 @@ let proxy: Awaited> | undefined beforeAll(async () => { await setupTestDatabase() const url = new URL(process.env["DATABASE_URL"]!) + if (config.preparation === "all-custom") { + // Only the lifecycle's freshly cloned database may receive this default. + // No role change, provisioning database change, or session-only shortcut. + const template = process.env["INLINE_TEST_DATABASE_TEMPLATE"] + const name = url.pathname.slice(1) + if (!template || !/^inline_test_template_[0-9]+_[0-9a-f]{32}$/.test(template) || + !new RegExp(`^test_db_${template.slice(-32)}_[0-9a-f]{12}$`).test(name) || + !["localhost", "127.0.0.1"].includes(url.hostname)) { + throw new Error("Refusing to configure an unowned benchmark database") + } + const [current] = await db.execute<{ name: string }>(sql`select current_database() as name`) + if (current?.name !== name) throw new Error("Benchmark database identity changed") + await db.$client.unsafe(`ALTER DATABASE "${name}" SET plan_cache_mode = force_custom_plan`) + } + process.env["INLINE_DATABASE_PREPARED_QUERIES"] = config.preparation proxy = await createWireProxy({ host: url.hostname, port: Number(url.port || 5432) }, 0) url.hostname = "127.0.0.1" url.port = String(proxy.port) await closeDb() initDb(url.toString()) + await validateDatabaseStartup() }) afterAll(async () => { try { await background.close() } finally { - try { await teardownTestDatabase() } finally { await proxy?.close() } + try { await teardownTestDatabase() } finally { + try { await proxy?.close() } finally { + if (previousPreparation === undefined) Reflect.deleteProperty(process.env, "INLINE_DATABASE_PREPARED_QUERIES") + else process.env["INLINE_DATABASE_PREPARED_QUERIES"] = previousPreparation + } + } } }) @@ -107,6 +130,24 @@ test("verified backend benchmark matrix", async () => { if (report.results.length !== selected.length * config.rtts.length || selected.some((entry) => !scenarios.includes(entry))) { throw new Error("Incomplete benchmark matrix") } + const [settings] = await db.execute<{ mode: string; statements: number; custom: string; generic: string }>(sql` + select current_setting('plan_cache_mode') as mode, count(*)::integer as statements, + coalesce(sum(custom_plans), 0)::text as custom, coalesce(sum(generic_plans), 0)::text as generic + from pg_prepared_statements + `) + if (!settings || !["auto", "force_custom_plan", "force_generic_plan"].includes(settings.mode)) { + throw new Error("Cannot record benchmark planning policy") + } + if (config.preparation === "all-custom" && settings.mode !== "force_custom_plan") { + throw new Error("Benchmark lost custom-plan qualification") + } + report.preparation = { + mode: config.preparation, planCacheMode: settings.mode as BenchmarkPreparation["planCacheMode"], + sampledBackendStatements: settings.statements, + sampledBackendCustomPlans: Number(settings.custom), sampledBackendGenericPlans: Number(settings.generic), + } + console.info(`Preparation ${report.preparation.mode}; plan_cache_mode=${report.preparation.planCacheMode}; ` + + `sampled backend: ${settings.statements} retained statements, ${settings.custom} custom / ${settings.generic} generic executions`) if (await sourceDigest() !== digest) throw new Error("Backend sources changed during the benchmark; rerun on a stable checkout") // The launcher promotes this only after Bun hooks and DB disposal also pass. writeFileSync(config.output, JSON.stringify({ ...report, status: "candidate" }, null, 2) + "\n", { flag: "wx" }) diff --git a/server/docs/database-observability.md b/server/docs/database-observability.md index c56acafb3..3fe146ecb 100644 --- a/server/docs/database-observability.md +++ b/server/docs/database-observability.md @@ -84,15 +84,26 @@ transactions and deployment overlap. The shared query client now prepares only tagged `sync.replay` and `sync.sidecars` access queries with at most 301 parameters (100 chat IDs used three times, plus one user). Their two sources yield at most 200 statement -shapes per connection. Larger batches and transaction-local queries retain the -existing driver behavior. Explicit `prepare: false` is respected. This caches +shapes per connection. Larger batches retain the existing driver behavior. The same bounded policy +also applies to transactions, savepoints and reserved clients. Explicit `prepare: false` is respected. This caches plans, never permissions or results. Production PgBouncer was verified to use transaction mode with `max_prepared_statements=200`; keep this capability when changing providers or pool configuration. Local integration tests exercise actual Drizzle prepared reads through PgBouncer, and authorization tests cover revocation after warm repeated reads. Array-binding alternatives were rejected -after live probes showed PostgreSQL kept replanning them instead of selecting -the much faster generic plan for the existing query shape. +after historical live probes showed PostgreSQL kept replanning them instead of +selecting the faster generic plan for that query shape. This is evidence for the +narrow family, not for enabling generic plans across all queries. + +Broader preparation requires explicit `INLINE_DATABASE_PREPARED_QUERIES=all-custom` +and an operator-established `force_custom_plan` default on the query role. The +startup and readiness probes inspect `plan_cache_mode` through the query pool; +startup rejects unqualified opt-ins, and readiness disables broad admission if +the sampled backend loses qualification. Establish the role default and recycle +existing PgBouncer backend sessions before enabling it: sampling one backend +does not prove the setting on every backend. No role or session settings are +changed by this code. See [the prepared-query policy](../BENCHMARKING.md#prepared-query-policy) +for the bounded second-observation admission policy and benchmark limits. Replay access checks execute the canonical SELECT directly on the shared client. The former single-query explicit transaction added BEGIN/SET/COMMIT diff --git a/server/scripts/benchmark-backend.test.ts b/server/scripts/benchmark-backend.test.ts index 3c7ef67f2..5a7da91b5 100644 --- a/server/scripts/benchmark-backend.test.ts +++ b/server/scripts/benchmark-backend.test.ts @@ -25,6 +25,7 @@ test.each([ ["--warmup", "NaN", "--warmup must"], ["--rtt", "0,,3", "Empty RTT"], ["--rtt", "-1", "RTT must"], + ["--preparation", "invalid", "--preparation must be access or all-custom"], ["--output", ".env.json", "never environment files"], ["--retry", "3", "Unknown option"], ])("invalid %s %s fails before provisioning: %s", async (flag, value, message) => { @@ -37,6 +38,7 @@ function report(): BenchmarkReport { return { schemaVersion: 1, status: "passed", createdAt: "2026-01-01T00:00:00.000Z", revision: { commit: "fixture", dirty: false, sourceDigest: "a".repeat(64) }, environment: { bun: "1.4.0", postgres: "15", postgresJs: "3.4.7", drizzle: "fixture", platform: "darwin", arch: "arm64", cpu: "fixture", poolMax: 10, clientPrepareDefault: true, concurrency: 1, cache: "cold-application-warm-pool" }, + preparation: { mode: "access", planCacheMode: "auto", sampledBackendStatements: 2, sampledBackendCustomPlans: 5, sampledBackendGenericPlans: 0 }, warmup: 1, sampleCount: 1, results: [{ id: "send.dm.open", version: 1, size: 1, rttMs: 3, maxCommands: 27, calibrationMs: [3], samples: [{ operationMs: 10, settledMs: 12, commandsBeforeReturn: 2, @@ -73,3 +75,36 @@ test("failed, empty, incomplete and incompatible runs cannot be compared as succ expect(() => compareReports(report(), changed)).toThrow() } }) + + +test.each(["access", "all-custom"])("explicit %s preparation can list without provisioning", async (mode) => { + const result = await run(`--preparation=${mode}`, "--list") + expect(result.code).toBe(0) + expect(result.output).toContain("send.public.100") +}) + +test("preparation comparisons identify the treatment and its actual planning policy", () => { + const before = report(), after = report() + after.preparation = { mode: "all-custom", planCacheMode: "force_custom_plan", sampledBackendStatements: 10, sampledBackendCustomPlans: 20, sampledBackendGenericPlans: 0 } + expect(compareReports(before, after)[0]).toMatchObject({ + preparationBefore: "access", preparationAfter: "all-custom", + planCacheModeBefore: "auto", planCacheModeAfter: "force_custom_plan", + }) +}) + +test("preparation metadata rejects unqualified modes and malformed backend counters", () => { + for (const metadata of [ + { mode: "unknown", planCacheMode: "auto", sampledBackendStatements: 0, sampledBackendCustomPlans: 0, sampledBackendGenericPlans: 0 }, + { mode: "all-custom", planCacheMode: "auto", sampledBackendStatements: 0, sampledBackendCustomPlans: 0, sampledBackendGenericPlans: 0 }, + { mode: "access", planCacheMode: "auto", sampledBackendStatements: -1, sampledBackendCustomPlans: 0, sampledBackendGenericPlans: 0 }, + { mode: "access", planCacheMode: "auto", sampledBackendStatements: 0 }, + ]) expect(() => validateReport({ ...report(), preparation: metadata })).toThrow("preparation") +}) + +test("legacy reports remain readable but cannot silently compare to a qualified mode", () => { + const legacy = report() + delete legacy.preparation + expect(() => validateReport(legacy)).not.toThrow() + expect(() => compareReports(legacy, report())).toThrow("Preparation policy is unknown") + expect(() => compareReports(legacy, structuredClone(legacy))).not.toThrow() +}) diff --git a/server/scripts/benchmark-backend.ts b/server/scripts/benchmark-backend.ts index ba4f9eb2e..6bcdb7aa3 100644 --- a/server/scripts/benchmark-backend.ts +++ b/server/scripts/benchmark-backend.ts @@ -11,16 +11,21 @@ const { values } = parseArgs({ args: process.argv.slice(2), strict: true, allowPositionals: false, options: { scenario: { type: "string", multiple: true }, all: { type: "boolean" }, list: { type: "boolean" }, - samples: { type: "string" }, warmup: { type: "string" }, rtt: { type: "string" }, + samples: { type: "string" }, warmup: { type: "string" }, rtt: { type: "string" }, preparation: { type: "string" }, output: { type: "string" }, compare: { type: "string" }, help: { type: "boolean", short: "h" }, }, }) +const preparation = values.preparation ?? "access" +if (preparation !== "access" && preparation !== "all-custom") { + throw new Error("--preparation must be access or all-custom") +} if (values.help) { console.info(`Usage: bun run bench:backend [--scenario ID ... | --all] [--rtt 0,3,5] --list List scenarios and reviewed SQL command budgets; no DB needed --samples N Measured, independently verified iterations (default 20, max 200) --warmup N Unreported, verified iterations per case (default 3, max 20) --rtt MS,... Requested additional loopback transport RTT (default 0, max 50) + --preparation MODE access (default) or all-custom; configures only disposable DB --output FILE.json Save a new report; never overwrite a baseline --compare FILE.json Compare a completed, compatible baseline with this run @@ -68,7 +73,7 @@ const environment = createTestEnvironment(process.env, databaseUrl) // must not silently change the deployment mode or invalidate that baseline. environment["INLINE_REALTIME_DISTRIBUTED"] = "0" environment["REALTIME_DISTRIBUTED"] = "0" -environment["INLINE_BACKEND_BENCH_CONFIG"] = JSON.stringify({ ids: selected.map((scenario) => scenario.id), samples, warmup, rtts, output: candidate }) +environment["INLINE_BACKEND_BENCH_CONFIG"] = JSON.stringify({ ids: selected.map((scenario) => scenario.id), samples, warmup, rtts, preparation, output: candidate }) let template: Awaited> | undefined let child: ReturnType | undefined diff --git a/server/src/__tests__/performance/catalog.ts b/server/src/__tests__/performance/catalog.ts index 618d3d191..3807a0c8b 100644 --- a/server/src/__tests__/performance/catalog.ts +++ b/server/src/__tests__/performance/catalog.ts @@ -18,10 +18,10 @@ const snapshotCommands = 3 const catalogAdmissionCommands = 2 const historySnapshotChatCommands = 1 export const scenarios: readonly ScenarioSpec[] = [ - { id: "send.dm.open", kind: "sendDm", size: 1, maxCommands: 26 }, + { id: "send.dm.open", kind: "sendDm", size: 1, maxCommands: 25 }, { id: "send.dm.closed", kind: "sendDm", size: 1, variant: "closed", maxCommands: 42 }, { id: "send.dm.retry", kind: "sendDm", size: 1, variant: "retry", maxCommands: 7 }, - ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `send.public.${size}`, kind: "sendThread", size, maxCommands: 32 + size })), + ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `send.public.${size}`, kind: "sendThread", size, maxCommands: 32 })), ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `send.reply.${size}`, kind: "sendThread", size, variant: "reply", maxCommands: 66 + 7 * size })), ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `getChats.${size}`, kind: "getChats", size, maxCommands: 11 + snapshotCommands + catalogAdmissionCommands })), { id: "history.50", kind: "history", size: 50, maxCommands: 7 + snapshotCommands + historySnapshotChatCommands }, @@ -29,7 +29,7 @@ export const scenarios: readonly ScenarioSpec[] = [ { id: "updates.replay.message", kind: "replay", size: 1, maxCommands: 23 }, { id: "updates.checkpoint", kind: "checkpoint", size: 1, maxCommands: 6 }, ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `updates.enqueue.${size}`, kind: "enqueue", size, maxCommands: 4 * size + 2 })), - ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `dialogs.noop.${size}`, kind: "dialog", size, maxCommands: size + 5 })), + ...[1, 10, 100].map((size): ScenarioSpec => ({ id: `dialogs.noop.${size}`, kind: "dialog", size, maxCommands: 5 })), { id: "read.advance", kind: "read", size: 1, maxCommands: 12 }, { id: "read.noop", kind: "read", size: 1, variant: "noop", maxCommands: 5 }, ] diff --git a/server/src/__tests__/performance/report.ts b/server/src/__tests__/performance/report.ts index 846136e27..dfbbb7109 100644 --- a/server/src/__tests__/performance/report.ts +++ b/server/src/__tests__/performance/report.ts @@ -10,6 +10,15 @@ export type BenchmarkResult = { samples: Sample[] } +export type BenchmarkPreparation = { + mode: "access" | "all-custom" + planCacheMode: "auto" | "force_custom_plan" | "force_generic_plan" + // pg_prepared_statements is local to one sampled backend, not a pool total. + sampledBackendStatements: number + sampledBackendCustomPlans: number + sampledBackendGenericPlans: number +} + export type BenchmarkReport = { schemaVersion: 1 status: "passed" @@ -20,6 +29,7 @@ export type BenchmarkReport = { platform: string; arch: string; cpu: string; poolMax: number clientPrepareDefault: boolean; concurrency: 1; cache: "cold-application-warm-pool" } + preparation?: BenchmarkPreparation warmup: number sampleCount: number results: BenchmarkResult[] @@ -38,6 +48,9 @@ export function summarize(result: BenchmarkResult) { export function compareReports(before: BenchmarkReport, after: BenchmarkReport) { validateReport(before) validateReport(after) + if (Boolean(before.preparation) !== Boolean(after.preparation)) { + throw new Error("Preparation policy is unknown in one report; rerun the baseline with explicit preparation metadata") + } if (JSON.stringify(before.environment) !== JSON.stringify(after.environment) || before.warmup !== after.warmup || before.sampleCount !== after.sampleCount) { throw new Error("Incompatible benchmark environment, warmup, or sample count; rerun the baseline with matching settings") } @@ -52,6 +65,10 @@ export function compareReports(before: BenchmarkReport, after: BenchmarkReport) const a = summarize(previous), b = summarize(result) return { scenario: result.id, rttMs: result.rttMs, + preparationBefore: before.preparation?.mode ?? "legacy-unknown", + preparationAfter: after.preparation?.mode ?? "legacy-unknown", + planCacheModeBefore: before.preparation?.planCacheMode ?? "unknown", + planCacheModeAfter: after.preparation?.planCacheMode ?? "unknown", commandsSaved: a.commands.median - b.commands.median, exchangesSaved: a.exchanges.median - b.exchanges.median, operationMsSaved: a.operationMs.median - b.operationMs.median, @@ -77,6 +94,15 @@ export function validateReport(value: unknown): asserts value is BenchmarkReport report.environment.cache !== "cold-application-warm-pool") { throw new Error("Missing benchmark provenance or measurement settings") } + const preparation = report.preparation + if (preparation !== undefined && (!preparation || + !["access", "all-custom"].includes(preparation.mode) || + !["auto", "force_custom_plan", "force_generic_plan"].includes(preparation.planCacheMode) || + (preparation.mode === "all-custom" && preparation.planCacheMode !== "force_custom_plan") || + [preparation.sampledBackendStatements, preparation.sampledBackendCustomPlans, preparation.sampledBackendGenericPlans] + .some((count) => !Number.isSafeInteger(count) || count < 0))) { + throw new Error("Invalid benchmark preparation policy or counters") + } const keys = new Set() for (const result of report.results) { const key = `${result.id}@${result.rttMs}` diff --git a/server/src/__tests__/performance/wire.postgres.test.ts b/server/src/__tests__/performance/wire.postgres.test.ts index f4a0518a6..c16bc76f8 100644 --- a/server/src/__tests__/performance/wire.postgres.test.ts +++ b/server/src/__tests__/performance/wire.postgres.test.ts @@ -1,6 +1,7 @@ import { expect, test } from "bun:test" import { sql } from "drizzle-orm" -import { db, closeDb, initDb } from "@in/server/db" +import { makeDatabaseClients } from "@in/server/db/connectionPolicy" +import { drizzle } from "drizzle-orm/postgres-js" import { setupTestLifecycle } from "../database" import { createWireProxy } from "./wire" import { measureOperation } from "./measure" @@ -13,9 +14,14 @@ test("real Drizzle queries expose protocol exchanges independently of SQL comman const proxy = await createWireProxy({ host: url.hostname, port: Number(url.port || 5432) }, 0) url.hostname = "127.0.0.1" url.port = String(proxy.port) - await closeDb() - initDb(url.toString()) + // This dedicated direct fixture qualifies the opt-in explicitly; production + // uses a role default because PgBouncer rejects this startup parameter. + url.searchParams.set("plan_cache_mode", "force_custom_plan") + const clients = makeDatabaseClients(url.toString(), { INLINE_DATABASE_PREPARED_QUERIES: "all-custom" }) + const db = drizzle(clients.queryClient) try { + await clients.validateStartup() + await db.execute(sql`select ${0}::int as value`) await db.execute(sql`select ${1}::int as value`) const sample = await measureOperation(db.$client.options, async () => { const result = await db.execute(sql`select ${2}::int as value`) @@ -23,11 +29,10 @@ test("real Drizzle queries expose protocol exchanges independently of SQL comman }, async () => {}, proxy) expect(sample.sql.commands).toBe(1) expect(sample.sql.catalogCommands).toBe(0) - // Drizzle currently calls unsafe() without per-query prepare:true. Its - // Describe/Flush and Bind/Execute/Sync produce two exchange boundaries. - // A future driver optimization may reduce this to one, never to zero. - expect(sample.wire!.exchangeBoundaries).toBeGreaterThanOrEqual(1) - expect(sample.wire!.exchangeBoundaries).toBeLessThanOrEqual(2) + // Warm parameterized Drizzle reads must skip Describe/Flush, not merely + // keep the SQL command count unchanged. + expect(sample.wire!.exchangeBoundaries).toBe(1) + expect(sample.wire!.frames["H"] ?? 0).toBe(0) expect(sample.wire!.frames["S"]).toBe(1) expect(sample.wire!.clientBytes).toBeGreaterThan(0) expect(sample.wire!.serverBytes).toBeGreaterThan(0) @@ -37,8 +42,7 @@ test("real Drizzle queries expose protocol exchanges independently of SQL comman expect(delayed.wire!.exchangeBoundaries).toBe(1) expect(delayed.operationMs).toBeGreaterThan(0) } finally { - await closeDb() + await clients.close() await proxy.close() - initDb(direct) } }) diff --git a/server/src/__tests__/performance/wire.test.ts b/server/src/__tests__/performance/wire.test.ts new file mode 100644 index 000000000..9a0bebb7f --- /dev/null +++ b/server/src/__tests__/performance/wire.test.ts @@ -0,0 +1,72 @@ +import { expect, test } from "bun:test" +import { connect, createServer, type Socket } from "node:net" +import { createWireProxy } from "./wire" + +function requestBytes() { + const startup = Buffer.alloc(8) + startup.writeUInt32BE(8) + startup.writeUInt32BE(196608, 4) + const query = Buffer.alloc(5 + 96 * 1024) + query[0] = "Q".charCodeAt(0) + query.writeUInt32BE(query.length - 1, 1) + query.fill(97, 5) + return { startup, query } +} + +test.each(["client", "server"] as const)("delayed proxy preserves queued bytes when %s half-closes first", async (first) => { + const { startup, query } = requestBytes() + const request = Buffer.concat([startup, query]) + const response = Buffer.alloc(256 * 1024, 87) + const received = Promise.withResolvers() + const peers = new Set() + const server = createServer({ allowHalfOpen: true }, (peer) => { + peers.add(peer) + const chunks: Buffer[] = [] + let replied = false + peer.on("data", (chunk: Buffer) => { + chunks.push(chunk) + if (first === "server" && !replied) { replied = true; peer.end(response) } + }) + peer.on("end", () => { + received.resolve(Buffer.concat(chunks)) + if (first === "client") peer.end(response) + }) + peer.on("error", received.reject) + peer.on("close", () => peers.delete(peer)) + }) + await new Promise((resolve, reject) => { + server.once("error", reject) + server.listen(0, "127.0.0.1", () => { server.off("error", reject); resolve() }) + }) + const address = server.address() + if (!address || typeof address === "string") throw new Error("Missing fixture port") + const proxy = await createWireProxy({ host: "127.0.0.1", port: address.port }, 50) + const client = connect({ host: "127.0.0.1", port: proxy.port, allowHalfOpen: true }) + const reply = new Promise((resolve, reject) => { + const chunks: Buffer[] = [] + client.on("data", (chunk: Buffer) => chunks.push(chunk)) + client.on("error", reject) + client.on("end", () => { + // Even after receiving the server's EOF, the request direction remains + // writable until the client explicitly finishes its remaining bytes. + if (first === "server") client.end(query) + resolve(Buffer.concat(chunks)) + }) + client.on("close", () => { + if (!client.readableEnded) reject(new Error("Proxy closed before response EOF")) + }) + }) + try { + if (first === "client") client.end(request) + else client.write(startup) + const [actualRequest, actualResponse] = await Promise.all([received.promise, reply]) + expect(actualRequest.equals(request)).toBe(true) + expect(actualResponse.equals(response)).toBe(true) + expect(proxy.snapshot()).toMatchObject({ clientBytes: request.length, serverBytes: response.length, frames: { Q: 1 } }) + } finally { + client.destroy() + await proxy.close() + for (const peer of peers) peer.destroy() + await new Promise((resolve, reject) => server.close((error) => error ? reject(error) : resolve())) + } +}) diff --git a/server/src/__tests__/performance/wire.ts b/server/src/__tests__/performance/wire.ts index d4014cf10..111e0b1fd 100644 --- a/server/src/__tests__/performance/wire.ts +++ b/server/src/__tests__/performance/wire.ts @@ -51,37 +51,52 @@ export async function createWireProxy(target: { host: string; port: number }, rt let counts: WireCounts const reset = () => { counts = { frames: {}, exchangeBoundaries: 0, clientBytes: 0, serverBytes: 0 } } reset() - const server = createServer((client) => { - const upstream = connect({ host: target.host, port: target.port }) + const server = createServer({ allowHalfOpen: true }, (client) => { + const upstream = connect({ host: target.host, port: target.port, allowHalfOpen: true }) for (const socket of [client, upstream]) { sockets.add(socket); socket.setNoDelay(true) } const frames = new FrontendFrames((kind) => { counts.frames[kind] = (counts.frames[kind] ?? 0) + 1 if (kind === "Q" || kind === "H" || kind === "S") counts.exchangeBoundaries++ }) - const forward = (source: Socket, destination: Socket, bytes: Buffer) => { - source.pause() - const write = () => { - if (destination.destroyed) return - destination.write(bytes, () => { if (!source.destroyed) source.resume() }) - } - if (!delayMs) write() - else { - const timer = setTimeout(() => { timers.delete(timer); write() }, delayMs / 2) - timers.add(timer) + const forward = (source: Socket, destination: Socket, observe: (bytes: Buffer) => void) => { + let pending = 0 + let ended = false + const finish = () => { + // Reading EOF must not end the destination before its delayed writes. + if (ended && !pending && !destination.destroyed) destination.end() } + source.on("data", (bytes: Buffer) => { + try { observe(bytes) } catch (error) { failures.push(error as Error); source.destroy(); return } + source.pause() + pending++ + const write = () => { + if (destination.destroyed) { pending--; return } + destination.write(bytes, () => { + pending-- + finish() + if (!source.destroyed) source.resume() + }) + } + if (!delayMs) write() + else { + const timer = setTimeout(() => { timers.delete(timer); write() }, delayMs / 2) + timers.add(timer) + } + }) + source.on("end", () => { ended = true; finish() }) } - client.on("data", (bytes: Buffer) => { - try { frames.accept(bytes) } catch (error) { failures.push(error as Error); client.destroy(); return } - counts.clientBytes += bytes.length - forward(client, upstream, bytes) - }) - upstream.on("data", (bytes: Buffer) => { counts.serverBytes += bytes.length; forward(upstream, client, bytes) }) + forward(client, upstream, (bytes) => { frames.accept(bytes); counts.clientBytes += bytes.length }) + forward(upstream, client, (bytes) => { counts.serverBytes += bytes.length }) client.on("error", (error) => { failures.push(error); upstream.destroy() }) upstream.on("error", (error) => { failures.push(error); client.destroy() }) - client.on("end", () => upstream.end()) - upstream.on("end", () => client.end()) - client.on("close", () => { sockets.delete(client); upstream.destroy() }) - upstream.on("close", () => { sockets.delete(upstream); client.destroy() }) + for (const [socket, peer] of [[client, upstream], [upstream, client]] as const) { + socket.on("close", () => { + sockets.delete(socket) + // A cleanly closed source can still have delayed bytes headed to its + // peer. Abort the pair only for an incomplete/failed connection. + if (!socket.readableEnded || !socket.writableFinished) peer.destroy() + }) + } }) await new Promise((resolve, reject) => { server.once("error", reject) diff --git a/server/src/db/connectionPolicy.integration.test.ts b/server/src/db/connectionPolicy.integration.test.ts index f7f6a844e..5e64171ff 100644 --- a/server/src/db/connectionPolicy.integration.test.ts +++ b/server/src/db/connectionPolicy.integration.test.ts @@ -57,9 +57,14 @@ suite("real transaction pooling safety (requires local PostgreSQL and PgBouncer await admin`alter role fixture_app set lock_timeout = '5s'` await admin`alter role fixture_app set idle_in_transaction_session_timeout = '15s'` await admin`create role fixture_unsafe login` + await admin`create role fixture_custom login` + await admin`alter role fixture_custom set statement_timeout = '30s'` + await admin`alter role fixture_custom set lock_timeout = '5s'` + await admin`alter role fixture_custom set idle_in_transaction_session_timeout = '15s'` + await admin`alter role fixture_custom set plan_cache_mode = force_custom_plan` await admin`create table public.pooling_fixture (id int primary key)` - await admin`grant select, insert on public.pooling_fixture to fixture_app` - await writeFile(join(directory, "users.txt"), '"fixture_app" ""\n"fixture_unsafe" ""\n') + await admin`grant select, insert on public.pooling_fixture to fixture_app, fixture_custom` + await writeFile(join(directory, "users.txt"), '"fixture_app" ""\n"fixture_unsafe" ""\n"fixture_custom" ""\n') await writeFile(join(directory, "pgbouncer.ini"), `[databases]\npostgres = host=127.0.0.1 port=${port} dbname=postgres\n[pgbouncer]\nlisten_addr = 127.0.0.1\nlisten_port = ${pooledPort}\nunix_socket_dir = ${directory}\nauth_type = trust\nauth_file = ${join(directory, "users.txt")}\npool_mode = transaction\ndefault_pool_size = 2\nmax_client_conn = 30\nmax_prepared_statements = 200\nlogfile = ${join(directory, "pgbouncer.log")}\npidfile = ${join(directory, "pgbouncer.pid")}\n`) bouncer = Bun.spawn([pgbouncer!, join(directory, "pgbouncer.ini")], { stdout: "ignore", stderr: "ignore" }) const deadline = Date.now() + 5_000 @@ -93,6 +98,46 @@ suite("real transaction pooling safety (requires local PostgreSQL and PgBouncer finally { await unsafe.close() } }) + test("all-custom rejects the query role even when the direct health role is qualified", async () => { + const unqualified = makeDatabaseClients(pooledUrl, { + DATABASE_CONNECTION_MODE: "pgbouncer", + DATABASE_DIRECT_URL: directUrl.replace("fixture_app", "fixture_custom"), + INLINE_DATABASE_PREPARED_QUERIES: "all-custom", + }) + try { await expect(unqualified.validateStartup()).rejects.toThrow("force_custom_plan") } + finally { await unqualified.close() } + }) + + test("all-custom validates role defaults through PgBouncer and keeps custom plans", async () => { + const qualified = makeDatabaseClients(pooledUrl.replace("fixture_app", "fixture_custom"), { + DATABASE_CONNECTION_MODE: "pgbouncer", + DATABASE_DIRECT_URL: directUrl, + INLINE_DATABASE_PREPARED_QUERIES: "all-custom", + }) + try { + await qualified.validateStartup() + const db = drizzle(qualified.queryClient) + for (let value = 0; value < 24; value++) { + const rows = await db.execute(sql`select ${value}::int as qualified_value`) + expect(rows[0]?.["qualified_value"]).toBe(value) + } + const plans = await qualified.queryClient`select generic_plans, custom_plans from pg_prepared_statements + where statement like '%as qualified_value%' and statement not like '%pg_prepared_statements%'` + expect(plans.length).toBeGreaterThan(0) + expect(plans.every((row) => Number(row["generic_plans"]) === 0)).toBe(true) + expect(plans.some((row) => Number(row["custom_plans"]) > 5)).toBe(true) + await qualified.checkHealth() + // Deliberately change the isolated fixture backend to exercise readiness + // losing its qualification; application queries return to narrow mode. + await qualified.queryClient`set plan_cache_mode = auto` + try { + await expect(qualified.checkHealth()).rejects.toThrow("force_custom_plan") + } finally { + await qualified.queryClient`set plan_cache_mode = force_custom_plan` + } + } finally { await qualified.close() } + }) + test("prepared queries reuse two server connections without losing role defaults", async () => { const results = await Promise.all(Array.from({ length: 40 }, (_, value) => clients!.queryClient.begin(async (tx) => { const rows = await tx`select ${value}::int as value, pg_backend_pid() as pid, current_setting('statement_timeout') as timeout` diff --git a/server/src/db/connectionPolicy.test.ts b/server/src/db/connectionPolicy.test.ts index 397cc6daa..9caac8d9f 100644 --- a/server/src/db/connectionPolicy.test.ts +++ b/server/src/db/connectionPolicy.test.ts @@ -74,6 +74,11 @@ describe("database connection safety", () => { } }) + test("invalid prepared modes fail without echoing their value", () => { + expect(() => makeDatabaseClients(direct, { INLINE_DATABASE_PREPARED_QUERIES: "private-secret" })) + .toThrow("INLINE_DATABASE_PREPARED_QUERIES must be access or all-custom.") + }) + test("pooled mode and a separate direct URL must be explicit", () => { expect(() => databaseConnectionPolicy(pooled)).toThrow("explicit PgBouncer mode") expect(() => databaseConnectionPolicy(pooled, { DATABASE_CONNECTION_MODE: "pgbouncer" })).toThrow("DATABASE_DIRECT_URL") diff --git a/server/src/db/connectionPolicy.ts b/server/src/db/connectionPolicy.ts index b2e920b2c..950e437c8 100644 --- a/server/src/db/connectionPolicy.ts +++ b/server/src/db/connectionPolicy.ts @@ -1,5 +1,5 @@ import postgres from "postgres" -import { prepareAccessQueries } from "./preparedAccessQueries" +import { prepareQueries, type PreparedQueryPolicy } from "./preparedQueries" export const QUERY_TIMEOUTS = { statement_timeout: 30_000, @@ -19,6 +19,7 @@ type DatabaseEnvironment = { DATABASE_DIRECT_URL?: string DATABASE_CONNECTION_MODE?: string INLINE_DATABASE_QUERY_POOL_MAX?: string + INLINE_DATABASE_PREPARED_QUERIES?: string } const parseDatabaseUrl = (value: string): URL => { @@ -143,6 +144,11 @@ export const makeDatabaseClients = (databaseUrl: string, environment: DatabaseEn if (!/^(?:[1-9]|10)$/.test(queryPoolMax) || queryPoolMax.trim() !== queryPoolMax) { throw new Error("INLINE_DATABASE_QUERY_POOL_MAX must be an integer from 1 through 10.") } + const preparedMode = environment.INLINE_DATABASE_PREPARED_QUERIES ?? "access" + if (preparedMode !== "access" && preparedMode !== "all-custom") { + throw new Error("INLINE_DATABASE_PREPARED_QUERIES must be access or all-custom.") + } + const preparation: PreparedQueryPolicy = { mode: "access" } const policy = databaseConnectionPolicy(databaseUrl, environment) const queryClient = postgres(policy.queryUrl, { max: Number(queryPoolMax), connect_timeout: 5, idle_timeout: 30, @@ -155,29 +161,36 @@ export const makeDatabaseClients = (databaseUrl: string, environment: DatabaseEn max: 1, connect_timeout: 2, idle_timeout: 30, connection: { application_name: "inline-health", ...HEALTH_TIMEOUTS }, }) - prepareAccessQueries(queryClient) + prepareQueries(queryClient, preparation) const close = () => Promise.all([queryClient.end({ timeout: 5 }), healthClient.end({ timeout: 5 })]).then(() => {}) const queryTimeoutSettings = () => queryClient` select name, setting, unit from pg_settings - where name in ('statement_timeout', 'lock_timeout', 'idle_in_transaction_session_timeout') + where name in ('statement_timeout', 'lock_timeout', 'idle_in_transaction_session_timeout', 'plan_cache_mode') `.execute() + const assertPreparation = (rows: readonly TimeoutRow[]) => { + if (preparedMode === "all-custom" && rows.find((row) => row.name === "plan_cache_mode")?.setting !== "force_custom_plan") { + preparation.mode = "access" + throw new Error("all-custom preparation requires force_custom_plan on the query role.") + } + } const checkHealth = () => { const direct = healthClient.unsafe( "SELECT (EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::double precision AS database_time_millis", ).execute() - if (policy.mode === "direct") return direct + if (policy.mode === "direct" && preparedMode !== "all-custom") return direct // A healthy direct endpoint must not hide a broken/queued pooled endpoint. // The health controller's existing deadline cancels both pending queries. const pooled = queryTimeoutSettings() const probes = combineHealthProbes(direct, [pooled]) return Object.assign(probes.then(async (clock) => { const rows = await pooled - assertQueryTimeouts(rows) + if (policy.mode === "pgbouncer") assertQueryTimeouts(rows) + assertPreparation(rows) return clock }), { cancel: probes.cancel }) } const validateStartup = async (): Promise => { - if (policy.mode === "direct") return + if (policy.mode === "direct" && preparedMode === "access") return const query = queryTimeoutSettings() let timer: ReturnType | undefined try { @@ -185,23 +198,28 @@ export const makeDatabaseClients = (databaseUrl: string, environment: DatabaseEn query, new Promise((_, reject) => { timer = setTimeout(() => { - reject(new Error("PgBouncer timeout qualification exceeded 5 seconds.")) + reject(new Error("Database role qualification exceeded 5 seconds.")) try { query.cancel() } catch { /* The startup deadline remains authoritative. */ } }, 5_000) }), ]) - assertQueryTimeouts(rows) - await healthClient`select 1` + if (policy.mode === "pgbouncer") assertQueryTimeouts(rows) + // The query path can use a different role from the direct health client. + // Operators must also recycle existing backends after setting role defaults. + assertPreparation(rows) + if (policy.mode === "pgbouncer") await healthClient`select 1` + preparation.mode = preparedMode } catch { await close() // Driver errors can include connection metadata; emit a stable safe error. - throw new Error("PgBouncer startup qualification failed; verify role timeout defaults and the direct health endpoint.") + throw new Error(`${policy.mode === "pgbouncer" ? "PgBouncer" : "Database"} startup qualification failed; verify role timeout defaults, the direct health endpoint, and force_custom_plan when all-custom preparation is enabled.`) } finally { clearTimeout(timer) } } const diagnosticsConfiguration = { mode: policy.mode, + preparedQueries: preparedMode, queryPoolMax: queryClient.options.max, healthPoolMax: healthClient.options.max, statementTimeoutMs: QUERY_TIMEOUTS.statement_timeout, diff --git a/server/src/db/models/users.ts b/server/src/db/models/users.ts index 050c64556..0efc0ec84 100644 --- a/server/src/db/models/users.ts +++ b/server/src/db/models/users.ts @@ -22,10 +22,12 @@ export class UsersModel { return [] } + // One binding keeps the SQL shape stable across recipient counts. bigint + // also lets unknown protocol IDs beyond int4 simply miss the indexed key. const rows = await db .select({ id: users.id }) .from(users) - .where(and(inArray(users.id, uniqueUserIds), userNotDeleted())) + .where(and(sql`${users.id} = any(${sql.param(uniqueUserIds)}::bigint[])`, userNotDeleted())) const activeUserIds = new Set(rows.map((row) => row.id)) return uniqueUserIds.filter((id) => activeUserIds.has(id)) diff --git a/server/src/db/preparedAccessQueries.ts b/server/src/db/preparedAccessQueries.ts deleted file mode 100644 index 5431123a9..000000000 --- a/server/src/db/preparedAccessQueries.ts +++ /dev/null @@ -1,22 +0,0 @@ -import type postgres from "postgres" -import { accessQueryComment } from "./databaseDiagnostics" - -const comments = [accessQueryComment("sync.replay"), accessQueryComment("sync.sidecars")] - -/** - * Drizzle's postgres-js adapter uses unsafe(), which disables preparation even - * when the client enables it. Opt in only user replay's access queries on the - * shared client: at most 100 chat IDs repeated in three clauses plus one user. - * The two sources therefore produce at most 200 statement shapes per client. - * Larger batches and transaction-owned clients retain their existing behavior. - * This caches plans, never authorization results. PgBouncer must support - * protocol-level prepared statements (production max_prepared_statements=200). - */ -export function prepareAccessQueries(client: postgres.Sql): void { - const unsafe = client.unsafe - client.unsafe = (query, parameters, options) => unsafe(query, parameters, - options?.prepare === undefined && parameters !== undefined && parameters.length <= 301 && - comments.some((comment) => query.trimEnd().endsWith(comment)) - ? { ...options, prepare: true } - : options) -} diff --git a/server/src/db/preparedQueries.postgres.test.ts b/server/src/db/preparedQueries.postgres.test.ts new file mode 100644 index 000000000..49982ff8e --- /dev/null +++ b/server/src/db/preparedQueries.postgres.test.ts @@ -0,0 +1,211 @@ +import { expect, test } from "bun:test" +import postgres from "postgres" +import { drizzle } from "drizzle-orm/postgres-js" +import { sql } from "drizzle-orm" +import { accessQueryComment } from "./databaseDiagnostics" +import { setupTestLifecycle } from "../__tests__/database" +import { createWireProxy } from "../__tests__/performance/wire" +import { MAX_BROAD_PREPARED_QUERY_SHAPES, MAX_PREPARED_QUERY_SHAPES, prepareQueries, type PreparedQueryPolicy } from "./preparedQueries" + +setupTestLifecycle() + +async function withClient(run: (client: postgres.Sql, proxy: Awaited>) => Promise, mode: PreparedQueryPolicy["mode"] = "all-custom") { + const url = new URL(process.env["DATABASE_URL"]!) + const proxy = await createWireProxy({ host: url.hostname, port: Number(url.port || 5432) }, 0) + url.hostname = "127.0.0.1" + url.port = String(proxy.port) + const client = postgres(url.toString(), { max: 1, ssl: false, connection: { plan_cache_mode: "force_custom_plan" } }) + prepareQueries(client, { mode }) + prepareQueries(client) // Installing twice must not stack wrappers or reset the shape budget. + try { + await client`select 1` + await run(client, proxy) + } finally { + await client.end() + await proxy.close() + } +} + +test("Drizzle reuses plans with new values in the pool, transactions and savepoints", async () => { + await withClient(async (client, proxy) => { + const db = drizzle(client) + const read = async (value: number) => { + const rows = await db.execute(sql`select ${value}::int as value`) + expect(rows[0]?.["value"]).toBe(value) + } + await read(0) + await read(1) + proxy.reset() + await read(2) + expect(proxy.snapshot().frames["H"] ?? 0).toBe(0) + expect(proxy.snapshot().frames["P"] ?? 0).toBe(0) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + + const nestedRead = (value: number) => db.transaction(async (tx) => { + const rows = await tx.execute(sql`select ${value}::int + 10 as value`) + expect(rows[0]?.["value"]).toBe(value + 10) + await tx.transaction(async (nested) => { + const inner = await nested.execute(sql`select ${value}::int + 20 as value`) + expect(inner[0]?.["value"]).toBe(value + 20) + }) + }, { isolationLevel: "repeatable read", accessMode: "read only" }) + await nestedRead(2) + await nestedRead(3) + proxy.reset() + await nestedRead(4) + expect(proxy.snapshot().frames["H"] ?? 0).toBe(0) + expect(proxy.snapshot().frames["P"] ?? 0).toBe(0) + + const reserved = await client.reserve() + try { + await reserved.unsafe("select $1::int + 30 as value", [4]) + await reserved.unsafe("select $1::int + 30 as value", [5]) + proxy.reset() + const rows = await reserved.unsafe("select $1::int + 30 as value", [6]) + expect(rows[0]?.["value"]).toBe(36) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + } finally { reserved.release() } + }) +}) + +test("prepared writes preserve savepoint and outer transaction rollback", async () => { + await withClient(async (client) => { + const db = drizzle(client) + await db.execute(sql`create temporary table prepared_write_test (id integer primary key)`) + await db.transaction(async (tx) => { + await tx.execute(sql`insert into prepared_write_test values (${1})`) + await expect(tx.transaction(async (nested) => { + await nested.execute(sql`insert into prepared_write_test values (${2})`) + throw new Error("rollback inner") + })).rejects.toThrow("rollback inner") + await tx.execute(sql`insert into prepared_write_test values (${3})`) + }) + await expect(db.transaction(async (tx) => { + await tx.execute(sql`insert into prepared_write_test values (${4})`) + throw new Error("rollback outer") + })).rejects.toThrow("rollback outer") + const rows = await db.execute(sql`select id from prepared_write_test order by id`) + expect(rows.map((row) => row["id"])).toEqual([1, 3]) + }) +}) + +test("explicit opt-out and shape-budget overflow still execute correctly without retaining new plans", async () => { + await withClient(async (client, proxy) => { + const query = "select $1::int as value" + await client.unsafe(query, [0]) + await client.unsafe(query, [1]) + proxy.reset() + expect((await client.unsafe(query, [2], { prepare: false }))[0]?.["value"]).toBe(2) + expect(proxy.snapshot().frames["H"]).toBe(1) + for (let i = 1; i < MAX_BROAD_PREPARED_QUERY_SHAPES; i++) { + await client.unsafe(`select $1::int + ${i} as value`, [0]) + await client.unsafe(`select $1::int + ${i} as value`, [1]) + } + const overflow = `select $1::int + ${MAX_PREPARED_QUERY_SHAPES} as value` + await client.unsafe(overflow, [1]) + proxy.reset() + expect((await client.unsafe(overflow, [2]))[0]?.["value"]).toBe(MAX_PREPARED_QUERY_SHAPES + 2) + expect(proxy.snapshot().frames["H"]).toBe(1) + const access = `select $1::int as access_value ${accessQueryComment("sync.replay")}` + await client.unsafe(access, [1]) + proxy.reset() + expect((await client.unsafe(access, [2]))[0]?.["access_value"]).toBe(2) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + proxy.reset() + expect((await client.unsafe(query, [3]))[0]?.["value"]).toBe(3) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + }) +}) + + +test("nullable scalar type changes and custom typed parameters keep their original values", async () => { + await withClient(async (client, proxy) => { + for (const value of [true, null, false, null]) { + expect((await client.unsafe("select $1::boolean as value", [value]))[0]?.["value"]).toBe(value) + } + const typed = () => client.unsafe("select $1::int as typed", [client.typed(7, 23)]) + await typed() + proxy.reset() + expect((await typed())[0]?.["typed"]).toBe(7) + expect(proxy.snapshot().frames["H"]).toBe(1) + }) +}) + +test("ordinary array bindings reuse plans with changed lengths and values", async () => { + await withClient(async (client, proxy) => { + // JSON readback avoids the driver's numeric-array NULL parser so this + // checks the bound database values, including NULL and an empty array. + const query = "select to_json($1::integer[]) as ids" + await client.unsafe(query, [[1]]) + await client.unsafe(query, [[1, 2]]) + for (const ids of [[3], [4, 5, 6], [], [null, 7]]) { + proxy.reset() + expect((await client.unsafe(query, [ids]))[0]?.["ids"]).toEqual(ids) + expect(proxy.snapshot().frames["P"] ?? 0).toBe(0) + expect(proxy.snapshot().frames["H"] ?? 0).toBe(0) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + } + }) +}) + + +test("default policy leaves unrelated partial-index predicates unnamed after repeated reads", async () => { + await withClient(async (client, proxy) => { + await client`set plan_cache_mode = auto` + await client`create temporary table prepared_partial_test (id int, stopped boolean)` + await client`insert into prepared_partial_test select id, id > 10 from generate_series(1, 10000) id` + await client`create index prepared_partial_due on prepared_partial_test (id) where stopped = false` + await client`analyze prepared_partial_test` + const query = "select id from prepared_partial_test where stopped = $1 order by id" + for (let run = 0; run < 12; run++) expect(await client.unsafe(query, [false])).toHaveLength(10) + proxy.reset() + expect(await client.unsafe(query, [false])).toHaveLength(10) + expect(proxy.snapshot().frames["H"]).toBe(1) + const retained = await client`select name from pg_prepared_statements where statement = ${query}` + expect(retained).toHaveLength(0) + const plan = await client.unsafe(`explain (format json) ${query}`, [false]) + expect(JSON.stringify(plan)).toContain("prepared_partial_due") + const access = `select $1::int as value ${accessQueryComment("sync.replay")}` + await client.unsafe(access, [1]) + proxy.reset() + expect((await client.unsafe(access, [2]))[0]?.["value"]).toBe(2) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + }, "access") +}) + +test("broad custom plans retain partial indexes and skip Describe after more than five executions", async () => { + await withClient(async (client, proxy) => { + await client`create temporary table prepared_custom_test (id int, stopped boolean)` + await client`insert into prepared_custom_test select id, id > 10 from generate_series(1, 10000) id` + await client`create index prepared_custom_due on prepared_custom_test (id) where stopped = false` + await client`analyze prepared_custom_test` + const query = "select id from prepared_custom_test where stopped = $1 order by id" + for (let run = 0; run < 12; run++) expect(await client.unsafe(query, [false])).toHaveLength(10) + proxy.reset() + expect(await client.unsafe(query, [false])).toHaveLength(10) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + const plans = await client`select name, generic_plans, custom_plans from pg_prepared_statements where statement = ${query}` + expect(plans).toHaveLength(1) + expect(Number(plans[0]?.["generic_plans"])).toBe(0) + expect(Number(plans[0]?.["custom_plans"])).toBeGreaterThan(5) + const name = String(plans[0]?.["name"]).replaceAll('"', '""') + const plan = await client.unsafe(`explain (format json) execute "${name}"(false)`) + expect(JSON.stringify(plan)).toContain("prepared_custom_due") + }) +}) + +test("one-off shapes do not exhaust broad admission and access keeps reserved capacity", async () => { + await withClient(async (client, proxy) => { + for (let i = 0; i < MAX_PREPARED_QUERY_SHAPES + 20; i++) { + await client.unsafe(`select $1::int + ${i} as one_off`, [1]) + } + const hot = "select $1::int as hot" + await client.unsafe(hot, [1]) + await client.unsafe(hot, [2]) + proxy.reset() + expect((await client.unsafe(hot, [3]))[0]?.["hot"]).toBe(3) + expect(proxy.snapshot().exchangeBoundaries).toBe(1) + const plans = await client`select count(*)::int as count from pg_prepared_statements where statement like '%as one_off'` + expect(plans[0]?.["count"]).toBe(0) + }) +}) diff --git a/server/src/db/preparedQueries.ts b/server/src/db/preparedQueries.ts new file mode 100644 index 000000000..ceade51f9 --- /dev/null +++ b/server/src/db/preparedQueries.ts @@ -0,0 +1,97 @@ +import type postgres from "postgres" +import { accessQueryComment } from "./databaseDiagnostics" + +// postgres.js retains admitted names until the physical connection closes. Do not +// pretend that evicting an application entry would deallocate those statements. +export const MAX_PREPARED_QUERY_SHAPES = 200 +export const MAX_BROAD_PREPARED_QUERY_SHAPES = 100 +export type PreparedQueryPolicy = { mode: "access" | "all-custom" } +const accessComments = [accessQueryComment("sync.replay"), accessQueryComment("sync.sidecars")] +const MAX_QUERY_LENGTH = 32_768 +const MAX_PARAMETERS = 301 +const installed = new WeakSet() +type TransactionCallback = (sql: postgres.TransactionSql) => T | Promise + +/** Drizzle uses unsafe(), whose call-level default disables preparation even + * when the pool enables it. Apply the same bounded policy to borrowed/transaction + * clients, including savepoints. Cache plans only, never rows or authority. + * Broad preparation requires force_custom_plan on every backend serving the + * application role. The connection policy activates it only after startup checks. + * Requires protocol-level prepared statement support when using PgBouncer. */ +export function prepareQueries(client: postgres.Sql, policy: PreparedQueryPolicy = { mode: "access" }): void { + const shapes = new Set() + const broadShapes = new Set() + // Only candidates are evicted: they have never created retained statements. + const candidates = new Set() + const wrap = (scope: T): T => { + if (installed.has(scope)) return scope + installed.add(scope) + const unsafe = scope.unsafe + scope.unsafe = (query, parameters, options) => { + if (options?.prepare !== undefined || !parameters?.length || parameters.length > MAX_PARAMETERS || + query.length > MAX_QUERY_LENGTH || !/^\s*(select|insert|update|delete|with)\b/i.test(query)) { + return unsafe(query, parameters, options) + } + const access = accessComments.some((comment) => query.trimEnd().endsWith(comment)) + if (!access && policy.mode !== "all-custom") return unsafe(query, parameters, options) + // Match the driver's ordinary scalar/array type categories. Custom typed + // parameters retain their original execution path rather than admitting + // unbounded type variants under one SQL-text entry. + const types = parameters.map(parameterType) + if (types.some((type) => type === undefined)) return unsafe(query, parameters, options) + const signature = `${types.join(",")}\0${query}` + if (shapes.has(signature)) return unsafe(query, parameters, { ...options, prepare: true }) + // In broad mode leave half the budget for replay access, even when many + // variable-length IN queries arrive first. Default access mode has all 200. + const capacity = shapes.size < MAX_PREPARED_QUERY_SHAPES && + (access || broadShapes.size < MAX_BROAD_PREPARED_QUERY_SHAPES) + if (!capacity) return unsafe(query, parameters, options) + if (!access && !candidates.has(signature)) { + candidates.add(signature) + if (candidates.size > MAX_PREPARED_QUERY_SHAPES) candidates.delete(candidates.values().next().value!) + return unsafe(query, parameters, options) + } + candidates.delete(signature) + shapes.add(signature) + if (!access) broadShapes.add(signature) + return unsafe(query, parameters, { ...options, prepare: true }) + } + if (typeof scope.begin === "function") { + const begin = scope.begin + scope.begin = (first: string | TransactionCallback, callback?: TransactionCallback) => { + if (typeof first === "function") return begin((tx) => first(wrap(tx))) + if (!callback) throw new Error("Missing transaction callback") + return begin(first, (tx) => callback(wrap(tx))) + } + } + if ("savepoint" in scope && typeof scope.savepoint === "function") { + const savepoint = scope.savepoint as postgres.TransactionSql["savepoint"] + scope.savepoint = (first: string | TransactionCallback, callback?: TransactionCallback) => { + if (typeof first === "function") return savepoint((tx) => first(wrap(tx))) + if (!callback) throw new Error("Missing savepoint callback") + return savepoint(first, (tx) => callback(wrap(tx))) + } + } + if (typeof scope.reserve === "function") { + const reserve = scope.reserve + scope.reserve = async () => wrap(await reserve()) + } + return scope + } + wrap(client) +} + +function parameterType(value: unknown): number | undefined { + // postgres.js infers ordinary arrays from their first element. Bound the + // traversal so malformed/cyclic inputs do not recurse inside this policy. + for (let depth = 0; Array.isArray(value); depth++) { + if (depth >= 32) return undefined + value = value[0] ?? null + } + if (value instanceof Date) return 1184 + if (value instanceof Uint8Array) return 17 + if (typeof value === "boolean") return 16 + if (typeof value === "bigint") return 20 + if (value === null || typeof value === "string" || typeof value === "number") return 0 + return undefined +} diff --git a/server/src/modules/dialogOpen.postgres.test.ts b/server/src/modules/dialogOpen.postgres.test.ts new file mode 100644 index 000000000..ec8a7da29 --- /dev/null +++ b/server/src/modules/dialogOpen.postgres.test.ts @@ -0,0 +1,80 @@ +import { expect, test } from "bun:test" +import { asc, eq } from "drizzle-orm" +import { db } from "@in/server/db" +import { dialogs, users } from "@in/server/db/schema" +import { UsersModel } from "@in/server/db/models/users" +import { setupTestLifecycle, testUtils } from "../__tests__/setup" +import { measureOperation } from "../__tests__/performance/measure" +import { setDialogOpenForUsers } from "./dialogOpen" + +setupTestLifecycle() + +test("one hundred unchanged dialogs use a constant query budget and preserve their state", async () => { + const actors = await db.insert(users).values(Array.from({ length: 100 }, (_, i) => ({ email: `batch-${i}@example.test` }))).returning() + const chat = await testUtils.createChat(null) + if (!chat) throw new Error("Missing chat") + await db.insert(dialogs).values(actors.map((user) => ({ + userId: user.id, chatId: chat.id, open: true, order: "a0", archived: false, + }))) + const before = await db.select().from(dialogs).where(eq(dialogs.chatId, chat.id)).orderBy(asc(dialogs.userId)) + const sample = await measureOperation(db.$client.options, async () => { + const result = await setDialogOpenForUsers({ chat, userIds: actors.map((user) => user.id), open: true }) + expect(result.changedDialogs).toEqual([]) + expect(result.dialogs.toSorted((a, b) => a.userId - b.userId)).toEqual(before) + }, async () => {}) + // Active owners, BEGIN, one sorted owner-lock batch, existing dialogs, COMMIT. + expect(sample.sql.commands).toBe(5) + const smallBatch = await measureOperation(db.$client.options, async () => { + const result = await setDialogOpenForUsers({ chat, userIds: [actors[0]!.id], open: true }) + expect(result.changedDialogs).toEqual([]) + expect(result.dialogs).toHaveLength(1) + }, async () => {}) + expect(smallBatch.sql.shapes).toEqual(sample.sql.shapes) + expect(await db.select().from(dialogs).where(eq(dialogs.chatId, chat.id)).orderBy(asc(dialogs.userId))).toEqual(before) +}) + +test("array-bound active owners preserve filtering and caller order", async () => { + const first = await testUtils.createUser("array-first@example.test") + const second = await testUtils.createUser("array-second@example.test") + const removed = await testUtils.createUser("array-removed@example.test") + await db.update(users).set({ deleted: true }).where(eq(users.id, removed.id)) + expect(await UsersModel.getActiveUserIds([ + second.id, removed.id, first.id, second.id, 2 ** 40, 0, -1, Number.NaN, + ])).toEqual([second.id, first.id]) + expect(await UsersModel.getActiveUserIds([])).toEqual([]) +}) + +test("overlapping reversed owner batches keep independent chat allocations ordered without deadlock", async () => { + const actors = await db.insert(users).values(Array.from({ length: 10 }, (_, i) => ({ email: `overlap-${i}@example.test` }))).returning() + const first = await testUtils.createChat(null) + const second = await testUtils.createChat(null) + if (!first || !second) throw new Error("Missing chats") + const ids = actors.map((user) => user.id) + const results = await Promise.allSettled([ + setDialogOpenForUsers({ chat: first, userIds: [...ids, ids[0]!], open: true }), + setDialogOpenForUsers({ chat: second, userIds: ids.toReversed(), open: true }), + ]) + for (const result of results) { + if (result.status === "rejected") throw result.reason + expect(result.value.dialogs).toHaveLength(10) + expect(result.value.changedDialogs).toHaveLength(10) + } + for (const user of actors) { + const rows = await db.select().from(dialogs).where(eq(dialogs.userId, user.id)) + expect(rows).toHaveLength(2) + expect(rows.every((row) => row.open === true && row.order !== null)).toBe(true) + expect(new Set(rows.map((row) => row.order)).size).toBe(2) + } +}) + +test("closing metadata-only state returns persisted cleared fields even without a visibility change", async () => { + const actor = await testUtils.createUser() + const chat = await testUtils.createChat(null) + if (!chat) throw new Error("Missing chat") + await db.insert(dialogs).values({ userId: actor.id, chatId: chat.id, open: false, order: "a0" }) + const result = await setDialogOpenForUsers({ chat, userIds: [actor.id], open: false }) + expect(result.changedDialogs).toEqual([]) + expect(result.dialogs[0]?.order).toBeNull() + const repeated = await setDialogOpenForUsers({ chat, userIds: [actor.id], open: false }) + expect(repeated.dialogs).toEqual(result.dialogs) +}) diff --git a/server/src/modules/dialogOpen.ts b/server/src/modules/dialogOpen.ts index 0bbc5804d..ccc99e6dd 100644 --- a/server/src/modules/dialogOpen.ts +++ b/server/src/modules/dialogOpen.ts @@ -2,7 +2,7 @@ import { db } from "@in/server/db" import { UsersModel } from "@in/server/db/models/users" import type { Transaction } from "@in/server/db/types" import { chats, dialogFolders, dialogs, users, type DbChat, type DbDialog, type DbNewDialog } from "@in/server/db/schema" -import { and, asc, desc, eq, inArray, isNotNull, isNull, or } from "drizzle-orm" +import { and, asc, desc, eq, isNotNull, isNull, or, sql } from "drizzle-orm" import { FractionalIndex } from "@in/server/modules/fractionalIndex" export type DialogOpenPlacement = "top" | "bottom" @@ -261,15 +261,17 @@ async function setDialogOpenForUsersInTransaction( // Acquire all user owners before touching any dialog rows. Projection // writers use the same users -> dialogs order, so overlapping batches do - // not deadlock or derive a stale fractional order. - for (const userId of userIds) { - await tx.select({ id: users.id }).from(users).where(eq(users.id, userId)).for("no key update").limit(1) - } + // not deadlock or derive a stale fractional order. Bind each ID list as one + // array so varying recipient counts do not consume new prepared SQL shapes. + await tx.select({ id: users.id }).from(users) + .where(sql`${users.id} = any(${sql.param(userIds)}::bigint[])`) + .orderBy(asc(users.id)) + .for("no key update") const existingDialogs = await tx .select() .from(dialogs) - .where(and(eq(dialogs.chatId, input.chat.id), inArray(dialogs.userId, userIds))) + .where(and(eq(dialogs.chatId, input.chat.id), sql`${dialogs.userId} = any(${sql.param(userIds)}::bigint[])`)) const existingUserIds = new Set(existingDialogs.map((dialog) => dialog.userId)) const changedUserIds = new Set() @@ -289,6 +291,9 @@ async function setDialogOpenForUsersInTransaction( return showInChatList && dialog.chatListHidden === true }) const missingUserIds = userIds.filter((userId) => !existingUserIds.has(userId)) + if (dialogsToOpen.length === 0 && missingUserIds.length === 0) { + return { dialogs: existingDialogs, changedDialogs: [] } + } for (const dialog of dialogsToOpen) { const placement = input.openPlacementByUserId?.get(dialog.userId) ?? defaultDialogOpenPlacement @@ -346,6 +351,9 @@ async function setDialogOpenForUsersInTransaction( .filter((dialog) => effectiveDialogOpenForDialog(dialog) || dialog.folderId != null) .map((dialog) => dialog.userId) const missingUserIds = userIds.filter((userId) => !existingUserIds.has(userId)) + if (persistCloseUserIds.length === 0 && missingUserIds.length === 0) { + return { dialogs: existingDialogs, changedDialogs: [] } + } if (persistCloseUserIds.length > 0) { await tx @@ -356,7 +364,7 @@ async function setDialogOpenForUsersInTransaction( order: null, folderId: null, }) - .where(and(eq(dialogs.chatId, input.chat.id), inArray(dialogs.userId, persistCloseUserIds))) + .where(and(eq(dialogs.chatId, input.chat.id), sql`${dialogs.userId} = any(${sql.param(persistCloseUserIds)}::bigint[])`)) changedCloseUserIds.forEach((userId) => changedUserIds.add(userId)) } @@ -381,7 +389,7 @@ async function setDialogOpenForUsersInTransaction( const finalDialogs = await tx .select() .from(dialogs) - .where(and(eq(dialogs.chatId, input.chat.id), inArray(dialogs.userId, userIds))) + .where(and(eq(dialogs.chatId, input.chat.id), sql`${dialogs.userId} = any(${sql.param(userIds)}::bigint[])`)) return { dialogs: finalDialogs, diff --git a/server/src/modules/updates/replayLoopGuard.test.ts b/server/src/modules/updates/replayLoopGuard.test.ts index 3b6836260..3b111c8e0 100644 --- a/server/src/modules/updates/replayLoopGuard.test.ts +++ b/server/src/modules/updates/replayLoopGuard.test.ts @@ -58,4 +58,23 @@ describe("replay loop guard", () => { expect(guard.allow("a", 4000)).toBe(true) expect(guard.allow("a", 0)).toBe(true) }) + + test("resets rewind backoff after passing the previous high-water mark on the same socket", () => { + let now = 0 + const guard = new ReplayLoopGuard(() => now) + for (const seq of [4000, 0, 4000]) expect(guard.allow("a", seq)).toBe(true) + expect(guard.allow("a", 0)).toBe(false) + + // Recovery is possible during backoff, without an idle expiry or reconnect. + expect(guard.allow("a", 4001)).toBe(true) + now += 1000 + expect(guard.allow("a", 0)).toBe(true) + expect(guard.allow("a", 4001)).toBe(true) + expect(guard.allow("a", 0)).toBe(false) + // A new failure episode starts at 60 seconds, not the old 120-second step. + now += 59_999 + expect(guard.allow("a", 0)).toBe(false) + now += 1 + expect(guard.allow("a", 0)).toBe(true) + }) }) diff --git a/server/src/modules/updates/replayLoopGuard.ts b/server/src/modules/updates/replayLoopGuard.ts index cf1c22ff8..0ec831bdc 100644 --- a/server/src/modules/updates/replayLoopGuard.ts +++ b/server/src/modules/updates/replayLoopGuard.ts @@ -25,6 +25,7 @@ export class ReplayLoopGuard { // A recovered client may continue beyond the old failed pass immediately. if (startSeq > state.highWater) { state.highWater = startSeq + state.rewinds = 0 state.blockedUntil = 0 } if (now < state.blockedUntil) return false