Skip to content

Commit 2b67a2f

Browse files
committed
fix(run-store,webapp): route env-scoped waitpoint tags and idempotency-key resets to the run-ops DB when the env mints run-ops ids
Two env-scoped writes with no owning run to route by defaulted to the legacy database in a fully-minted-new environment: waitpoint tags (createWaitpointTag minted no id, so id-shape sent them to legacy) and idempotency-key resets (clearIdempotencyKey by predicate fanned out to both databases). Both now read the env mint kind and pin to the run-ops database, so they no longer write to the draining legacy database; a predicate reset with no mint signal still fans out.
1 parent c6e3ecd commit 2b67a2f

7 files changed

Lines changed: 152 additions & 22 deletions

File tree

apps/webapp/app/models/waitpointTag.server.ts

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,22 +8,30 @@ export async function createWaitpointTag({
88
tag,
99
environmentId,
1010
projectId,
11+
residency,
1112
}: {
1213
tag: string;
1314
environmentId: string;
1415
projectId: string;
16+
// Residency from the env mint kind: a tag has no owning run, so a minted-new env pins it to NEW
17+
// instead of defaulting to the draining legacy DB.
18+
residency?: "NEW" | "LEGACY";
1519
}) {
1620
if (tag.trim().length === 0) return;
1721

1822
let attempts = 0;
1923

2024
while (attempts < MAX_RETRIES) {
2125
try {
22-
return await runStore.upsertWaitpointTag({
23-
environmentId,
24-
name: tag,
25-
projectId,
26-
});
26+
return await runStore.upsertWaitpointTag(
27+
{
28+
environmentId,
29+
name: tag,
30+
projectId,
31+
},
32+
undefined,
33+
residency
34+
);
2735
} catch (error) {
2836
if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") {
2937
// Handle unique constraint violation (conflict)

apps/webapp/app/routes/api.v1.waitpoints.tokens.ts

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,16 @@ const { action } = createActionApiRoute(
5959

6060
const timeout = await parseDelay(body.timeout);
6161

62+
// A token (and its tags) has no owning run, so it can't co-locate. Resolve the env mint kind so a
63+
// minted-new env creates them on the run-ops DB (NEW) instead of defaulting to the draining LEGACY
64+
// DB by their cuid id-shape.
65+
const mintKind = await resolveRunIdMintKind({
66+
organizationId: authentication.environment.organizationId,
67+
id: authentication.environment.id,
68+
orgFeatureFlags: authentication.environment.organization.featureFlags,
69+
});
70+
const residency = mintKind === "runOpsId" ? "NEW" : "LEGACY";
71+
6272
//upsert tags
6373
let tags: { id: string; name: string }[] = [];
6474
const bodyTags = typeof body.tags === "string" ? [body.tags] : body.tags;
@@ -75,30 +85,22 @@ const { action } = createActionApiRoute(
7585
tag,
7686
environmentId: authentication.environment.id,
7787
projectId: authentication.environment.projectId,
88+
residency,
7889
});
7990
if (tagRecord) {
8091
tags.push(tagRecord);
8192
}
8293
}
8394
}
8495

85-
// A token has no owning run, so it can't co-locate. Resolve the env mint kind so a minted-new
86-
// env creates the token on the run-ops DB (NEW) instead of defaulting to the draining LEGACY DB
87-
// by its cuid id-shape.
88-
const mintKind = await resolveRunIdMintKind({
89-
organizationId: authentication.environment.organizationId,
90-
id: authentication.environment.id,
91-
orgFeatureFlags: authentication.environment.organization.featureFlags,
92-
});
93-
9496
const result = await engine.createManualWaitpoint({
9597
environmentId: authentication.environment.id,
9698
projectId: authentication.environment.projectId,
9799
idempotencyKey: body.idempotencyKey,
98100
idempotencyKeyExpiresAt,
99101
timeout,
100102
tags: bodyTags,
101-
standaloneResidency: mintKind === "runOpsId" ? "NEW" : "LEGACY",
103+
standaloneResidency: residency,
102104
});
103105

104106
const $responseHeaders = await responseHeaders(authentication.environment);

apps/webapp/app/v3/services/resetIdempotencyKey.server.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,30 @@ import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
22
import { BaseService, ServiceValidationError } from "./baseService.server";
33
import { logger } from "~/services/logger.server";
44
import { getMollifierBuffer } from "~/v3/mollifier/mollifierBuffer.server";
5+
import { resolveRunIdMintKind } from "~/v3/engineVersion.server";
56

67
export class ResetIdempotencyKeyService extends BaseService {
78
public async call(
89
idempotencyKey: string,
910
taskIdentifier: string,
1011
authenticatedEnv: AuthenticatedEnvironment
1112
): Promise<{ id: string }> {
13+
// The predicate has no run id to route by. When the env mints run-ops ids its runs live on NEW,
14+
// so pin the reset to NEW and skip the wrong-DB (0-row) write to the draining legacy DB.
15+
const mintKind = await resolveRunIdMintKind({
16+
organizationId: authenticatedEnv.organizationId,
17+
id: authenticatedEnv.id,
18+
orgFeatureFlags: authenticatedEnv.organization.featureFlags,
19+
});
20+
const residency = mintKind === "runOpsId" ? "NEW" : "LEGACY";
21+
1222
const { count: pgCount } = await this.runStore.clearIdempotencyKey(
1323
{
1424
byPredicate: {
1525
idempotencyKey,
1626
taskIdentifier,
1727
runtimeEnvironmentId: authenticatedEnv.id,
28+
residency,
1829
},
1930
},
2031
this._prisma
@@ -80,6 +91,7 @@ export class ResetIdempotencyKeyService extends BaseService {
8091
idempotencyKey,
8192
taskIdentifier,
8293
runtimeEnvironmentId: authenticatedEnv.id,
94+
residency,
8395
},
8496
},
8597
this._prisma

internal-packages/run-store/src/PostgresRunStore.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2535,7 +2535,9 @@ export class PostgresRunStore implements RunStore {
25352535

25362536
async upsertWaitpointTag(
25372537
data: { environmentId: string; name: string; projectId: string; id?: string },
2538-
tx?: PrismaClientOrTransaction
2538+
tx?: PrismaClientOrTransaction,
2539+
// `residency` selects the store at the router; a single store has one client and ignores it.
2540+
_residency?: "NEW" | "LEGACY"
25392541
): Promise<WaitpointTag> {
25402542
const prisma = tx ?? this.prisma;
25412543

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
import { describe, expect, it } from "vitest";
2+
import { RoutingRunStore } from "./runOpsStore.js";
3+
import type { RunStore } from "./types.js";
4+
5+
// Env-scoped writes with no owning run (waitpoint tags; idempotency-key reset by predicate) must
6+
// route to NEW when the env mints run-ops ids, instead of defaulting to LEGACY / fanning a wrong-DB
7+
// write. Pure routing: fake RunStore slots record which store the router dispatches to.
8+
9+
type Call = { method: string; args: unknown[] };
10+
type FakeStore = RunStore & { slot: "new" | "legacy"; calls: Call[] };
11+
12+
function fakeStore(slot: "new" | "legacy"): FakeStore {
13+
const calls: Call[] = [];
14+
const rec =
15+
(method: string, result: unknown) =>
16+
(...args: unknown[]) => {
17+
calls.push({ method, args });
18+
return Promise.resolve(result);
19+
};
20+
return {
21+
slot,
22+
calls,
23+
upsertWaitpointTag: rec("upsertWaitpointTag", { id: slot, slot }),
24+
clearIdempotencyKey: rec("clearIdempotencyKey", { count: slot === "new" ? 1 : 0 }),
25+
} as unknown as FakeStore;
26+
}
27+
28+
function buildRouter() {
29+
const newStore = fakeStore("new");
30+
const legacyStore = fakeStore("legacy");
31+
const router = new RoutingRunStore({
32+
new: newStore,
33+
legacy: legacyStore,
34+
classify: (id: string) => (id.startsWith("new") ? "NEW" : "LEGACY"),
35+
});
36+
return { router, newStore, legacyStore };
37+
}
38+
39+
describe("RoutingRunStore.upsertWaitpointTag — residency hint for a tag with no minted id", () => {
40+
it("routes to NEW when residency is NEW and no id is supplied", async () => {
41+
const { router, newStore, legacyStore } = buildRouter();
42+
await router.upsertWaitpointTag({ environmentId: "env", name: "t", projectId: "p" }, undefined, "NEW");
43+
expect(newStore.calls.map((c) => c.method)).toEqual(["upsertWaitpointTag"]);
44+
expect(legacyStore.calls).toHaveLength(0);
45+
});
46+
47+
it("still falls back to LEGACY when no id and no residency are supplied", async () => {
48+
const { router, newStore, legacyStore } = buildRouter();
49+
await router.upsertWaitpointTag({ environmentId: "env", name: "t", projectId: "p" });
50+
expect(legacyStore.calls.map((c) => c.method)).toEqual(["upsertWaitpointTag"]);
51+
expect(newStore.calls).toHaveLength(0);
52+
});
53+
});
54+
55+
describe("RoutingRunStore.clearIdempotencyKey — predicate routes NEW when the env mints new", () => {
56+
it("routes a byPredicate reset to NEW only when residency is NEW (no legacy fan-out)", async () => {
57+
const { router, newStore, legacyStore } = buildRouter();
58+
const result = await router.clearIdempotencyKey({
59+
byPredicate: {
60+
idempotencyKey: "k",
61+
taskIdentifier: "task",
62+
runtimeEnvironmentId: "env",
63+
residency: "NEW",
64+
},
65+
});
66+
expect(newStore.calls.map((c) => c.method)).toEqual(["clearIdempotencyKey"]);
67+
expect(legacyStore.calls).toHaveLength(0);
68+
expect(result.count).toBe(1);
69+
});
70+
71+
it("still fans out a byPredicate reset with no residency (mixed residency)", async () => {
72+
const { router, newStore, legacyStore } = buildRouter();
73+
await router.clearIdempotencyKey({
74+
byPredicate: { idempotencyKey: "k", taskIdentifier: "task", runtimeEnvironmentId: "env" },
75+
});
76+
expect(newStore.calls).toHaveLength(1);
77+
expect(legacyStore.calls).toHaveLength(1);
78+
});
79+
80+
it("routes byId to the owning store (unchanged)", async () => {
81+
const { router, newStore, legacyStore } = buildRouter();
82+
await router.clearIdempotencyKey({ byId: { runId: "new_run", idempotencyKey: "k" } });
83+
expect(newStore.calls.map((c) => c.method)).toEqual(["clearIdempotencyKey"]);
84+
expect(legacyStore.calls).toHaveLength(0);
85+
});
86+
});

internal-packages/run-store/src/runOpsStore.ts

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -513,7 +513,12 @@ export class RoutingRunStore implements RunStore {
513513
const store = this.#route(params.byId.runId);
514514
return store.clearIdempotencyKey(params, undefined);
515515
}
516-
// `byFriendlyIds` / `byPredicate` can span mixed residency — fan out and sum.
516+
// A `byPredicate` whose env mints run-ops ids has its matching runs on NEW, so route to NEW only
517+
// and skip the wrong-DB (0-row) write to the draining legacy DB. Without that hint (or byFriendlyIds)
518+
// the predicate can span mixed residency — fan out and sum.
519+
if ("byPredicate" in params && params.byPredicate?.residency === "NEW") {
520+
return this.#new.clearIdempotencyKey(params, undefined);
521+
}
517522
return Promise.all([
518523
this.#new.clearIdempotencyKey(params),
519524
this.#legacy.clearIdempotencyKey(params),
@@ -1874,10 +1879,13 @@ export class RoutingRunStore implements RunStore {
18741879
// residency-aware, findManyWaitpointTags must de-dupe by (environmentId, name) or names will duplicate.
18751880
upsertWaitpointTag(
18761881
data: { environmentId: string; name: string; projectId: string; id?: string },
1877-
tx?: PrismaClientOrTransaction
1882+
tx?: PrismaClientOrTransaction,
1883+
residency?: Residency
18781884
): Promise<WaitpointTag> {
1879-
const { store, tx: routedTx } = this.#routeWaitpointWrite(data.id, tx);
1880-
return store.upsertWaitpointTag(data, routedTx);
1885+
// No owning run; route by a minted id-shape when present, else the env's residency hint, else
1886+
// fall back to LEGACY (same precedence as a standalone waitpoint). Caller tx is never forwarded.
1887+
const store = this.#waitpointWriteStore(undefined, residency, data.id);
1888+
return store.upsertWaitpointTag(data, undefined);
18811889
}
18821890

18831891
// A tag keyed by (environmentId, name) can exist on BOTH DBs for one env (dual-resident, no

internal-packages/run-store/src/types.ts

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,15 @@ export type FinalizeRunData = {
257257
export type ClearIdempotencyKeyInput =
258258
| { byId: { runId: string; idempotencyKey: string }; byPredicate?: never; byFriendlyIds?: never }
259259
| {
260-
byPredicate: { idempotencyKey: string; taskIdentifier: string; runtimeEnvironmentId: string };
260+
byPredicate: {
261+
idempotencyKey: string;
262+
taskIdentifier: string;
263+
runtimeEnvironmentId: string;
264+
// A predicate has no run id to route by, so it fans out to both stores. When the env mints
265+
// run-ops ids its matching runs live on NEW, so `residency: "NEW"` routes to NEW only and
266+
// avoids a wrong-DB (0-row) write to the draining legacy DB. Omit to fan out (mixed residency).
267+
residency?: Residency;
268+
};
261269
byId?: never;
262270
byFriendlyIds?: never;
263271
}
@@ -874,7 +882,11 @@ export interface RunStore {
874882
// de-dupes by id in case tag ids ever become residency-aware.
875883
upsertWaitpointTag(
876884
data: { environmentId: string; name: string; projectId: string; id?: string },
877-
tx?: PrismaClientOrTransaction
885+
tx?: PrismaClientOrTransaction,
886+
// A tag has no owning run to co-locate with; when no minted `id` pins it by id-shape, a
887+
// minted-new env's tags read this residency (NEW) so they land with the env's tokens/runs
888+
// instead of defaulting to LEGACY. Single-store impls ignore it.
889+
residency?: Residency
878890
): Promise<WaitpointTag>;
879891
findManyWaitpointTags(
880892
args: {

0 commit comments

Comments
 (0)