Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .server-changes/realtime-run-reads-from-primary.md
Original file line number Diff line number Diff line change
@@ -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.
1 change: 1 addition & 0 deletions apps/webapp/app/env.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand Down
Original file line number Diff line number Diff line change
@@ -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";
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
14 changes: 13 additions & 1 deletion apps/webapp/app/services/realtime/replicaLagEstimator.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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`
);
Expand Down
12 changes: 7 additions & 5 deletions apps/webapp/app/services/realtime/runReader.server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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[];
}
Expand All @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
Expand Down
41 changes: 41 additions & 0 deletions apps/webapp/test/realtime/replicaLagEstimator.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { afterEach, describe, expect, it } from "vitest";
import {
AuroraReplicaLagSource,
FirstSupportedReplicaLagSource,
ReplicaLagEstimator,
type ReplicaLagSource,
Expand Down Expand Up @@ -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);
});
});
5 changes: 4 additions & 1 deletion apps/webapp/test/realtime/runReaderProjection.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down
18 changes: 9 additions & 9 deletions apps/webapp/test/realtime/runReaderReadThrough.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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"}');
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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([
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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]);
Expand All @@ -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([]);
Expand Down
4 changes: 2 additions & 2 deletions apps/webapp/test/realtimeServices.replicaLag.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
});
Expand Down Expand Up @@ -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,
});
Expand Down