v3: Cancel awaited subtasks and reliable rate-limit recovery (#1200)

* v3: cancel subtasks when parent task runs are cancelled

* v3: recover from server rate limiting errors in a more reliable way

- Changing from sliding window to token bucket in the API rate limiter, to help smooth out traffic
- Adding spans to the API Client core & SDK functions
- Added waiting spans when retrying in the API Client
- Retrying in the API Client now respects the x-ratelimit-reset
- Retrying ApiError’s in tasks now respects the x-ratelimit-reset
- Added AbortTaskRunError that when thrown will stop retries
- Added idempotency keys SDK functions and automatically injecting the run ID when inside a task
- Added the ability to configure ApiRequestOptions (retries only for now) globally and on specific calls
- Implement the maxAttempts TaskRunOption (it wasn’t doing anything before)

* Adding some docs about the request options

* Fix type error

* Remove context propagation through graphile jobs

* Remove logger

* only select a subset of task run columns

* limit columns selected in batchTrigger as well

* added idempotency doc

* allow scoped idempotency keys, and fixed an issue with the unique index on BatchTaskRun and TaskRun

* Removed old cancel task run children code

Author: ericallam

.changeset/lemon-sloths-hide.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/core": patch
+---
+
+v3: recover from server rate limiting errors in a more reliable way

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

@@ -16,7 +16,6 @@ import {
 } from "./components/primitives/OperatingSystemProvider";
 import { getSharedSqsEventConsumer } from "./services/events/sqsEventConsumer";
 import { singleton } from "./utils/singleton";
-import { logger } from "./services/logger.server";
 
 const ABORT_DELAY = 30000;
 
@@ -186,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/env.server.ts

@@ -97,8 +97,9 @@ const EnvironmentSchema = z.object({
    * @example "1000ms"
    * @example "1000s"
    */
-  API_RATE_LIMIT_WINDOW: z.string().default("60s"),
-  API_RATE_LIMIT_MAX: z.coerce.number().int().default(600),
+  API_RATE_LIMIT_REFILL_INTERVAL: z.string().default("10s"), // refill 250 tokens every 10 seconds
+  API_RATE_LIMIT_MAX: z.coerce.number().int().default(750), // allow bursts of 750 requests
+  API_RATE_LIMIT_REFILL_RATE: z.coerce.number().int().default(250), // refix 250 tokens every 10 seconds
   API_RATE_LIMIT_REQUEST_LOGS_ENABLED: z.string().default("0"),
   API_RATE_LIMIT_REJECTION_LOGS_ENABLED: z.string().default("1"),
 

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

@@ -1,3 +1,4 @@
+import { SpanKind, SpanStatusCode, trace } from "@opentelemetry/api";
 import type {
   CronItem,
   CronItemOptions,
@@ -11,20 +12,19 @@ import type {
   WorkerUtils,
 } from "graphile-worker";
 import {
+  Logger as GraphileLogger,
   run as graphileRun,
   makeWorkerUtils,
   parseCronItems,
-  Logger as GraphileLogger,
 } from "graphile-worker";
-import { SpanKind, trace } from "@opentelemetry/api";
 
+import { flattenAttributes } from "@trigger.dev/core/v3";
 import omit from "lodash.omit";
 import { z } from "zod";
 import { $replica, PrismaClient, PrismaClientOrTransaction } from "~/db.server";
+import { env } from "~/env.server";
 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";
 
 const tracer = trace.getTracer("zodWorker", "3.0.0.dp.1");
 
@@ -338,11 +338,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", {
@@ -401,6 +435,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)}`
       );
@@ -422,9 +462,7 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       const job = AddJobResultsSchema.safeParse(result);
 
       if (!job.success) {
-        logger.debug("results returned from remove_job could not be parsed", {
-          error: job.error.flatten(),
-          result,
+        logger.debug("could not remove job, job_key did not exist", {
           jobKey,
         });
 

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/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/apiRateLimit.server.ts

@@ -106,8 +106,10 @@ export function authorizationRateLimitMiddleware({
       hashedAuthorizationValue
     );
 
+    const $remaining = Math.max(0, remaining); // remaining can be negative if the user has exceeded the limit, so clamp it to 0
+
     res.set("x-ratelimit-limit", limit.toString());
-    res.set("x-ratelimit-remaining", remaining.toString());
+    res.set("x-ratelimit-remaining", $remaining.toString());
     res.set("x-ratelimit-reset", reset.toString());
 
     if (success) {
@@ -122,12 +124,12 @@ export function authorizationRateLimitMiddleware({
           title: "Rate Limit Exceeded",
           status: 429,
           type: "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/429",
-          detail: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
+          detail: `Rate limit exceeded ${$remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
           reset,
           limit,
           remaining,
           secondsUntilReset,
-          error: `Rate limit exceeded ${remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
+          error: `Rate limit exceeded ${$remaining}/${limit} requests remaining. Retry in ${secondsUntilReset} seconds.`,
         },
         null,
         2
@@ -138,7 +140,11 @@ export function authorizationRateLimitMiddleware({
 
 export const apiRateLimiter = authorizationRateLimitMiddleware({
   keyPrefix: "api",
-  limiter: Ratelimit.slidingWindow(env.API_RATE_LIMIT_MAX, env.API_RATE_LIMIT_WINDOW as Duration),
+  limiter: Ratelimit.tokenBucket(
+    env.API_RATE_LIMIT_REFILL_RATE,
+    env.API_RATE_LIMIT_REFILL_INTERVAL as Duration,
+    env.API_RATE_LIMIT_MAX
+  ),
   pathMatchers: [/^\/api/],
   // Allow /api/v1/tasks/:id/callback/:secret
   pathWhiteList: [
@@ -152,6 +158,8 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({
     /^\/api\/v1\/sources\/http\/[^\/]+$/, // /api/v1/sources/http/$id
     /^\/api\/v1\/endpoints\/[^\/]+\/[^\/]+\/index\/[^\/]+$/, // /api/v1/endpoints/$environmentId/$endpointSlug/index/$indexHookIdentifier
     "/api/v1/timezones",
+    "/api/v1/usage/ingest",
+    /^\/api\/v1\/runs\/[^\/]+\/attempts$/, // /api/v1/runs/$runFriendlyId/attempts
   ],
   log: {
     rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1",

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 };
   }
 );