refactor(run-engine,webapp): route TaskRun writes through a new RunStore adapter (#3981)

Author: d-cs

apps/webapp/app/routes/api.v1.runs.$runId.tags.ts

@@ -9,6 +9,7 @@ import { getRequestAbortSignal } from "~/services/httpAsyncStorage.server";
 import { logger } from "~/services/logger.server";
 import { publishChangeRecord } from "~/services/realtime/runChangeNotifierInstance.server";
 import { mutateWithFallback } from "~/v3/mollifier/mutateWithFallback.server";
+import { runStore } from "~/v3/runStore.server";
 
 // Pull the existing tags out of a buffer entry's serialised payload so
 // the buffer-path response can dedup against them, matching the
@@ -84,14 +85,7 @@ export async function action({ request, params }: ActionFunctionArgs) {
         if (newTags.length === 0) {
           return json({ message: "No new tags to add" }, { status: 200 });
         }
-        const updated = await prisma.taskRun.update({
-          where: {
-            id: taskRun.id,
-            runtimeEnvironmentId: env.id,
-          },
-          data: { runTags: { push: newTags } },
-          select: { updatedAt: true },
-        });
+        const updated = await runStore.pushTags(taskRun.id, newTags, { runtimeEnvironmentId: env.id }, prisma);
         // Publish a run-changed record with the NEW tag set so tag feeds reindex
         // (no-op unless enabled). updatedAt is the read-your-writes watermark.
         publishChangeRecord({

apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.append.ts

@@ -6,6 +6,7 @@ import { $replica, prisma } from "~/db.server";
 import { getRealtimeStreamInstance } from "~/services/realtime/v1StreamsGlobal.server";
 import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server";
 import { ServiceValidationError } from "~/v3/services/common.server";
+import { runStore } from "~/v3/runStore.server";
 
 const ParamsSchema = z.object({
   runId: z.string(),
@@ -87,16 +88,7 @@ const { action } = createActionApiRoute(
     }
 
     if (!targetRun.realtimeStreams.includes(params.streamId)) {
-      await prisma.taskRun.update({
-        where: {
-          id: targetRun.id,
-        },
-        data: {
-          realtimeStreams: {
-            push: params.streamId,
-          },
-        },
-      });
+      await runStore.pushRealtimeStream(targetRun.id, params.streamId, prisma);
     }
 
     const part = await request.text();

apps/webapp/app/routes/realtime.v1.streams.$runId.$target.$streamId.ts

@@ -6,6 +6,7 @@ import {
   createActionApiRoute,
   createLoaderApiRoute,
 } from "~/services/routeBuilders/apiBuilder.server";
+import { runStore } from "~/v3/runStore.server";
 
 const ParamsSchema = z.object({
   runId: z.string(),
@@ -86,12 +87,7 @@ const { action } = createActionApiRoute(
       }
 
       if (!target.realtimeStreams.includes(params.streamId)) {
-        await prisma.taskRun.update({
-          where: { id: target.id },
-          data: {
-            realtimeStreams: { push: params.streamId },
-          },
-        });
+        await runStore.pushRealtimeStream(target.id, params.streamId, prisma);
       }
 
       const realtimeStream = getRealtimeStreamInstance(

apps/webapp/app/runEngine/concerns/idempotencyKeys.server.ts

@@ -10,6 +10,7 @@ import { getMollifierBuffer } from "~/v3/mollifier/mollifierBuffer.server";
 import { findRunByIdWithMollifierFallback } from "~/v3/mollifier/readFallback.server";
 import { claimOrAwait } from "~/v3/mollifier/idempotencyClaim.server";
 import { makeResolveMollifierFlag } from "~/v3/mollifier/mollifierGate.server";
+import { runStore } from "~/v3/runStore.server";
 import type { TraceEventConcern, TriggerTaskRequest } from "../types";
 
 // In-memory per-org mollifier-enabled check, shared with `evaluateGate`
@@ -190,10 +191,10 @@ export class IdempotencyKeyConcern {
         });
 
         // Update the existing run to remove the idempotency key
-        await this.prisma.taskRun.updateMany({
-          where: { id: existingRun.id, idempotencyKey },
-          data: { idempotencyKey: null, idempotencyKeyExpiresAt: null },
-        });
+        await runStore.clearIdempotencyKey(
+          { byId: { runId: existingRun.id, idempotencyKey } },
+          this.prisma
+        );
 
         return { isCached: false, idempotencyKey, idempotencyKeyExpiresAt };
       }
@@ -207,10 +208,10 @@ export class IdempotencyKeyConcern {
         });
 
         // Update the existing run to remove the idempotency key
-        await this.prisma.taskRun.updateMany({
-          where: { id: existingRun.id, idempotencyKey },
-          data: { idempotencyKey: null, idempotencyKeyExpiresAt: null },
-        });
+        await runStore.clearIdempotencyKey(
+          { byId: { runId: existingRun.id, idempotencyKey } },
+          this.prisma
+        );
 
         return { isCached: false, idempotencyKey, idempotencyKeyExpiresAt };
       }

apps/webapp/app/services/metadata/updateMetadata.server.ts

@@ -13,6 +13,7 @@ import { Effect, Schedule, Duration, Fiber } from "effect";
 import { type RuntimeFiber } from "effect/Fiber";
 import { setTimeout } from "timers/promises";
 import { Logger, LogLevel } from "@trigger.dev/core/logger";
+import type { RunStore } from "@internal/run-store";
 
 const RUN_UPDATABLE_WINDOW_MS = 60 * 60 * 1000; // 1 hour
 
@@ -24,6 +25,7 @@ type BufferedRunMetadataChangeOperation = {
 
 export type UpdateMetadataServiceOptions = {
   prisma: PrismaClientOrTransaction;
+  runStore: RunStore;
   flushIntervalMs?: number;
   flushEnabled?: boolean;
   flushLoggingEnabled?: boolean;
@@ -49,6 +51,7 @@ export class UpdateMetadataService {
   private _bufferedOperations: Map<string, BufferedRunMetadataChangeOperation[]> = new Map();
   private _flushFiber: RuntimeFiber<void> | null = null;
   private readonly _prisma: PrismaClientOrTransaction;
+  private readonly _runStore: RunStore;
   private readonly flushIntervalMs: number;
   private readonly flushEnabled: boolean;
   private readonly flushLoggingEnabled: boolean;
@@ -57,6 +60,7 @@ export class UpdateMetadataService {
 
   constructor(private readonly options: UpdateMetadataServiceOptions) {
     this._prisma = options.prisma;
+    this._runStore = options.runStore;
     this.flushIntervalMs = options.flushIntervalMs ?? 5000;
     this.flushEnabled = options.flushEnabled ?? true;
     this.flushLoggingEnabled = options.flushLoggingEnabled ?? true;
@@ -260,17 +264,16 @@ export class UpdateMetadataService {
       const writeTime = new Date();
       const result = yield* _(
         Effect.tryPromise(() =>
-          this._prisma.taskRun.updateMany({
-            where: {
-              id: runId,
-              metadataVersion: run.metadataVersion,
-            },
-            data: {
-              metadata: newMetadataPacket.data,
+          this._runStore.updateMetadata(
+            runId,
+            {
+              metadata: newMetadataPacket.data!,
               metadataVersion: { increment: 1 },
               updatedAt: writeTime,
             },
-          })
+            { expectedMetadataVersion: run.metadataVersion },
+            this._prisma
+          )
         )
       );
 
@@ -469,20 +472,19 @@ export class UpdateMetadataService {
       // Update with optimistic locking; updatedAt stamped explicitly so the caller can
       // publish the exact committed watermark without a follow-up read.
       const writeTime = new Date();
-      const result = await this._prisma.taskRun.updateMany({
-        where: {
-          id: runId,
-          metadataVersion: run.metadataVersion,
-        },
-        data: {
-          metadata: newMetadataPacket.data,
+      const result = await this._runStore.updateMetadata(
+        runId,
+        {
+          metadata: newMetadataPacket.data!,
           metadataType: newMetadataPacket.dataType,
           metadataVersion: {
             increment: 1,
           },
           updatedAt: writeTime,
         },
-      });
+        { expectedMetadataVersion: run.metadataVersion },
+        this._prisma
+      );
 
       if (result.count === 0) {
         if (this.flushLoggingEnabled) {
@@ -564,19 +566,19 @@ export class UpdateMetadataService {
       // Update the metadata without version check; updatedAt stamped explicitly so the
       // caller can publish the exact committed watermark.
       const writeTime = new Date();
-      await this._prisma.taskRun.update({
-        where: {
-          id: runId,
-        },
-        data: {
-          metadata: metadataPacket?.data,
+      await this._runStore.updateMetadata(
+        runId,
+        {
+          metadata: metadataPacket?.data!,
           metadataType: metadataPacket?.dataType,
           metadataVersion: {
             increment: 1,
           },
           updatedAt: writeTime,
         },
-      });
+        {},
+        this._prisma
+      );
       updatedAtMs = writeTime.getTime();
     }
 

apps/webapp/app/services/metadata/updateMetadataInstance.server.ts

@@ -2,13 +2,15 @@ import { singleton } from "~/utils/singleton";
 import { env } from "~/env.server";
 import { UpdateMetadataService } from "./updateMetadata.server";
 import { prisma } from "~/db.server";
+import { runStore } from "~/v3/runStore.server";
 import { publishChangeRecord } from "~/services/realtime/runChangeNotifierInstance.server";
 
 export const updateMetadataService = singleton(
   "update-metadata-service",
   () =>
     new UpdateMetadataService({
       prisma,
+      runStore,
       flushIntervalMs: env.BATCH_METADATA_OPERATIONS_FLUSH_INTERVAL_MS,
       flushEnabled: env.BATCH_METADATA_OPERATIONS_FLUSH_ENABLED === "1",
       flushLoggingEnabled: env.BATCH_METADATA_OPERATIONS_FLUSH_LOGGING_ENABLED === "1",

apps/webapp/app/v3/runStore.server.ts

@@ -0,0 +1,8 @@
+import { PostgresRunStore } from "@internal/run-store";
+import { $replica, prisma } from "~/db.server";
+import { singleton } from "~/utils/singleton";
+
+export const runStore = singleton(
+  "PostgresRunStore",
+  () => new PostgresRunStore({ prisma, readOnlyPrisma: $replica })
+);

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

@@ -1,16 +1,19 @@
 import { Span, SpanKind } from "@opentelemetry/api";
+import type { RunStore } from "@internal/run-store";
 import { $replica, PrismaClientOrTransaction, prisma } from "~/db.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { attributesFromAuthenticatedEnv, tracer } from "../tracer.server";
 import { engine, RunEngine } from "../runEngine.server";
+import { runStore as defaultRunStore } from "../runStore.server";
 import { ServiceValidationError } from "./common.server";
 
 export { ServiceValidationError };
 
 export abstract class BaseService {
   constructor(
     protected readonly _prisma: PrismaClientOrTransaction = prisma,
-    protected readonly _replica: PrismaClientOrTransaction = $replica
+    protected readonly _replica: PrismaClientOrTransaction = $replica,
+    protected readonly runStore: RunStore = defaultRunStore
   ) {}
 
   protected async traceWithEnv<T>(

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

@@ -408,10 +408,10 @@ export class BatchTriggerV3Service extends BaseService {
 
     // Expire the cached runs that are no longer valid
     if (expiredRunIds.size) {
-      await this._prisma.taskRun.updateMany({
-        where: { friendlyId: { in: Array.from(expiredRunIds) } },
-        data: { idempotencyKey: null },
-      });
+      await this.runStore.clearIdempotencyKey(
+        { byFriendlyIds: Array.from(expiredRunIds) },
+        this._prisma
+      );
     }
 
     return runs;

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

@@ -17,15 +17,14 @@ export class RescheduleTaskRunService extends BaseService {
       throw new ServiceValidationError(`Invalid delay: ${body.delay}`);
     }
 
-    const updatedRun = await this._prisma.taskRun.update({
-      where: {
-        id: taskRun.id,
-      },
-      data: {
+    const updatedRun = await this.runStore.rescheduleRun(
+      taskRun.id,
+      {
         delayUntil: delay,
         queueTimestamp: delay,
       },
-    });
+      this._prisma
+    );
 
     if (updatedRun.engine === "V1") {
       await EnqueueDelayedRunService.reschedule(taskRun.id, delay);