Merge remote-tracking branch 'origin/main' into samejr/Settings-page-layouts

Author: samejr

.server-changes/realtime-run-reads-from-primary.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: improvement
+---
+
+Realtime run subscriptions can now be configured to read run data straight from the primary database, so a run's latest state is never served from a lagging replica. Off by default; replica reads are unchanged unless you turn it on.

apps/webapp/app/env.server.ts

@@ -462,6 +462,7 @@ const EnvironmentSchema = z
     // TTL/size of the per-org realtimeBackend flag cache used to pick the serving backend.
     REALTIME_BACKEND_FLAG_CACHE_TTL_MS: z.coerce.number().int().default(30_000),
     REALTIME_BACKEND_FLAG_CACHE_MAX_ENTRIES: z.coerce.number().int().default(50_000),
+    REALTIME_BACKEND_NATIVE_RUN_READS_FROM_PRIMARY: z.string().default("0"),
     // "1" enables the read-your-writes gate: wake hydrates wait out the measured replica lag
     // (anchored to the change record's updatedAtMs) and stale reads are retried.
     REALTIME_BACKEND_NATIVE_REPLICA_LAG_GATE_ENABLED: z.string().default("1"),

apps/webapp/app/services/realtime/nativeRealtimeClientInstance.server.ts

@@ -1,5 +1,5 @@
 import { getMeter } from "@internal/tracing";
-import { $replica } from "~/db.server";
+import { $replica, prisma } from "~/db.server";
 import { runStore } from "~/v3/runStore.server";
 import { env } from "~/env.server";
 import { singleton } from "~/utils/singleton";
@@ -130,9 +130,11 @@ function initializeNativeRealtimeClient(): NativeRealtimeClient {
         })
       : undefined;
 
+  const runReadsFromPrimary = env.REALTIME_BACKEND_NATIVE_RUN_READS_FROM_PRIMARY === "1";
+
   // One RunHydrator shared by the router and the client, so its single-flight + short-TTL cache covers both.
   const runReader = new RunHydrator({
-    replica: $replica,
+    readClient: runReadsFromPrimary ? prisma : $replica,
     runStore,
     cacheTtlMs: env.REALTIME_BACKEND_NATIVE_RUN_CACHE_TTL_MS,
     maxCacheEntries: env.REALTIME_BACKEND_NATIVE_RUN_CACHE_MAX_ENTRIES,
@@ -142,7 +144,7 @@ function initializeNativeRealtimeClient(): NativeRealtimeClient {
   // when idle) and the router delays wake hydrates by it, anchored to each record's
   // updatedAtMs — so a publish racing the replica's apply is waited out, not read stale.
   const lagEstimator =
-    env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_GATE_ENABLED === "1"
+    !runReadsFromPrimary && env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_GATE_ENABLED === "1"
       ? new ReplicaLagEstimator({
           source: createPostgresReplicaLagSource($replica),
           sampleIntervalMs: env.REALTIME_BACKEND_NATIVE_REPLICA_LAG_SAMPLE_INTERVAL_MS,

apps/webapp/app/services/realtime/replicaLagEstimator.server.ts

@@ -25,13 +25,25 @@ export interface ReplicaLagSource {
 /** Aurora: replicas share the storage layer and reject every standard WAL function;
  * `aurora_replica_status()` is the only live lag source. Max across readers, since the
  * `$replica` pool balances over all of them. No reader rows = `$replica` is the writer =
- * no lag. Throws on non-Aurora (the function doesn't exist). */
+ * no lag. Throws on non-Aurora, but detects that with a `to_regproc` lookup rather than by
+ * letting the call fail — an unresolvable function is a query error the driver reports to the
+ * error log (and Sentry) regardless of the app-level catch, once per sample. */
 export class AuroraReplicaLagSource implements ReplicaLagSource {
   readonly name = "aurora";
+  #available: boolean | undefined;
 
   constructor(private readonly db: RawQueryable) {}
 
   async sampleLagMs(): Promise<number | undefined> {
+    if (this.#available === undefined) {
+      const probe = await this.db.$queryRawUnsafe<{ available: boolean | null }[]>(
+        `SELECT to_regproc('aurora_replica_status') IS NOT NULL AS available`
+      );
+      this.#available = probe[0]?.available === true;
+    }
+    if (!this.#available) {
+      throw new Error("aurora_replica_status() is not available on this database");
+    }
     const rows = await this.db.$queryRawUnsafe<{ lag: number | null }[]>(
       `SELECT max(replica_lag_in_msec)::float8 AS lag FROM aurora_replica_status() WHERE session_id <> 'MASTER_SESSION_ID' AND replica_lag_in_msec IS NOT NULL`
     );

apps/webapp/app/services/realtime/runReader.server.ts

@@ -82,9 +82,11 @@ export interface RunListResolver {
 }
 
 export type RunHydratorOptions = {
-  /** A read-replica Prisma client (`$replica`). Always Postgres. */
-  replica: Pick<PrismaClient, "taskRun">;
-  /** RunStore the reads are routed through; `replica` is passed as the read client. */
+  /** The Prisma client handed to the RunStore as the read client. Always Postgres. A branded
+   * replica (`$replica`) keeps routed reads on each store's replica; an unbranded writer
+   * (`prisma`) escalates them to each store's own primary. */
+  readClient: Pick<PrismaClient, "taskRun">;
+  /** RunStore the reads are routed through. */
   runStore: RunStore;
   /** Read-through cache TTL (ms) collapsing duplicate refetches for the same run. Set 0 to disable. Defaults to 250ms. */
   cacheTtlMs?: number;
@@ -154,7 +156,7 @@ export class RunHydrator {
         },
         select: buildHydratorSelect(skipColumns),
       },
-      this.options.replica as PrismaClientOrTransaction
+      this.options.readClient as PrismaClientOrTransaction
     );
     return rows as unknown as RealtimeRunRow[];
   }
@@ -166,7 +168,7 @@ export class RunHydrator {
         runtimeEnvironmentId: environmentId,
       },
       { select: RUN_HYDRATOR_SELECT },
-      this.options.replica as PrismaClientOrTransaction
+      this.options.readClient as PrismaClientOrTransaction
     );
 
     return (run ?? null) as RealtimeRunRow | null;

apps/webapp/app/services/realtime/shadowRealtimeClientInstance.server.ts

@@ -21,7 +21,7 @@ function initializeShadowRealtimeClient(): ShadowRealtimeClient {
   });
 
   const comparator = new RealtimeShadowComparator({
-    runReader: new RunHydrator({ replica: $replica, runStore }),
+    runReader: new RunHydrator({ readClient: $replica, runStore }),
     runListResolver: new ClickHouseRunListResolver({
       getClickhouse: (organizationId) =>
         clickhouseFactory.getClickhouseForOrganization(organizationId, "realtime"),

apps/webapp/test/realtime/replicaLagEstimator.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, describe, expect, it } from "vitest";
 import {
+  AuroraReplicaLagSource,
   FirstSupportedReplicaLagSource,
   ReplicaLagEstimator,
   type ReplicaLagSource,
@@ -162,3 +163,43 @@ describe("FirstSupportedReplicaLagSource", () => {
     expect(composed.name).toBe("flaky");
   });
 });
+
+describe("AuroraReplicaLagSource", () => {
+  function fakeDb(available: boolean) {
+    const queries: string[] = [];
+    return {
+      queries,
+      db: {
+        async $queryRawUnsafe<T = unknown>(query: string): Promise<T> {
+          queries.push(query);
+          if (query.includes("to_regproc")) {
+            return [{ available: available ? true : null }] as T;
+          }
+          return [{ lag: 12.5 }] as T;
+        },
+      },
+    };
+  }
+
+  it("never puts aurora_replica_status() on the wire when the function is absent", async () => {
+    const { db, queries } = fakeDb(false);
+    const aurora = new AuroraReplicaLagSource(db);
+
+    await expect(aurora.sampleLagMs()).rejects.toThrow(/not available/);
+    await expect(aurora.sampleLagMs()).rejects.toThrow(/not available/);
+
+    expect(queries.filter((q) => q.includes("FROM aurora_replica_status()"))).toHaveLength(0);
+    expect(queries.filter((q) => q.includes("to_regproc"))).toHaveLength(1);
+  });
+
+  it("samples lag on Aurora and probes for the function only once", async () => {
+    const { db, queries } = fakeDb(true);
+    const aurora = new AuroraReplicaLagSource(db);
+
+    expect(await aurora.sampleLagMs()).toBe(12.5);
+    expect(await aurora.sampleLagMs()).toBe(12.5);
+
+    expect(queries.filter((q) => q.includes("to_regproc"))).toHaveLength(1);
+    expect(queries.filter((q) => q.includes("FROM aurora_replica_status()"))).toHaveLength(2);
+  });
+});

apps/webapp/test/realtime/runReaderProjection.test.ts

@@ -56,7 +56,10 @@ describe("RunHydrator.hydrateByIds column projection", () => {
       },
     } as any;
     const runStore = new PostgresRunStore({ prisma: replica, readOnlyPrisma: replica });
-    return { hydrator: new RunHydrator({ replica, runStore }), getSelect: () => capturedSelect };
+    return {
+      hydrator: new RunHydrator({ readClient: replica, runStore }),
+      getSelect: () => capturedSelect,
+    };
   }
 
   it("projects the SELECT by skipColumns", async () => {

apps/webapp/test/realtime/runReaderReadThrough.test.ts

@@ -218,7 +218,7 @@ describe("RunHydrator read-route through the runStore seam (legacy + new)", () =
       });
 
       const runStore = makeRoutingShapedStore({ newStore, legacyStore });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore });
 
       const rows = await hydrator.hydrateByIds(envId, [newRunId, legacyRunId]);
       expect(rows.map((r) => r.id).sort()).toEqual([legacyRunId, newRunId].sort());
@@ -267,7 +267,7 @@ describe("RunHydrator read-route through the runStore seam (legacy + new)", () =
       });
 
       const runStore = makeRoutingShapedStore({ newStore, legacyStore });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore });
 
       const row = await hydrator.getRunById(envId, migratedRunId);
       expect(row?.id).toBe(migratedRunId);
@@ -300,7 +300,7 @@ describe("RunHydrator read-route through the runStore seam (legacy + new)", () =
       });
 
       const runStore = makeRoutingShapedStore({ newStore, legacyStore });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore });
 
       const byId = await hydrator.getRunById(envId, oldRunId);
       expect(byId?.payload).toBe('{"era":"old"}');
@@ -338,7 +338,7 @@ describe("RunHydrator read-route through the runStore seam (legacy + new)", () =
 
       // A generic legacy replica would miss the NEW row entirely — the metadata must come off NEW.
       const runStore = makeRoutingShapedStore({ newStore, legacyStore });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore, cacheTtlMs: 0 });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore, cacheTtlMs: 0 });
 
       const snapshot = await hydrator.getRunById(envId, terminalRunId);
       expect(snapshot?.id).toBe(terminalRunId);
@@ -385,7 +385,7 @@ describe("RunHydrator read-route through the runStore seam (legacy + new)", () =
       // Use a 0ms TTL so each getRunById re-reads through the seam (no cached stale row across the
       // crossing). Single-flight/TTL are proven separately below.
       const runStore = makeRoutingShapedStore({ newStore, legacyStore, classify });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore, cacheTtlMs: 0 });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore, cacheTtlMs: 0 });
 
       // Before migration: served from LEGACY.
       const before = await hydrator.getRunById(envId, runId);
@@ -434,7 +434,7 @@ describe("RunHydrator single-flight + TTL cache intact across the seam", () => {
       });
 
       const runStore = makeRoutingShapedStore({ newStore, legacyStore });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore, cacheTtlMs: 60_000 });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore, cacheTtlMs: 60_000 });
 
       // Two concurrent calls -> single-flight collapses to ONE underlying read.
       const [a, b] = await Promise.all([
@@ -466,7 +466,7 @@ describe("RunHydrator single-flight + TTL cache intact across the seam", () => {
       const missingRunId = newId("missing_run");
 
       const runStore = makeRoutingShapedStore({ newStore, legacyStore });
-      const hydrator = new RunHydrator({ replica: prisma14, runStore, cacheTtlMs: 60_000 });
+      const hydrator = new RunHydrator({ readClient: prisma14, runStore, cacheTtlMs: 60_000 });
 
       const first = await hydrator.getRunById(envId, missingRunId);
       expect(first).toBeNull();
@@ -504,7 +504,7 @@ describe("RunHydrator single-DB passthrough (one PostgresRunStore over one clien
         });
       }
 
-      const hydrator = new RunHydrator({ replica: prisma, runStore: store, cacheTtlMs: 60_000 });
+      const hydrator = new RunHydrator({ readClient: prisma, runStore: store, cacheTtlMs: 60_000 });
 
       // hydrateByIds returns both rows from the single client.
       const rows = await hydrator.hydrateByIds(envId, [runIdA, runIdB]);
@@ -523,7 +523,7 @@ describe("RunHydrator single-DB passthrough (one PostgresRunStore over one clien
   postgresTest("empty id-set returns [] without touching the store", async ({ prisma }) => {
     const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
     const findRunsSpy = vi.spyOn(store, "findRuns");
-    const hydrator = new RunHydrator({ replica: prisma, runStore: store });
+    const hydrator = new RunHydrator({ readClient: prisma, runStore: store });
 
     const rows = await hydrator.hydrateByIds("env_none", []);
     expect(rows).toEqual([]);

apps/webapp/test/realtimeServices.replicaLag.test.ts

@@ -216,7 +216,7 @@ describe("realtime-svc — replica-lag guards", () => {
       const replica = laggingReplica(prisma, [{ model: "taskRun", mode: "missing" }]);
       // The hydrator holds the real store; hydrateByIds passes `options.replica` as the read client.
       const hydrator = new RunHydrator({
-        replica: replica.client as PrismaClient,
+        readClient: replica.client as PrismaClient,
         runStore: writerStore,
         cacheTtlMs: 0,
       });
@@ -263,7 +263,7 @@ describe("realtime-svc — replica-lag guards", () => {
 
       const replica = laggingReplica(prisma, [{ model: "taskRun", mode: "missing" }]);
       const hydrator = new RunHydrator({
-        replica: replica.client as PrismaClient,
+        readClient: replica.client as PrismaClient,
         runStore: writerStore,
         cacheTtlMs: 0,
       });