Jyotisha/frontend/tests/bench-local-pg-roundtrips.ts

/**
 * Count local-Postgres adapter roundtrips for a fictional homepage plus consult write.
 * Run from frontend/: npx tsx tests/bench-local-pg-roundtrips.ts
 * Prints JSON. Does not print connection strings.
 */
import { spawnSync } from "node:child_process";

import type { Pool, PoolClient } from "pg";

import {
  closeLocalPostgresDataPools,
  clearLocalPostgresMetadataCacheForTests,
  createLocalPostgresDataClient,
  localPostgresPoolForTests,
} from "../src/lib/db/local-postgres-client-core.ts";
import { startPostgresFixture } from "./helpers/postgres-fixture.ts";

const userId = "018f4e6d-7a11-7000-8000-0000000000bb";
const sessionId = "018f4e6d-7a11-7000-8000-0000000000cc";
const chartId = "018f4e6d-7a11-7000-8000-0000000000dd";

const schemaSql = `
create table if not exists public.profiles (
  id uuid primary key,
  display_name text
);
create table if not exists public.chart_profiles (
  id uuid primary key,
  user_id uuid not null,
  label text,
  updated_at timestamptz not null default now()
);
create table if not exists public.chat_sessions (
  id uuid primary key,
  user_id uuid not null,
  title text,
  updated_at timestamptz not null default now()
);
grant usage on schema public to authenticated, service_role;
grant select, insert, update, delete on public.profiles, public.chart_profiles, public.chat_sessions to authenticated, service_role;
insert into public.profiles (id, display_name) values ('${userId}', 'Fictional Person')
  on conflict (id) do nothing;
insert into public.chart_profiles (id, user_id, label) values ('${chartId}', '${userId}', 'home')
  on conflict (id) do nothing;
insert into public.chat_sessions (id, user_id, title) values ('${sessionId}', '${userId}', 'consult')
  on conflict (id) do nothing;
create or replace function public.session_probe()
returns table(db_user text)
language sql stable
as $$ select current_user::text $$;
grant execute on function public.session_probe() to service_role, authenticated;
`;

type Mode = { cache: "on" | "off"; merged: "on" | "off" };

function percentile(values: number[], ratio: number): number {
  const sorted = [...values].sort((left, right) => left - right);
  const index = Math.min(sorted.length - 1, Math.max(0, Math.ceil(ratio * sorted.length) - 1));
  return sorted[index] ?? 0;
}

function median(values: number[]): number {
  const sorted = [...values].sort((left, right) => left - right);
  return sorted[Math.floor(sorted.length / 2)] ?? 0;
}

const benchCounter = { roundtrips: 0 };
const queryWrapped = new WeakSet<object>();

async function withCountingPools(
  urls: readonly string[],
  run: () => Promise<unknown>,
): Promise<{ roundtrips: number; elapsedMs: number; payload: string }> {
  const restores: Array<() => void> = [];
  benchCounter.roundtrips = 0;
  for (const url of urls) {
    const pool = localPostgresPoolForTests(url) as Pool & {
      connect: () => Promise<PoolClient>;
    };
    const original = pool.connect.bind(pool);
    pool.connect = async () => {
      const client = await original();
      if (!queryWrapped.has(client)) {
        queryWrapped.add(client);
        const query = client.query.bind(client) as (input: unknown, values?: unknown) => Promise<unknown>;
        client.query = ((input: unknown, values?: unknown) => {
          benchCounter.roundtrips += 1;
          return values === undefined ? query(input) : query(input, values);
        }) as PoolClient["query"];
      }
      return client;
    };
    restores.push(() => {
      pool.connect = original;
    });
  }
  const started = performance.now();
  try {
    const payload = await run();
    return {
      roundtrips: benchCounter.roundtrips,
      elapsedMs: performance.now() - started,
      payload: JSON.stringify(payload),
    };
  } finally {
    for (const restore of restores) restore();
  }
}

async function sequence(url: string, serviceUrl: string): Promise<unknown[]> {
  const user = createLocalPostgresDataClient(url, { id: userId, email: "person@example.com" }, "authenticated");
  const service = createLocalPostgresDataClient(serviceUrl, null, "service_role");
  const profile = await user.from("profiles").select("id,display_name").eq("id", userId).maybeSingle();
  const charts = await user.from("chart_profiles").select("id,label").eq("user_id", userId).order("updated_at", { ascending: false }).limit(40);
  const sessions = await user.from("chat_sessions").select("id,title").eq("user_id", userId).order("updated_at", { ascending: false }).limit(40);
  const updated = await user.from("chat_sessions").update({ title: "consult-bench" }).eq("id", sessionId);
  const upsertConflict = await user.from("chat_sessions").upsert(
    { id: sessionId, user_id: userId, title: "consult-bench" },
    { onConflict: "id" },
  );
  const upsertKey = await user.from("chat_sessions").upsert({ id: sessionId, user_id: userId, title: "consult-bench" });
  const rpc = await service.rpc("session_probe");
  const serviceRead = await service.from("profiles").select("id").eq("id", userId).maybeSingle();
  return [profile, charts, sessions, updated, upsertConflict, upsertKey, rpc, serviceRead].map((result) => ({
    data: result.data ?? null,
    error: result.error ?? null,
    count: "count" in result ? result.count ?? null : null,
  }));
}

async function measure(url: string, serviceUrl: string, mode: Mode): Promise<{
  roundtrips: number;
  medianMs: number;
  p95Ms: number;
  payload: string;
  groups: Array<{ roundtrips: number; medianMs: number }>;
}> {
  process.env.JYOTISHA_PG_METADATA_CACHE = mode.cache;
  process.env.JYOTISHA_PG_MERGED_SESSION_SETUP = mode.merged;
  clearLocalPostgresMetadataCacheForTests();
  await sequence(url, serviceUrl);
  const groupMedians: number[] = [];
  const groupRoundtrips: number[] = [];
  const allTimes: number[] = [];
  let payload = "";
  for (let index = 0; index < 3; index += 1) {
    const group: number[] = [];
    let roundtrips = 0;
    for (let repeat = 0; repeat < 50; repeat += 1) {
      const sample = await withCountingPools([url, serviceUrl], () => sequence(url, serviceUrl));
      group.push(sample.elapsedMs);
      allTimes.push(sample.elapsedMs);
      roundtrips += sample.roundtrips;
      payload = sample.payload;
    }
    groupMedians.push(median(group));
    groupRoundtrips.push(Math.round(roundtrips / 50));
  }
  return {
    roundtrips: median(groupRoundtrips),
    medianMs: median(groupMedians),
    p95Ms: percentile(allTimes, 0.95),
    payload,
    groups: groupRoundtrips.map((roundtrips, index) => ({
      roundtrips,
      medianMs: Number((groupMedians[index] ?? 0).toFixed(2)),
    })),
  };
}

async function main(): Promise<void> {
  const docker = spawnSync("docker", ["info"], { encoding: "utf8" });
  if (docker.status !== 0) {
    console.log(JSON.stringify({ ok: false, reason: "docker-unavailable" }));
    process.exit(2);
  }
  const fixture = startPostgresFixture();
  const appUrl = fixture.connectionUrl("app_runtime", "app-runtime-test-password");
  const serviceUrl = fixture.connectionUrl("service_runtime", "service-runtime-test-password");
  try {
    fixture.psqlScriptAs("postgres", "postgres-test-password", schemaSql);
    const modes: Record<string, Mode> = {
      A: { cache: "off", merged: "off" },
      B: { cache: "on", merged: "on" },
      B1: { cache: "on", merged: "off" },
      B2: { cache: "off", merged: "on" },
    };
    const results: Record<string, { roundtrips: number; medianMs: number; p95Ms: number; groups: Array<{ roundtrips: number; medianMs: number }> }> = {};
    let payloadA = "";
    let payloadMismatch = false;
    for (const [name, mode] of Object.entries(modes)) {
      const measured = await measure(appUrl, serviceUrl, mode);
      results[name] = {
        roundtrips: measured.roundtrips,
        medianMs: Number(measured.medianMs.toFixed(2)),
        p95Ms: Number(measured.p95Ms.toFixed(2)),
        groups: measured.groups,
      };
      if (name === "A") payloadA = measured.payload;
      else if (measured.payload !== payloadA) payloadMismatch = true;
      clearLocalPostgresMetadataCacheForTests();
      await closeLocalPostgresDataPools();
    }
    console.log(JSON.stringify({
      ok: true,
      repeats: 3,
      repsPerGroup: 50,
      payloadMismatch,
      results,
      node: process.version,
    }));
  } finally {
    await closeLocalPostgresDataPools();
    fixture.stop();
    delete process.env.JYOTISHA_PG_METADATA_CACHE;
    delete process.env.JYOTISHA_PG_MERGED_SESSION_SETUP;
  }
}

main().catch((error: unknown) => {
  const message = error instanceof Error ? error.message : "bench failed";
  console.log(JSON.stringify({ ok: false, reason: message }));
  process.exit(1);
});