v3: cancel subtasks when parent task runs are cancelled

Author: ericallam

apps/webapp/app/components/runs/v3/CancelRunDialog.tsx

@@ -23,8 +23,7 @@ export function CancelRunDialog({ runFriendlyId, redirectPath }: CancelRunDialog
     <DialogContent key="cancel">
       <DialogHeader>Cancel this run?</DialogHeader>
       <DialogDescription>
-        Canceling a run will stop execution. If you want to run this later you will have to replay
-        the entire run with the original payload.
+        Canceling a run will stop execution, along with any executing subtasks.
       </DialogDescription>
       <DialogFooter>
         <Form action={`/resources/taskruns/${runFriendlyId}/cancel`} method="post">

apps/webapp/app/entry.server.tsx

@@ -185,6 +185,7 @@ export { apiRateLimiter } from "./services/apiRateLimit.server";
 export { socketIo } from "./v3/handleSocketIo.server";
 export { wss } from "./v3/handleWebsockets.server";
 export { registryProxy } from "./v3/registryProxy.server";
+export { runWithHttpContext } from "./services/httpAsyncStorage.server";
 import { eventLoopMonitor } from "./eventLoopMonitor.server";
 import { env } from "./env.server";
 

apps/webapp/app/platform/zodWorker.server.ts

@@ -16,7 +16,7 @@ import {
   parseCronItems,
   Logger as GraphileLogger,
 } from "graphile-worker";
-import { SpanKind, trace } from "@opentelemetry/api";
+import { SpanKind, SpanStatusCode, context, propagation, trace } from "@opentelemetry/api";
 
 import omit from "lodash.omit";
 import { z } from "zod";
@@ -25,13 +25,20 @@ import { PgListenService } from "~/services/db/pgListen.server";
 import { workerLogger as logger } from "~/services/logger.server";
 import { flattenAttributes } from "@trigger.dev/core/v3";
 import { env } from "~/env.server";
+import { getHttpContext } from "~/services/httpAsyncStorage.server";
 
 const tracer = trace.getTracer("zodWorker", "3.0.0.dp.1");
 
 export interface MessageCatalogSchema {
   [key: string]: z.ZodFirstPartySchemaTypes | z.ZodDiscriminatedUnion<any, any>;
 }
 
+const ZodWorkerMessageSchema = z.object({
+  version: z.literal("1"),
+  payload: z.unknown(),
+  context: z.record(z.string().optional()).optional(),
+});
+
 const RawCronPayloadSchema = z.object({
   _cron: z.object({
     ts: z.coerce.date(),
@@ -338,11 +345,45 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       }
     }
 
-    const { job, durationInMs } = await this.#addJob(
-      identifier as string,
-      payload,
-      spec,
-      options?.tx ?? this.#prisma
+    const { job, durationInMs } = await tracer.startActiveSpan(
+      `Enqueue ${identifier as string}`,
+      {
+        kind: SpanKind.PRODUCER,
+        attributes: {
+          "job.task_identifier": identifier as string,
+          "job.payload": payload,
+          "job.priority": spec.priority,
+          "job.run_at": spec.runAt?.toISOString(),
+          "job.jobKey": spec.jobKey,
+          "job.flags": spec.flags,
+          "job.max_attempts": spec.maxAttempts,
+          "worker.name": this.#name,
+        },
+      },
+      async (span) => {
+        try {
+          const results = await this.#addJob(
+            identifier as string,
+            payload,
+            spec,
+            options?.tx ?? this.#prisma
+          );
+
+          return results;
+        } catch (error) {
+          if (error instanceof Error) {
+            span.recordException(error);
+          } else {
+            span.recordException(new Error(String(error)));
+          }
+
+          span.setStatus({ code: SpanStatusCode.ERROR });
+
+          throw error;
+        } finally {
+          span.end();
+        }
+      }
     );
 
     logger.debug("Enqueued worker task", {
@@ -375,6 +416,18 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
   ) {
     const now = performance.now();
 
+    let $payload = payload;
+
+    if (!getHttpContext()) {
+      const $context = {};
+      propagation.inject(context.active(), $context);
+      $payload = {
+        version: "1",
+        payload,
+        context: $context,
+      };
+    }
+
     const results = await tx.$queryRawUnsafe(
       `SELECT * FROM ${this.graphileWorkerSchema}.add_job(
           identifier => $1::text,
@@ -387,7 +440,7 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
           job_key_mode => $8::text
         )`,
       identifier,
-      JSON.stringify(payload),
+      JSON.stringify($payload),
       spec.runAt || null,
       spec.maxAttempts || null,
       spec.jobKey || null,
@@ -401,6 +454,12 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
     const rows = AddJobResultsSchema.safeParse(results);
 
     if (!rows.success) {
+      logger.debug("results returned from add_job could not be parsed", {
+        identifier,
+        $payload,
+        spec,
+      });
+
       throw new Error(
         `Failed to add job to queue, zod parsing error: ${JSON.stringify(rows.error)}`
       );
@@ -543,7 +602,11 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       throw new Error(`Unknown message type: ${String(typeName)}`);
     }
 
-    const payload = messageSchema.parse(rawPayload);
+    const messagePayload = ZodWorkerMessageSchema.safeParse(rawPayload);
+
+    const payload = messageSchema.parse(
+      messagePayload.success ? messagePayload.data.payload : rawPayload
+    );
     const job = helpers.job;
 
     logger.debug("Received worker task, calling handler", {
@@ -558,6 +621,10 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       throw new Error(`No task for message type: ${String(typeName)}`);
     }
 
+    const activeContext = messagePayload.success
+      ? propagation.extract(context.active(), messagePayload.data.context ?? {})
+      : undefined;
+
     await tracer.startActiveSpan(
       `Run ${typeName as string}`,
       {
@@ -581,6 +648,7 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
           "worker.name": this.#name,
         },
       },
+      activeContext ?? context.active(),
       async (span) => {
         try {
           await task.handler(payload, job, helpers);

apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts

@@ -6,6 +6,7 @@ import { env } from "~/env.server";
 import { authenticateApiRequest } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { parseRequestJsonAsync } from "~/utils/parseRequestJson.server";
+import { ServiceValidationError } from "~/v3/services/baseService.server";
 import { TriggerTaskService } from "~/v3/services/triggerTask.server";
 import { startActiveSpan } from "~/v3/tracer.server";
 
@@ -92,18 +93,12 @@ export async function action({ request, params }: ActionFunctionArgs) {
       traceContext,
     });
 
-    const run = await service.call(
-      taskId,
-      authenticationResult.environment,
-      { ...body.data },
-      // { ...body.data, payload: (anyBody as any).payload },
-      {
-        idempotencyKey: idempotencyKey ?? undefined,
-        triggerVersion: triggerVersion ?? undefined,
-        traceContext,
-        spanParentAsLink: spanParentAsLink === 1,
-      }
-    );
+    const run = await service.call(taskId, authenticationResult.environment, body.data, {
+      idempotencyKey: idempotencyKey ?? undefined,
+      triggerVersion: triggerVersion ?? undefined,
+      traceContext,
+      spanParentAsLink: spanParentAsLink === 1,
+    });
 
     if (!run) {
       return json({ error: "Task not found" }, { status: 404 });
@@ -113,7 +108,9 @@ export async function action({ request, params }: ActionFunctionArgs) {
       id: run.friendlyId,
     });
   } catch (error) {
-    if (error instanceof Error) {
+    if (error instanceof ServiceValidationError) {
+      return json({ error: error.message }, { status: 422 });
+    } else if (error instanceof Error) {
       return json({ error: error.message }, { status: 400 });
     }
 

apps/webapp/app/routes/metrics.ts

@@ -1,7 +1,10 @@
 import { LoaderFunctionArgs } from "@remix-run/server-runtime";
 import { metricsRegister } from "~/metrics.server";
+import { logger } from "~/services/logger.server";
 
 export async function loader({ request }: LoaderFunctionArgs) {
+  logger.debug("Getting metrics from the metrics register");
+
   return new Response(await metricsRegister.metrics(), {
     headers: {
       "Content-Type": metricsRegister.contentType,

apps/webapp/app/routes/resources.orgs.$organizationSlug.projects.v3.$projectParam.runs.$runParam.spans.$spanParam/route.tsx

@@ -261,7 +261,7 @@ function RunActionButtons({ span }: { span: Span }) {
 
   if (span.isPartial) {
     return (
-      <Dialog>
+      <Dialog key="in-progress">
         <LinkButton
           to={v3RunDownloadLogsPath({ friendlyId: runParam })}
           LeadingIcon={CloudArrowDownIcon}
@@ -290,7 +290,7 @@ function RunActionButtons({ span }: { span: Span }) {
   }
 
   return (
-    <Dialog>
+    <Dialog key="complete">
       <LinkButton
         to={v3RunDownloadLogsPath({ friendlyId: runParam })}
         LeadingIcon={CloudArrowDownIcon}

apps/webapp/app/services/httpAsyncStorage.server.ts

@@ -0,0 +1,19 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+
+export type HttpLocalStorage = {
+  requestId: string;
+  path: string;
+  host: string;
+};
+
+const httpLocalStorage = new AsyncLocalStorage<HttpLocalStorage>();
+
+export type RunWithHttpContextFunction = <T>(context: HttpLocalStorage, fn: () => T) => T;
+
+export function runWithHttpContext<T>(context: HttpLocalStorage, fn: () => T): T {
+  return httpLocalStorage.run(context, fn);
+}
+
+export function getHttpContext(): HttpLocalStorage | undefined {
+  return httpLocalStorage.getStore();
+}

apps/webapp/app/services/logger.server.ts

@@ -2,6 +2,7 @@ import type { LogLevel } from "@trigger.dev/core-backend";
 import { Logger } from "@trigger.dev/core-backend";
 import { sensitiveDataReplacer } from "./sensitiveDataReplacer";
 import { AsyncLocalStorage } from "async_hooks";
+import { getHttpContext } from "./httpAsyncStorage.server";
 
 const currentFieldsStore = new AsyncLocalStorage<Record<string, unknown>>();
 
@@ -16,7 +17,8 @@ export const logger = new Logger(
   sensitiveDataReplacer,
   () => {
     const fields = currentFieldsStore.getStore();
-    return fields ? { ...fields } : {};
+    const httpContext = getHttpContext();
+    return { ...fields, http: httpContext };
   }
 );
 

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

@@ -18,6 +18,10 @@ export abstract class BaseService {
         try {
           return await fn(span);
         } catch (e) {
+          if (e instanceof ServiceValidationError) {
+            throw e;
+          }
+
           if (e instanceof Error) {
             span.recordException(e);
           } else {