diff --git a/.server-changes/fix-run-store-routing-read-your-writes.md b/.server-changes/fix-run-store-routing-read-your-writes.md new file mode 100644 index 0000000000..1b70c8aeb1 --- /dev/null +++ b/.server-changes/fix-run-store-routing-read-your-writes.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +Fix run store routing dropping the caller's read client, which downgraded read-your-writes reads (execution snapshots, waitpoints, batches) to the read replica and could fail dequeues under replica lag diff --git a/apps/webapp/app/db.server.ts b/apps/webapp/app/db.server.ts index 0284c19246..7805628f72 100644 --- a/apps/webapp/app/db.server.ts +++ b/apps/webapp/app/db.server.ts @@ -8,6 +8,7 @@ import { type PrismaTransactionOptions, } from "@trigger.dev/database"; import { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { markReadReplicaClient } from "@internal/run-store"; import invariant from "tiny-invariant"; import { z } from "zod"; import { env } from "./env.server"; @@ -170,7 +171,11 @@ export const prisma = singleton("prisma", () => export const $replica: PrismaReplicaClient = singleton("replica", () => { const replica = getReplicaClient(); - return replica ? captureInfrastructureErrors(tagDatasource("replica", replica)) : prisma; + // Brand ONLY a real replica so the run-store routing layer keeps replica reads off the primary. + // No replica configured → fall back to the writer `prisma`, which must stay UNBRANDED. + return replica + ? markReadReplicaClient(captureInfrastructureErrors(tagDatasource("replica", replica))) + : prisma; }); export type RunOpsClients = { writer: PrismaClient; replica: PrismaReplicaClient }; @@ -254,9 +259,14 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => { captureInfraErrorsRunOps( tagDatasourceRunOps("writer", buildRunOpsWriterClient({ url, clientType })) ), + // Brand the run-ops replica (only built for a real replica URL) so routed replica reads stay + // off the primary. When no replica URL is set, selectRunOpsTopology reuses the writer here — + // which this callback never touches, so the writer stays unbranded. buildNewReplica: (url, clientType) => - captureInfraErrorsRunOps( - tagDatasourceRunOps("replica", buildRunOpsReplicaClient({ url, clientType })) + markReadReplicaClient( + captureInfraErrorsRunOps( + tagDatasourceRunOps("replica", buildRunOpsReplicaClient({ url, clientType })) + ) ), } ); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 843f4c3b32..b797fae9b9 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -507,6 +507,13 @@ export class PostgresRunStore implements RunStore { this.schemaVariant = options.schemaVariant ?? "legacy"; } + // The writer handle in read-client form, so the routing layer can honor a caller-passed client + // (read-your-writes) with THIS store's own primary instead of leaking the caller's client across + // DBs. Cast mirrors runInTransaction: the generated clients differ only in delegates reads use. + get primaryReadClient(): ReadClient { + return this.prisma as unknown as ReadClient; + } + // Open ONE interactive transaction on this store's OWN writer client and run `fn` against THIS store // (so subclass overrides survive) with the tx as the client to thread into the inner writes. `runId` // is ignored here — a single store has one connection — but is in the contract so the router can diff --git a/internal-packages/run-store/src/index.ts b/internal-packages/run-store/src/index.ts index 81f88ed6d1..160f9cdada 100644 --- a/internal-packages/run-store/src/index.ts +++ b/internal-packages/run-store/src/index.ts @@ -1,3 +1,4 @@ export * from "./types.js"; export * from "./PostgresRunStore.js"; export * from "./runOpsStore.js"; +export * from "./readReplicaClient.js"; diff --git a/internal-packages/run-store/src/readReplicaClient.ts b/internal-packages/run-store/src/readReplicaClient.ts new file mode 100644 index 0000000000..d5f2028cb1 --- /dev/null +++ b/internal-packages/run-store/src/readReplicaClient.ts @@ -0,0 +1,26 @@ +// A writer and a read-replica Prisma client are structurally identical at runtime (a replica is a +// `new PrismaClient(...)` too, so it also exposes `$transaction`). The routing layer therefore +// cannot tell a caller-passed replica from a writer by shape, yet it must — a writer/tx read has to +// reach the owning store's PRIMARY (read-your-writes) while a replica read stays on a replica (read +// scaling). So the client builder brands replica handles and the routing store reads the brand. +export const READ_REPLICA_BRAND: unique symbol = Symbol.for("trigger.dev/run-store/read-replica"); + +// Brand a replica client (returns the same object). MUST only be called on a genuine replica: the +// routing layer trusts the brand to mean "do not escalate this read to the primary". An unbranded +// replica just escalates as before — a scaling miss, never a correctness bug. +export function markReadReplicaClient(client: T): T { + try { + (client as Record)[READ_REPLICA_BRAND] = true; + } catch { + // Frozen/exotic clients may reject the assignment; only costs the optimization, not correctness. + } + return client; +} + +export function isReadReplicaClient(client: unknown): boolean { + return ( + !!client && + typeof client === "object" && + (client as Record)[READ_REPLICA_BRAND] === true + ); +} diff --git a/internal-packages/run-store/src/runOpsStore.routedReadPrimary.test.ts b/internal-packages/run-store/src/runOpsStore.routedReadPrimary.test.ts new file mode 100644 index 0000000000..4e9b191de5 --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.routedReadPrimary.test.ts @@ -0,0 +1,375 @@ +// RED→GREEN repro for the routed-read CLIENT DROP: RoutingRunStore's runId-routed / fan-out reads +// accept `client?: ReadClient` but dropped it, so the sub-store fell back to its REPLICA. The +// run-engine passes its writer (`tx ?? this.$.prisma`) into these reads for read-your-writes +// consistency (dequeue re-reads the just-written QUEUED snapshot), so the drop surfaces in cloud as +// TASK_DEQUEUED_INVALID_STATE / "No execution snapshot found for TaskRun ...". The fix routes a +// caller-passed client to the OWNING store's OWN primary (never forwarded verbatim — it is bound to +// the control-plane DB); no client keeps the replica default. +// +// Deterministic harness: `heteroPostgresTest` hands two PHYSICALLY separate postgres containers +// over the same full schema. The owning store WRITES to one and its `readOnlyPrisma` points at the +// other, which stays EMPTY (a replica with unbounded lag) — so a replica-routed read MISSES and a +// primary-routed read finds the row, with no replica==primary aliasing to mask the drop. + +import { heteroPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; +import { markReadReplicaClient } from "./readReplicaClient.js"; +import type { CreateRunInput } from "./types.js"; + +// ownerEngine classifies by internal-id LENGTH: 25 chars → cuid → LEGACY, 27 → ksuid → NEW. +const CUID_25 = "c".repeat(25); +const KSUID_27 = "k".repeat(27); + +// Router topology where the OWNING store (the one the test's run ids route to) writes to `writer` +// but reads by default from `lagging` — a physically separate, never-written DB. The other store +// lives entirely on the lagging DB so fan-out legs can't accidentally see rows. Both DBs carry the +// full schema (the forwarding under test is residency-agnostic; dedicated-subset parity is covered +// by the sibling suites), so both stores use the "legacy" variant. +function splitTopology( + residency: "LEGACY" | "NEW", + writer: PrismaClient, + lagging: PrismaClient +): { owningStore: PostgresRunStore; router: RoutingRunStore } { + const owningStore = new PostgresRunStore({ + prisma: writer, + readOnlyPrisma: lagging, + schemaVariant: "legacy", + }); + const otherStore = new PostgresRunStore({ + prisma: lagging, + readOnlyPrisma: lagging, + schemaVariant: "legacy", + }); + const router = new RoutingRunStore( + residency === "LEGACY" + ? { new: otherStore, legacy: owningStore } + : { new: owningStore, legacy: otherStore } + ); + return { owningStore, router }; +} + +async function seedEnvironment(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +describe("run-ops split — routed reads honor a caller-passed client via the owning store's PRIMARY", () => { + // The outage path: dequeue writes the QUEUED snapshot then re-reads it via + // `getLatestExecutionSnapshot(this.$.prisma, ...)`. The router must not downgrade that read to + // the replica. Covers the whole snapshot read family on the LEGACY (cuid) routing arm. + heteroPostgresTest( + "LEGACY cuid: snapshot reads with a client resolve on the owning primary; without, on the replica", + async ({ prisma14, prisma17 }) => { + const { router } = splitTopology("LEGACY", prisma14, prisma17); + const seed = await seedEnvironment(prisma14, "snap_leg"); + const runId = `run_${CUID_25}`; + await router.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_snap_leg", + organizationId: seed.organization.id, + projectId: seed.project.id, + runtimeEnvironmentId: seed.environment.id, + }) + ); + + // findLatestExecutionSnapshot — client passed → owning primary finds the fresh snapshot. + const latest = await router.findLatestExecutionSnapshot(runId, prisma14); + expect(latest).not.toBeNull(); + expect(latest?.executionStatus).toBe("RUN_CREATED"); + // No client → the (empty) replica, unchanged behavior. + expect(await router.findLatestExecutionSnapshot(runId)).toBeNull(); + + const snapshotId = latest!.id; + + // findExecutionSnapshot (runId-routed, the warm-restart shape). + const one = await router.findExecutionSnapshot( + { where: { runId, isValid: true }, select: { id: true } }, + prisma14 + ); + expect(one?.id).toBe(snapshotId); + expect(await router.findExecutionSnapshot({ where: { runId, isValid: true } })).toBeNull(); + + // findManyExecutionSnapshots (runId-routed). + const many = await router.findManyExecutionSnapshots( + { where: { runId }, select: { id: true } }, + prisma14 + ); + expect(many.map((s) => s.id)).toEqual([snapshotId]); + expect(await router.findManyExecutionSnapshots({ where: { runId } })).toEqual([]); + + // findSnapshotCompletedWaitpointIds (the resume-payload join read). Seed a completed + // waitpoint + its `_completedWaitpoints` join on the writer only. + const waitpoint = await prisma14.waitpoint.create({ + data: { + friendlyId: "waitpoint_snap_leg", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem_snap_leg", + userProvidedIdempotencyKey: false, + projectId: seed.project.id, + environmentId: seed.environment.id, + }, + }); + // Link via the Prisma relation API (not a raw insert into the implicit join table) so a + // relation rename fails at compile time rather than silently seeding nothing. + await prisma14.taskRunExecutionSnapshot.update({ + where: { id: snapshotId }, + data: { completedWaitpoints: { connect: { id: waitpoint.id } } }, + }); + expect(await router.findSnapshotCompletedWaitpointIds(snapshotId, prisma14)).toEqual([ + waitpoint.id, + ]); + expect(await router.findSnapshotCompletedWaitpointIds(snapshotId)).toEqual([]); + } + ); + + // NEW (ksuid) routing arm. The caller's client here is the CONTROL-PLANE writer — the wrong + // physical DB for a NEW-resident run — so this also pins that the client is never forwarded + // verbatim: the read must resolve on the owning NEW store's OWN primary. + heteroPostgresTest( + "NEW ksuid: a control-plane client routes the snapshot read to the NEW store's OWN primary", + async ({ prisma14, prisma17 }) => { + // Owning (NEW) store writes to prisma14; the control-plane/other store is prisma17. + const { router } = splitTopology("NEW", prisma14, prisma17); + const seed = await seedEnvironment(prisma14, "snap_new"); + const runId = `run_${KSUID_27}`; + await router.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_snap_new", + organizationId: seed.organization.id, + projectId: seed.project.id, + runtimeEnvironmentId: seed.environment.id, + }) + ); + + // Control-plane writer (prisma17 side of this topology) passed as the client: the row can + // only be found on the NEW store's own primary (prisma14) — verbatim forwarding would miss. + const latest = await router.findLatestExecutionSnapshot(runId, prisma17); + expect(latest).not.toBeNull(); + expect(latest?.executionStatus).toBe("RUN_CREATED"); + // No client → the NEW store's (empty) replica. + expect(await router.findLatestExecutionSnapshot(runId)).toBeNull(); + } + ); + + heteroPostgresTest( + "LEGACY cuid: findRuns fan-out and batch friendlyId probe honor a caller client", + async ({ prisma14, prisma17 }) => { + const { router } = splitTopology("LEGACY", prisma14, prisma17); + const seed = await seedEnvironment(prisma14, "runs_leg"); + const runId = `run_${CUID_25}`; + await router.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_runs_leg", + organizationId: seed.organization.id, + projectId: seed.project.id, + runtimeEnvironmentId: seed.environment.id, + }) + ); + + // findRuns (bounded id-set fan-out). + const rows = await router.findRuns( + { where: { id: { in: [runId] } }, select: { id: true } }, + prisma14 + ); + expect(rows.map((r) => r.id)).toEqual([runId]); + expect( + await router.findRuns({ where: { id: { in: [runId] } }, select: { id: true } }) + ).toEqual([]); + + // findBatchTaskRunByFriendlyId (env-scoped fan-out probe; the one batch read that defaults + // to the replica). + const batch = await prisma14.batchTaskRun.create({ + data: { + id: `batch_${CUID_25}`, + friendlyId: "batch_runs_leg", + runtimeEnvironmentId: seed.environment.id, + }, + }); + const viaPrimary = await router.findBatchTaskRunByFriendlyId( + batch.friendlyId, + seed.environment.id, + undefined, + prisma14 + ); + expect(viaPrimary?.id).toBe(batch.id); + expect( + await router.findBatchTaskRunByFriendlyId(batch.friendlyId, seed.environment.id) + ).toBeNull(); + } + ); + + // Read scaling: a caller that passes an explicit READ REPLICA (e.g. `$replica`) must NOT be + // escalated to the primary — only true read-your-writes (a writer/tx) should. A replica is a + // full PrismaClient at runtime (it has `$transaction` too), so shape can't distinguish it; the + // client builder brands it and the router honors the brand. Proven here by branding a client and + // showing the read stays on the owning store's (empty) replica — same as passing no client — + // while an unbranded writer escalates and finds the fresh row. + heteroPostgresTest( + "a branded read-replica client stays on the replica; a writer escalates to the primary", + async ({ prisma14, prisma17 }) => { + const { router } = splitTopology("LEGACY", prisma14, prisma17); + const seed = await seedEnvironment(prisma14, "replica_leg"); + const runId = `run_${CUID_25}`; + await router.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_replica_leg", + organizationId: seed.organization.id, + projectId: seed.project.id, + runtimeEnvironmentId: seed.environment.id, + }) + ); + + // The router discards the caller's client object (it can't cross DBs) and reads the brand + // only, so a branded marker faithfully stands in for a passed `$replica`. + const replicaClient = markReadReplicaClient({} as unknown as PrismaClient); + + // findRun (readYourWrites path): branded replica → owning replica (empty) → miss. + expect( + await router.findRun({ id: runId }, { select: { id: true } }, replicaClient) + ).toBeNull(); + // Control: an unbranded writer escalates to the owning primary → finds the fresh row. + expect((await router.findRun({ id: runId }, { select: { id: true } }, prisma14))?.id).toBe( + runId + ); + + // findLatestExecutionSnapshot (#ownPrimary path): same replica-stays-on-replica invariant. + expect(await router.findLatestExecutionSnapshot(runId, replicaClient)).toBeNull(); + expect((await router.findLatestExecutionSnapshot(runId, prisma14))?.executionStatus).toBe( + "RUN_CREATED" + ); + // No client behaves identically to the branded replica. + expect(await router.findLatestExecutionSnapshot(runId)).toBeNull(); + } + ); + + heteroPostgresTest( + "LEGACY cuid: waitpoint reads honor a caller client", + async ({ prisma14, prisma17 }) => { + const { router } = splitTopology("LEGACY", prisma14, prisma17); + const seed = await seedEnvironment(prisma14, "wp_leg"); + const runId = `run_${CUID_25}`; + await router.createRun( + buildCreateRunInput({ + runId, + friendlyId: "run_wp_leg", + organizationId: seed.organization.id, + projectId: seed.project.id, + runtimeEnvironmentId: seed.environment.id, + }) + ); + const waitpoint = await prisma14.waitpoint.create({ + data: { + id: `waitpoint_${CUID_25}`, + friendlyId: "waitpoint_wp_leg", + type: "MANUAL", + status: "PENDING", + idempotencyKey: "idem_wp_leg", + userProvidedIdempotencyKey: false, + projectId: seed.project.id, + environmentId: seed.environment.id, + }, + }); + await prisma14.taskRunWaitpoint.create({ + data: { taskRunId: runId, waitpointId: waitpoint.id, projectId: seed.project.id }, + }); + + // findWaitpoint (by id, resolve + scalar read). + const found = await router.findWaitpoint({ where: { id: waitpoint.id } }, prisma14); + expect(found?.id).toBe(waitpoint.id); + expect(await router.findWaitpoint({ where: { id: waitpoint.id } })).toBeNull(); + + // findManyWaitpoints (both-store fan-out). + const manyOnPrimary = await router.findManyWaitpoints( + { where: { id: { in: [waitpoint.id] } } }, + prisma14 + ); + expect(manyOnPrimary.map((w) => w.id)).toEqual([waitpoint.id]); + expect(await router.findManyWaitpoints({ where: { id: { in: [waitpoint.id] } } })).toEqual( + [] + ); + + // countPendingWaitpoints (both-store fan-out sum). + expect(await router.countPendingWaitpoints([waitpoint.id], prisma14)).toBe(1); + expect(await router.countPendingWaitpoints([waitpoint.id])).toBe(0); + + // findManyTaskRunWaitpoints (the blocked-run edge fan-out). + const edges = await router.findManyTaskRunWaitpoints( + { where: { taskRunId: runId } }, + prisma14 + ); + expect(edges).toHaveLength(1); + expect(edges[0]?.waitpointId).toBe(waitpoint.id); + expect(await router.findManyTaskRunWaitpoints({ where: { taskRunId: runId } })).toEqual([]); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.test.ts b/internal-packages/run-store/src/runOpsStore.test.ts index d8a4a11c86..6d65616dea 100644 --- a/internal-packages/run-store/src/runOpsStore.test.ts +++ b/internal-packages/run-store/src/runOpsStore.test.ts @@ -1257,6 +1257,39 @@ describe("RoutingRunStore.findRuns split-mode fan-out + drain", () => { } ); + // forWaitpointCompletion selects the store a subsequent WRITE (updateManyWaitpoints) lands on, + // so its resolution probe must read each store's PRIMARY — not the replica. Here the owning + // (NEW) store's replica lags (points at an empty DB), so a replica probe would miss the fresh + // waitpoint and mis-resolve to the id-shape's default (LEGACY), stranding the run. + heteroPostgresTest( + "forWaitpointCompletion probes the primary, resolving the owner even under replica lag", + async ({ prisma14, prisma17 }) => { + const legacyStore = new PostgresRunStore({ prisma: prisma14, readOnlyPrisma: prisma14 }); + // NEW store writes to prisma17 but its replica is the (empty, w.r.t. this waitpoint) prisma14. + const newStore = new PostgresRunStore({ prisma: prisma17, readOnlyPrisma: prisma14 }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + const seed17 = await seedEnvironment(prisma17, "wpc_lag17"); + + const cuidWaitpointId = `waitpoint_${"c".repeat(25)}`; // id-shape → LEGACY (the wrong owner) + await prisma17.waitpoint.create({ + data: { + id: cuidWaitpointId, + friendlyId: "waitpoint_wpc_lag", + type: "MANUAL", + idempotencyKey: "wpc-lag-key", + userProvidedIdempotencyKey: false, + projectId: seed17.project.id, + environmentId: seed17.environment.id, + }, + }); + + // Only the NEW store's PRIMARY (prisma17) has the row; a replica probe (prisma14) misses it. + expect(await router.forWaitpointCompletion(cuidWaitpointId, { routeKind: "MANUAL" })).toBe( + newStore + ); + } + ); + // A waitpoint must be born on the same DB as its run (cuid → LEGACY, ksuid → NEW) so that // completion and the blocking edge — which already routes by run id — line up. A cuid // waitpoint landing on NEW is the regression that strands a non-opted org's wait forever. diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 66bd4a976d..583cbdbb79 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -26,6 +26,7 @@ import type { TaskRunWithWaitpoint, WaitpointColocationOptions, } from "./types.js"; +import { isReadReplicaClient } from "./readReplicaClient.js"; /** * Run-ops routing substrate for the TaskRun-core method group. Implements {@link RunStore} @@ -47,6 +48,23 @@ export class RoutingRunStore implements RunStore { this.#classify = options.classify ?? ownerEngine; } + // A routing store spans two databases and has no single primary — routed reads resolve the + // OWNING sub-store's primary internally (#ownPrimary). + get primaryReadClient(): ReadClient { + throw new Error( + "RoutingRunStore has no single primary read client; routed reads use the owning sub-store's primary" + ); + } + + // Map a caller-passed read client onto a routed store. The caller's client is bound to the + // control-plane connection — the wrong database for a NEW-resident row — so it is never forwarded + // verbatim. A WRITER/tx signals read-your-writes (the just-written row must beat replica lag), so + // the routed read runs on the owning store's OWN primary. A caller-passed REPLICA (branded) or no + // client keeps the owning store's replica, preserving read scaling. + static #ownPrimary(store: RunStore, client: ReadClient | undefined): ReadClient | undefined { + return client != null && !isReadReplicaClient(client) ? store.primaryReadClient : undefined; + } + // An unclassifiable id is treated as LEGACY (probe the control-plane DB rather than drop a // real run), matching the read-through layer's policy. #classifySafe(id: string): Residency { @@ -58,10 +76,18 @@ export class RoutingRunStore implements RunStore { } // A `findRuns` caller bound to the given store (preserves `this`; the overload set isn't - // assignable to a single call signature, so it's cast through the implementation shape). - #findManyOn(store: RunStore): (args: unknown) => Promise>> { - const fn = store.findRuns as (args: unknown) => Promise>>; - return fn.bind(store); + // assignable to a single call signature, so it's cast through the implementation shape). A + // caller-passed client resolves to the store's own primary (#ownPrimary) on every call. + #findManyOn( + store: RunStore, + client: ReadClient | undefined + ): (args: unknown) => Promise>> { + const fn = store.findRuns as ( + args: unknown, + client?: ReadClient + ) => Promise>>; + const resolved = RoutingRunStore.#ownPrimary(store, client); + return (args: unknown) => fn.call(store, args, resolved); } // Route an existing run-ops id by residency. Throws on an unclassifiable id. @@ -122,17 +148,26 @@ export class RoutingRunStore implements RunStore { // Resolve which store ACTUALLY holds a waitpoint id: drain-on-read can relocate a cuid // waitpoint onto NEW while keeping its id, so probe the id-shape's home then the other. - async #resolveWaitpointStore(id: string | undefined): Promise { + // `onPrimary` probes each store's own primary (read-your-writes callers; a fresh row may not + // be on the replica yet, which would mis-resolve the store). + async #resolveWaitpointStore(id: string | undefined, onPrimary = false): Promise { const home = typeof id === "string" && this.#classifySafe(id) === "NEW" ? this.#new : this.#legacy; if (typeof id !== "string") { return home; } - if (await home.findWaitpoint({ where: { id } })) { + if ( + await home.findWaitpoint({ where: { id } }, onPrimary ? home.primaryReadClient : undefined) + ) { return home; } const other = home === this.#new ? this.#legacy : this.#new; - return (await other.findWaitpoint({ where: { id } })) ? other : home; + return (await other.findWaitpoint( + { where: { id } }, + onPrimary ? other.primaryReadClient : undefined + )) + ? other + : home; } static #waitpointId(clause: unknown): string | undefined { @@ -195,9 +230,9 @@ export class RoutingRunStore implements RunStore { ): Promise { // Pass through only the select/include args; the caller's actual client object is never // forwarded to the routed store (the control-plane writer can't query the NEW DB). But its - // IDENTITY is the read-your-writes signal: a WRITER means the caller just wrote this run and - // needs to beat replica lag, so route to the OWNING store's own primary (writer). A replica / - // nothing keeps the default — the owning store's replica. + // PRESENCE is the read-your-writes signal: a client means the caller just wrote this run and + // needs to beat replica lag, so route to the OWNING store's own primary (writer). Nothing + // keeps the default — the owning store's replica. const args = selectOrIncludeArgs(argsOrClient); const onPrimary = readYourWrites(argsOrClient, _client); const id = idFromWhere(where); @@ -271,18 +306,18 @@ export class RoutingRunStore implements RunStore { skip?: number; cursor?: Prisma.TaskRunWhereUniqueInput; }, - _client?: ReadClient + client?: ReadClient ): Promise { // SPLIT-mode fan-out across NEW + LEGACY. A `findRuns` `where` can span ids of mixed // residency, so we resolve each owning store and merge, preserving orderBy/take/skip. - // The caller's client is intentionally NOT forwarded (it is the control-plane client / - // a primary handle); each store reads its own replica, as the prior delegate did. NEW - // wins on id collisions (the copy->fence migration window) so a half-migrated run is - // never double-reported. - return this.#findRunsRouted(args); + // The caller's client is never forwarded verbatim (it is the control-plane client); its + // presence routes each leg to that store's OWN primary (read-your-writes), else each store + // reads its own replica as before. NEW wins on id collisions (the copy->fence migration + // window) so a half-migrated run is never double-reported. + return this.#findRunsRouted(args, client); } - async #findRunsRouted(args: FindRunsArgs): Promise { + async #findRunsRouted(args: FindRunsArgs, client?: ReadClient): Promise { if (args.cursor) { // No caller paginates findRuns by Prisma cursor in split mode (the runs list // paginates in ClickHouse and hydrates a bounded id set). Merging cursor windows @@ -293,20 +328,24 @@ export class RoutingRunStore implements RunStore { } const idList = idListFromWhere(args.where); - return idList ? this.#findRunsByIdSet(args, idList) : this.#findRunsOpen(args); + return idList ? this.#findRunsByIdSet(args, idList, client) : this.#findRunsOpen(args, client); } // Bounded id-set (the list hydrate + engine sweeps). Query NEW for the whole set first // (it holds ksuid runs); probe LEGACY only for the ids NEW missed that could still live // there (cuid). The two id sets are disjoint by construction, so the merge needs no dedupe. - async #findRunsByIdSet(args: FindRunsArgs, ids: string[]): Promise { + async #findRunsByIdSet( + args: FindRunsArgs, + ids: string[], + client?: ReadClient + ): Promise { const { args: selArgs, addedFields } = ensureProjected(args); // The id set already bounds the per-store result, so never push take/skip down — doing // so would truncate a store's page before the merge knows membership and mis-attribute // rows. take/skip are applied once, globally, in finalizeRows. const fan = { ...selArgs, take: undefined, skip: undefined }; - const findNew = this.#findManyOn(this.#new); - const findLegacy = this.#findManyOn(this.#legacy); + const findNew = this.#findManyOn(this.#new, client); + const findLegacy = this.#findManyOn(this.#legacy, client); const newRows = await findNew(fan); const foundIds = new Set(newRows.map((r) => r.id as string)); @@ -324,11 +363,11 @@ export class RoutingRunStore implements RunStore { // Open predicate (e.g. `{ batchId }`, `{ status, runtimeEnvironmentId }`): no id set to // partition, so query both stores and dedupe by id (NEW wins). - async #findRunsOpen(args: FindRunsArgs): Promise { + async #findRunsOpen(args: FindRunsArgs, client?: ReadClient): Promise { const { args: selArgs, addedFields } = ensureProjected(args); const fan = widenForMerge(selArgs); - const findNew = this.#findManyOn(this.#new); - const findLegacy = this.#findManyOn(this.#legacy); + const findNew = this.#findManyOn(this.#new, client); + const findLegacy = this.#findManyOn(this.#legacy, client); const [newRows, legacyRows] = await Promise.all([findNew(fan), findLegacy(fan)]); const byId = new Map>(); for (const r of legacyRows) byId.set(r.id as string, r); @@ -582,8 +621,8 @@ export class RoutingRunStore implements RunStore { argsOrClient?: { select?: unknown; include?: unknown } | ReadClient, _client?: ReadClient ): Promise { - // The caller's client is not forwarded, but a WRITER signals read-your-writes → the owning - // store's primary (writer); a replica / nothing → its replica (see findRun). + // The caller's client is not forwarded, but its presence signals read-your-writes → the + // owning store's primary (writer); nothing → its replica (see findRun). const args = selectOrIncludeArgs(argsOrClient); const onPrimary = readYourWrites(argsOrClient, _client); const id = idFromWhere(where); @@ -689,11 +728,15 @@ export class RoutingRunStore implements RunStore { include: { completedWaitpoints: true; checkpoint: true }; }> | null> { const owningStore = this.#routeOrNew(runId); - const snapshot = await owningStore.findLatestExecutionSnapshot(runId); + const snapshot = await owningStore.findLatestExecutionSnapshot( + runId, + RoutingRunStore.#ownPrimary(owningStore, client) + ); if (snapshot) { await this.#reresolveCompletedWaitpointsCrossDb( snapshot as Record, - owningStore + owningStore, + client ); } return snapshot; @@ -707,7 +750,8 @@ export class RoutingRunStore implements RunStore { // cross-DB and appended, so a cuid token completing a ksuid run keeps its OUTPUT on the resume. async #reresolveCompletedWaitpointsCrossDb( snapshot: Record, - owningStore: RunStore + owningStore: RunStore, + client?: ReadClient ): Promise { const snapshotId = snapshot.id; if (typeof snapshotId !== "string") { @@ -719,14 +763,18 @@ export class RoutingRunStore implements RunStore { const present = new Set(completed.map((w) => w.id as string)); // The join is co-resident with the snapshot, so read it from the OWNING store (the snapshot's // own id is a cuid and would mis-route the both-DB `findSnapshotCompletedWaitpointIds`). - const joinIds = await owningStore.findSnapshotCompletedWaitpointIds(snapshotId); + const joinIds = await owningStore.findSnapshotCompletedWaitpointIds( + snapshotId, + RoutingRunStore.#ownPrimary(owningStore, client) + ); const missing = joinIds.filter((id) => !present.has(id)); if (missing.length === 0) { return; // all completed tokens co-resident → owning-store hydration is complete } - const recovered = (await this.findManyWaitpoints({ - where: { id: { in: missing } }, - })) as Record[]; + const recovered = (await this.findManyWaitpoints( + { where: { id: { in: missing } } }, + client + )) as Record[]; snapshot.completedWaitpoints = [...completed, ...recovered]; } @@ -741,10 +789,17 @@ export class RoutingRunStore implements RunStore { ): Promise | null> { const runId = snapshotWhereRunId(args); if (runId !== undefined) { - return this.#routeOrNew(runId).findExecutionSnapshot(args); + const store = this.#routeOrNew(runId); + return store.findExecutionSnapshot(args, RoutingRunStore.#ownPrimary(store, client)); } - const fromNew = await this.#new.findExecutionSnapshot(args); - return fromNew ?? this.#legacy.findExecutionSnapshot(args); + const fromNew = await this.#new.findExecutionSnapshot( + args, + RoutingRunStore.#ownPrimary(this.#new, client) + ); + return ( + fromNew ?? + this.#legacy.findExecutionSnapshot(args, RoutingRunStore.#ownPrimary(this.#legacy, client)) + ); } // Snapshot reads route by OWNING run id; merge both DBs for an open/cross-residency where. @@ -754,11 +809,15 @@ export class RoutingRunStore implements RunStore { ): Promise[]> { const runId = snapshotWhereRunId(args); if (runId !== undefined) { - return this.#routeOrNew(runId).findManyExecutionSnapshots(args); + const store = this.#routeOrNew(runId); + return store.findManyExecutionSnapshots(args, RoutingRunStore.#ownPrimary(store, client)); } const [fromNew, fromLegacy] = await Promise.all([ - this.#new.findManyExecutionSnapshots(args), - this.#legacy.findManyExecutionSnapshots(args), + this.#new.findManyExecutionSnapshots(args, RoutingRunStore.#ownPrimary(this.#new, client)), + this.#legacy.findManyExecutionSnapshots( + args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), ]); return [...fromNew, ...fromLegacy]; } @@ -772,7 +831,11 @@ export class RoutingRunStore implements RunStore { // A snapshot lives with its run; route by the snapshot id's residency. findSnapshotCompletedWaitpointIds(snapshotId: string, client?: ReadClient): Promise { - return this.#routeOrNew(snapshotId).findSnapshotCompletedWaitpointIds(snapshotId); + const store = this.#routeOrNew(snapshotId); + return store.findSnapshotCompletedWaitpointIds( + snapshotId, + RoutingRunStore.#ownPrimary(store, client) + ); } // Keyed by waitpointId, but the WaitpointRunConnection / CompletedWaitpoint join co-locates with the @@ -781,8 +844,14 @@ export class RoutingRunStore implements RunStore { // row on each leg. async findWaitpointConnectedRunIds(waitpointId: string, client?: ReadClient): Promise { const [fromNew, fromLegacy] = await Promise.all([ - this.#new.findWaitpointConnectedRunIds(waitpointId), - this.#legacy.findWaitpointConnectedRunIds(waitpointId), + this.#new.findWaitpointConnectedRunIds( + waitpointId, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.findWaitpointConnectedRunIds( + waitpointId, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), ]); return uniqueStrings([...fromNew, ...fromLegacy]); } @@ -792,8 +861,14 @@ export class RoutingRunStore implements RunStore { client?: ReadClient ): Promise { const [fromNew, fromLegacy] = await Promise.all([ - this.#new.findWaitpointCompletedSnapshotIds(waitpointId), - this.#legacy.findWaitpointCompletedSnapshotIds(waitpointId), + this.#new.findWaitpointCompletedSnapshotIds( + waitpointId, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.findWaitpointCompletedSnapshotIds( + waitpointId, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), ]); return uniqueStrings([...fromNew, ...fromLegacy]); } @@ -814,8 +889,14 @@ export class RoutingRunStore implements RunStore { // each and sum rather than assume one home. async countPendingWaitpoints(waitpointIds: string[], client?: ReadClient): Promise { const [fromNew, fromLegacy] = await Promise.all([ - this.#new.countPendingWaitpoints(waitpointIds), - this.#legacy.countPendingWaitpoints(waitpointIds), + this.#new.countPendingWaitpoints( + waitpointIds, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.countPendingWaitpoints( + waitpointIds, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), ]); return fromNew + fromLegacy; } @@ -874,20 +955,26 @@ export class RoutingRunStore implements RunStore { const id = RoutingRunStore.#waitpointId((args as { where?: unknown }).where); const store = id !== undefined - ? await this.#resolveWaitpointStore(id) + ? await this.#resolveWaitpointStore(id, client !== undefined) : opts?.coLocateWithRunId !== undefined ? this.#routeOrNew(opts.coLocateWithRunId) : undefined; const row = store !== undefined - ? ((await store.findWaitpoint(scalarArgs as typeof args)) as Record | null) - : (((await this.#new.findWaitpoint(scalarArgs as typeof args)) ?? - (await this.#legacy.findWaitpoint(scalarArgs as typeof args))) as Record< - string, - unknown - > | null); + ? ((await store.findWaitpoint( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(store, client) + )) as Record | null) + : (((await this.#new.findWaitpoint( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#new, client) + )) ?? + (await this.#legacy.findWaitpoint( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ))) as Record | null); if (row) { - await this.#reresolveWaitpointRelationsCrossDb(row, relations); + await this.#reresolveWaitpointRelationsCrossDb(row, relations, client); } return row as Prisma.WaitpointGetPayload | null; } @@ -899,7 +986,7 @@ export class RoutingRunStore implements RunStore { args: Prisma.SelectSubset ): Promise | null> { const id = RoutingRunStore.#waitpointId((args as { where?: unknown }).where); - const store = id !== undefined ? await this.#resolveWaitpointStore(id) : this.#new; + const store = id !== undefined ? await this.#resolveWaitpointStore(id, true) : this.#new; return store.findWaitpointOnPrimary(args); } @@ -911,8 +998,14 @@ export class RoutingRunStore implements RunStore { args as Record ); const [fromNew, fromLegacy] = await Promise.all([ - this.#new.findManyWaitpoints(scalarArgs as typeof args), - this.#legacy.findManyWaitpoints(scalarArgs as typeof args), + this.#new.findManyWaitpoints( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.findManyWaitpoints( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), ]); // A token mirrored onto both DBs during drain appears in BOTH legs; dedup by id with NEW-wins // (the NEW copy is authoritative once a run migrates), matching the router's NEW-wins invariant @@ -927,7 +1020,11 @@ export class RoutingRunStore implements RunStore { } const rows = [...byId.values(), ...passthrough]; for (const row of rows) { - await this.#reresolveWaitpointRelationsCrossDb(row as Record, relations); + await this.#reresolveWaitpointRelationsCrossDb( + row as Record, + relations, + client + ); } return rows; } @@ -938,7 +1035,8 @@ export class RoutingRunStore implements RunStore { // group-A relation was requested (the byte-identical scalar path). async #reresolveWaitpointRelationsCrossDb( row: Record, - relations: Partial> + relations: Partial>, + client?: ReadClient ): Promise { const waitpointId = row.id; if (typeof waitpointId !== "string") { @@ -947,19 +1045,22 @@ export class RoutingRunStore implements RunStore { if ("blockingTaskRuns" in relations) { row.blockingTaskRuns = await this.#reresolveBlockingTaskRunsCrossDb( waitpointId, - relations.blockingTaskRuns + relations.blockingTaskRuns, + client ); } if ("connectedRuns" in relations) { row.connectedRuns = await this.#reresolveConnectedRunsCrossDb( waitpointId, - relations.connectedRuns + relations.connectedRuns, + client ); } if ("completedExecutionSnapshots" in relations) { row.completedExecutionSnapshots = await this.#reresolveCompletedExecutionSnapshotsCrossDb( waitpointId, - relations.completedExecutionSnapshots + relations.completedExecutionSnapshots, + client ); } } @@ -969,27 +1070,32 @@ export class RoutingRunStore implements RunStore { // with the run, so a single store misses a cross-DB run's edge; the both-DB read recovers it. async #reresolveBlockingTaskRunsCrossDb( waitpointId: string, - projection: SubProjection + projection: SubProjection, + client?: ReadClient ): Promise { const edgeArgs = projectionAsArgs(projection) ?? {}; - return this.findManyTaskRunWaitpoints({ - ...(edgeArgs as Prisma.TaskRunWaitpointFindManyArgs), - where: { waitpointId }, - }); + return this.findManyTaskRunWaitpoints( + { + ...(edgeArgs as Prisma.TaskRunWaitpointFindManyArgs), + where: { waitpointId }, + }, + client + ); } // connectedRuns: the WaitpointRunConnection join co-locates with the run, so read the connected run // ids from EACH store, then resolve the TaskRun rows across BOTH DBs (findRun routes by id). async #reresolveConnectedRunsCrossDb( waitpointId: string, - projection: SubProjection + projection: SubProjection, + client?: ReadClient ): Promise { - const runIds = await this.findWaitpointConnectedRunIds(waitpointId); + const runIds = await this.findWaitpointConnectedRunIds(waitpointId, client); const findRun = (this.findRun as (...rest: unknown[]) => Promise).bind(this); const args = projectionAsArgs(projection); const runs: unknown[] = []; for (const runId of runIds) { - const run = await findRun({ id: runId }, args); + const run = await findRun({ id: runId }, args, client); if (run != null) { runs.push(run); } @@ -1001,17 +1107,21 @@ export class RoutingRunStore implements RunStore { // the snapshot ids from EACH store, then resolve the snapshot rows across BOTH DBs. async #reresolveCompletedExecutionSnapshotsCrossDb( waitpointId: string, - projection: SubProjection + projection: SubProjection, + client?: ReadClient ): Promise { - const snapshotIds = await this.findWaitpointCompletedSnapshotIds(waitpointId); + const snapshotIds = await this.findWaitpointCompletedSnapshotIds(waitpointId, client); if (snapshotIds.length === 0) { return []; } const findArgs = projectionAsArgs(projection) ?? {}; - return this.findManyExecutionSnapshots({ - ...(findArgs as Prisma.TaskRunExecutionSnapshotFindManyArgs), - where: { id: { in: snapshotIds } }, - }); + return this.findManyExecutionSnapshots( + { + ...(findArgs as Prisma.TaskRunExecutionSnapshotFindManyArgs), + where: { id: { in: snapshotIds } }, + }, + client + ); } async updateWaitpoint( @@ -1065,12 +1175,17 @@ export class RoutingRunStore implements RunStore { : this.#legacy; // Resolve to where the waitpoint ACTUALLY lives: a migrated run's waitpoint can be on NEW // with a LEGACY-classified id (or vice versa), so verify and fall back rather than route - // by id-shape alone and miss it (which leaves the blocked run stuck forever). - if (await preferred.findWaitpoint({ where: { id: waitpointId } })) { + // by id-shape alone and miss it (which leaves the blocked run stuck forever). This guard + // selects the store a WRITE (updateManyWaitpoints) then lands on, so it must probe each + // store's PRIMARY (mirroring #resolveWaitpointStore's onPrimary): a just-created waitpoint the + // replica hasn't caught up on would otherwise mis-resolve the owner and strand the run. + if ( + await preferred.findWaitpoint({ where: { id: waitpointId } }, preferred.primaryReadClient) + ) { return preferred; } const other = preferred === this.#new ? this.#legacy : this.#new; - if (await other.findWaitpoint({ where: { id: waitpointId } })) { + if (await other.findWaitpoint({ where: { id: waitpointId } }, other.primaryReadClient)) { return other; } return preferred; @@ -1095,16 +1210,22 @@ export class RoutingRunStore implements RunStore { ); const [fromNew, fromLegacy] = await Promise.all([ - this.#new.findManyTaskRunWaitpoints(scalarArgs as typeof args), - this.#legacy.findManyTaskRunWaitpoints(scalarArgs as typeof args), + this.#new.findManyTaskRunWaitpoints( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#new, client) + ), + this.#legacy.findManyTaskRunWaitpoints( + scalarArgs as typeof args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ), ]); const edges = dedupeEdgesById([...fromNew, ...fromLegacy]) as Record[]; if (waitpoint) { - await this.#hydrateEdgeWaitpointsCrossDb(edges, waitpoint); + await this.#hydrateEdgeWaitpointsCrossDb(edges, waitpoint, client); } if (taskRun) { - await this.#hydrateEdgeTaskRunsCrossDb(edges, taskRun); + await this.#hydrateEdgeTaskRunsCrossDb(edges, taskRun, client); } return edges as Prisma.TaskRunWaitpointGetPayload[]; } @@ -1114,15 +1235,17 @@ export class RoutingRunStore implements RunStore { // run would otherwise hang forever (or be wrongly treated as completed) on a null status. async #hydrateEdgeWaitpointsCrossDb( edges: Record[], - projection: SubProjection + projection: SubProjection, + client?: ReadClient ): Promise { const ids = uniqueStrings(edges.map((e) => e.waitpointId)); if (ids.length === 0) { return; } - const waitpoints = (await this.findManyWaitpoints({ - where: { id: { in: ids } }, - })) as Record[]; + const waitpoints = (await this.findManyWaitpoints( + { where: { id: { in: ids } } }, + client + )) as Record[]; const byId = new Map(waitpoints.map((w) => [w.id as string, w])); for (const edge of edges) { const id = edge.waitpointId as string | undefined; @@ -1143,7 +1266,8 @@ export class RoutingRunStore implements RunStore { // blocked-run resume path keys off `waitpoint`). async #hydrateEdgeTaskRunsCrossDb( edges: Record[], - projection: SubProjection + projection: SubProjection, + client?: ReadClient ): Promise { // Bind to `this`: findRun reaches the private #routeOrNew/#findRunUnrouted members, so an unbound // reference loses `this` and throws on the first private access. @@ -1151,7 +1275,7 @@ export class RoutingRunStore implements RunStore { const args = projectionAsArgs(projection); for (const edge of edges) { const id = edge.taskRunId as string | undefined; - const run = id ? await findRun({ id }, args) : null; + const run = id ? await findRun({ id }, args, client) : null; edge.taskRun = applyEdgeProjection((run as Record) ?? null, projection); } } @@ -1182,21 +1306,27 @@ export class RoutingRunStore implements RunStore { const runId = whereFieldString(args.where?.taskRunId as Prisma.TaskRunWhereInput["id"]); if (runId !== undefined) { // Residency-classifiable run id present: route to the owning store. Never forward the - // caller's client (it is the control-plane handle); the owning store reads its OWN DB. - return this.#routeOrNew(runId).findTaskRunAttempt(args); + // caller's client verbatim (it is the control-plane handle); its presence resolves to the + // owning store's OWN primary. + const store = this.#routeOrNew(runId); + return store.findTaskRunAttempt(args, RoutingRunStore.#ownPrimary(store, client)); } // No classifiable run id (no taskRunId, or complex filter): fan out NEW-first → LEGACY. - return this.#findTaskRunAttemptUnrouted(args); + return this.#findTaskRunAttemptUnrouted(args, client); } async #findTaskRunAttemptUnrouted( - args: Prisma.SelectSubset + args: Prisma.SelectSubset, + client?: ReadClient ): Promise | null> { - const fromNew = await this.#new.findTaskRunAttempt(args); + const fromNew = await this.#new.findTaskRunAttempt( + args, + RoutingRunStore.#ownPrimary(this.#new, client) + ); if (fromNew != null) { return fromNew; } - return this.#legacy.findTaskRunAttempt(args); + return this.#legacy.findTaskRunAttempt(args, RoutingRunStore.#ownPrimary(this.#legacy, client)); } // Co-locate the checkpoint with its OWNING run so the run-routed snapshot's `checkpointId` FK @@ -1250,11 +1380,19 @@ export class RoutingRunStore implements RunStore { args?: { include?: T }, client?: ReadClient ): Promise | null> { - // Never forward the caller's client (it is the control-plane handle): a cross-DB probe with one - // shared client can only reach one DB, so each sub-store must read its OWN DB (5434 vs 5432). - const fromNew = await this.#new.findBatchTaskRunById(id, args); + // Never forward the caller's client verbatim (a cross-DB probe with one shared client can + // only reach one DB); its presence resolves each leg to that store's OWN primary. + const fromNew = await this.#new.findBatchTaskRunById( + id, + args, + RoutingRunStore.#ownPrimary(this.#new, client) + ); if (fromNew != null) return fromNew; - return this.#legacy.findBatchTaskRunById(id, args); + return this.#legacy.findBatchTaskRunById( + id, + args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ); } // Env-scoped friendlyId probe; no id-routing because cuid-on-NEW window batches exist. @@ -1264,10 +1402,21 @@ export class RoutingRunStore implements RunStore { args?: { include?: T }, client?: ReadClient ): Promise | null> { - // Never forward the caller's client (control-plane handle): each sub-store reads its OWN DB. - const fromNew = await this.#new.findBatchTaskRunByFriendlyId(friendlyId, environmentId, args); + // Never forward the caller's client verbatim; its presence resolves each leg to that + // store's OWN primary. + const fromNew = await this.#new.findBatchTaskRunByFriendlyId( + friendlyId, + environmentId, + args, + RoutingRunStore.#ownPrimary(this.#new, client) + ); if (fromNew != null) return fromNew; - return this.#legacy.findBatchTaskRunByFriendlyId(friendlyId, environmentId, args); + return this.#legacy.findBatchTaskRunByFriendlyId( + friendlyId, + environmentId, + args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ); } // --------------------------------------------------------------------------- @@ -1283,14 +1432,21 @@ export class RoutingRunStore implements RunStore { args?: { include?: T }, client?: ReadClient ): Promise | null> { - // Never forward the caller's client (control-plane handle): each sub-store reads its OWN DB. + // Never forward the caller's client verbatim; its presence resolves each leg to that + // store's OWN primary. const fromNew = await this.#new.findBatchTaskRunByIdempotencyKey( environmentId, idempotencyKey, - args + args, + RoutingRunStore.#ownPrimary(this.#new, client) ); if (fromNew != null) return fromNew; - return this.#legacy.findBatchTaskRunByIdempotencyKey(environmentId, idempotencyKey, args); + return this.#legacy.findBatchTaskRunByIdempotencyKey( + environmentId, + idempotencyKey, + args, + RoutingRunStore.#ownPrimary(this.#legacy, client) + ); } // Route by `where.id` when scalar; else (e.g. status filter) fan out to both and sum. @@ -1315,10 +1471,11 @@ export class RoutingRunStore implements RunStore { where: { batchTaskRunId: string; status?: BatchTaskRunItemStatus }, client?: ReadClient ): Promise { - // Never forward the caller's client (it is the control-plane handle): a ksuid batch routes to - // NEW, so a forwarded control-plane client would count items on the wrong DB (→ 0/wrong count). - // The routed store reads its OWN DB. Mirrors the probe-read sweep. - return this.#routeOrNew(where.batchTaskRunId).countBatchTaskRunItems(where); + // Never forward the caller's client verbatim (a ksuid batch routes to NEW, so a forwarded + // control-plane client would count items on the wrong DB → 0/wrong count); its presence + // resolves to the owning store's OWN primary. + const store = this.#routeOrNew(where.batchTaskRunId); + return store.countBatchTaskRunItems(where, RoutingRunStore.#ownPrimary(store, client)); } // Route by item `id` or `batchTaskRunId` when scalar; else fan out to both and sum. @@ -1373,25 +1530,18 @@ function selectOrIncludeArgs( return undefined; } -// Writers (PrismaClient / RunOpsPrismaClient) expose `$transaction`; a read replica -// (PrismaReplicaClient) does not. That is the identity the routing store keys read-your-writes on. -function isWriterClient(value: unknown): boolean { - return ( - !!value && - typeof value === "object" && - typeof (value as { $transaction?: unknown }).$transaction === "function" - ); -} - -// A read-your-writes call passes a WRITER (the just-written run must be read back before the -// replica has it). Recover the caller's client from the overloaded read args — slot two when it -// isn't a `{ select | include }` object, else slot three — and report whether it is a writer. +// A read-your-writes call passes a WRITER or ambient tx, whose just-written row must be read back +// before the replica has it. Recover the caller's client from the overloaded read args — slot two +// when it isn't a `{ select | include }` object, else slot three — and report whether it warrants +// escalation to the owning primary. A branded replica does NOT: it can't be forwarded across DBs, +// but it signals a replica-intended read, so the owning store keeps its own replica (read scaling). function readYourWrites( argsOrClient: { select?: unknown; include?: unknown } | ReadClient | unknown, client: ReadClient | undefined ): boolean { - const passedClient = selectOrIncludeArgs(argsOrClient) === undefined ? argsOrClient : client; - return isWriterClient(passedClient); + const passedClient = + selectOrIncludeArgs(argsOrClient) === undefined ? (argsOrClient ?? client) : client; + return passedClient != null && !isReadReplicaClient(passedClient); } // Read a plain scalar string field off a create-data object (e.g. `data.completedByTaskRunId`). diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 32e5644114..ce5e141de0 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -485,6 +485,13 @@ export interface RunStore { ): Promise; // Read + + // This store's own PRIMARY (writer) handle in read-client form. The routing layer passes it as + // the `client` for a routed read when the CALLER supplied one: the caller's client is bound to + // the control-plane DB (the wrong database for a NEW-resident row), so read-your-writes is + // honored by reading the OWNING store's own primary instead of its replica. + readonly primaryReadClient: ReadClient; + findRun( where: Prisma.TaskRunWhereInput, args: { select: S },