WIP

Author: ericallam

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;
 

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

@@ -5,12 +5,18 @@ import { prisma } from "~/db.server";
 import { env } from "~/env.server";
 import { ZodWorker } from "~/platform/zodWorker.server";
 import { eventRepository } from "~/v3/eventRepository.server";
+import { MarqsConcurrencyMonitor } from "~/v3/marqs/concurrencyMonitor.server";
+import { RequeueV2Message } from "~/v3/marqs/requeueV2Message.server";
+import { reportUsageEvent } from "~/v3/openMeter.server";
 import { RequeueTaskRunService } from "~/v3/requeueTaskRun.server";
 import { DeliverAlertService } from "~/v3/services/alerts/deliverAlert.server";
 import { PerformDeploymentAlertsService } from "~/v3/services/alerts/performDeploymentAlerts.server";
 import { PerformTaskAttemptAlertsService } from "~/v3/services/alerts/performTaskAttemptAlerts.server";
 import { PerformBulkActionService } from "~/v3/services/bulk/performBulkAction.server";
+import { CancelTaskAttemptDependenciesService } from "~/v3/services/cancelTaskAttemptDependencies.server";
+import { EnqueueDelayedRunService } from "~/v3/services/enqueueDelayedRun.server";
 import { ExecuteTasksWaitingForDeployService } from "~/v3/services/executeTasksWaitingForDeploy";
+import { ExpireEnqueuedRunService } from "~/v3/services/expireEnqueuedRun.server";
 import { IndexDeploymentService } from "~/v3/services/indexDeployment.server";
 import { ResumeBatchRunService } from "~/v3/services/resumeBatchRun.server";
 import { ResumeTaskDependencyService } from "~/v3/services/resumeTaskDependency.server";
@@ -44,11 +50,6 @@ import { DeliverWebhookRequestService } from "./sources/deliverWebhookRequest.se
 import { PerformTaskOperationService } from "./tasks/performTaskOperation.server";
 import { ProcessCallbackTimeoutService } from "./tasks/processCallbackTimeout.server";
 import { ResumeTaskService } from "./tasks/resumeTask.server";
-import { RequeueV2Message } from "~/v3/marqs/requeueV2Message.server";
-import { MarqsConcurrencyMonitor } from "~/v3/marqs/concurrencyMonitor.server";
-import { reportUsageEvent } from "~/v3/openMeter.server";
-import { EnqueueDelayedRunService } from "~/v3/services/enqueueDelayedRun.server";
-import { ExpireEnqueuedRunService } from "~/v3/services/expireEnqueuedRun.server";
 
 const workerCatalog = {
   indexEndpoint: z.object({
@@ -185,6 +186,9 @@ const workerCatalog = {
   "v3.expireRun": z.object({
     runId: z.string(),
   }),
+  "v3.cancelTaskAttemptDependencies": z.object({
+    attemptId: z.string(),
+  }),
 };
 
 const executionWorkerCatalog = {
@@ -698,6 +702,15 @@ function getWorkerQueue() {
           return await service.call(payload.runId);
         },
       },
+      "v3.cancelTaskAttemptDependencies": {
+        priority: 0,
+        maxAttempts: 8,
+        handler: async (payload, job) => {
+          const service = new CancelTaskAttemptDependenciesService();
+
+          return await service.call(payload.attemptId);
+        },
+      },
     },
   });
 }

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

@@ -1411,7 +1411,9 @@ function filteredAttributes(attributes: Attributes, prefix: string): Attributes
 }
 
 function calculateDurationFromStart(startTime: bigint, endTime: Date = new Date()) {
-  return Number(BigInt(endTime.getTime() * 1_000_000) - startTime);
+  const $endtime = typeof endTime === "string" ? new Date(endTime) : endTime;
+
+  return Number(BigInt($endtime.getTime() * 1_000_000) - startTime);
 }
 
 function getNowInNanoseconds(): bigint {

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

@@ -1,12 +1,12 @@
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
-import { eventRepository } from "../eventRepository.server";
+import { logger } from "~/services/logger.server";
 import { marqs } from "~/v3/marqs/index.server";
+import { eventRepository } from "../eventRepository.server";
 import { BaseService } from "./baseService.server";
-import { logger } from "~/services/logger.server";
 
 import { PrismaClientOrTransaction, prisma } from "~/db.server";
-import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
 import { isCancellableRunStatus } from "../taskStatus";
+import { ResumeTaskRunDependenciesService } from "./resumeTaskRunDependencies.server";
 
 export class CancelAttemptService extends BaseService {
   public async call(
@@ -43,6 +43,14 @@ export class CancelAttemptService extends BaseService {
         return;
       }
 
+      if (taskRunAttempt.status === "CANCELED") {
+        logger.warn("Task run attempt is already cancelled", {
+          attemptId,
+        });
+
+        return;
+      }
+
       await marqs?.acknowledgeMessage(taskRunId);
 
       await this._prisma.taskRunAttempt.update({

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

@@ -0,0 +1,55 @@
+import { PrismaClientOrTransaction } from "~/db.server";
+import { workerQueue } from "~/services/worker.server";
+import { BaseService } from "./baseService.server";
+import { logger } from "~/services/logger.server";
+import { CancelTaskRunService } from "./cancelTaskRun.server";
+
+export class CancelTaskAttemptDependenciesService extends BaseService {
+  public async call(attemptId: string) {
+    const taskAttempt = await this._prisma.taskRunAttempt.findUnique({
+      where: { id: attemptId },
+      include: {
+        dependencies: {
+          include: {
+            taskRun: true,
+          },
+        },
+      },
+    });
+
+    if (!taskAttempt) {
+      return;
+    }
+
+    if (!taskAttempt.dependencies.length) {
+      return;
+    }
+
+    if (taskAttempt.status !== "CANCELED") {
+      logger.debug("Task attempt is not cancelled, continuing anyway", {
+        attemptId,
+        status: taskAttempt.status,
+      });
+    }
+
+    const cancelRunService = new CancelTaskRunService();
+
+    for (const dependency of taskAttempt.dependencies) {
+      await cancelRunService.call(dependency.taskRun);
+    }
+  }
+
+  static async enqueue(attemptId: string, tx: PrismaClientOrTransaction, runAt?: Date) {
+    return await workerQueue.enqueue(
+      "v3.cancelTaskAttemptDependencies",
+      {
+        attemptId,
+      },
+      {
+        tx,
+        runAt,
+        jobKey: `cancelTaskAttemptDependencies:${attemptId}`,
+      }
+    );
+  }
+}

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

@@ -8,6 +8,7 @@ import { devPubSub } from "../marqs/devPubSub.server";
 import { BaseService } from "./baseService.server";
 import { CancelAttemptService } from "./cancelAttempt.server";
 import { CANCELLABLE_ATTEMPT_STATUSES, isCancellableRunStatus } from "../taskStatus";
+import { CancelTaskAttemptDependenciesService } from "./cancelTaskAttemptDependencies.server";
 
 type ExtendedTaskRun = Prisma.TaskRunGetPayload<{
   include: {
@@ -66,6 +67,16 @@ export class CancelTaskRunService extends BaseService {
           },
           include: {
             backgroundWorker: true,
+            dependencies: {
+              include: {
+                taskRun: true,
+              },
+            },
+            batchTaskRunItems: {
+              include: {
+                taskRun: true,
+              },
+            },
           },
         },
         runtimeEnvironment: true,
@@ -93,6 +104,18 @@ export class CancelTaskRunService extends BaseService {
       await this.#cancelRemainingRunWorkers(cancelledTaskRun);
     }
 
+    const cancelService = new CancelTaskRunService();
+
+    // Cancel any dependent task runs
+    for (const attempt of cancelledTaskRun.attempts) {
+      for (const dependency of attempt.dependencies) {
+        await cancelService.call(dependency.taskRun, {
+          ...opts,
+          reason: `Parent task run was cancelled`,
+        });
+      }
+    }
+
     return {
       id: cancelledTaskRun.id,
     };
@@ -103,6 +126,8 @@ export class CancelTaskRunService extends BaseService {
     attempts: ExtendedTaskRunAttempt[]
   ) {
     for (const attempt of attempts) {
+      await CancelTaskAttemptDependenciesService.enqueue(attempt.id, this._prisma);
+
       if (run.runtimeEnvironment.type === "DEVELOPMENT") {
         // Signal the task run attempt to stop
         await devPubSub.publish(

references/v3-catalog/src/trigger/subtasks.ts

@@ -19,7 +19,7 @@ export const simpleChildTask = task({
   run: async (payload: { message: string }, { ctx }) => {
     logger.log("Simple child task payload", { payload, ctx });
 
-    await wait.for({ seconds: 6 });
+    await wait.for({ seconds: 30 });
   },
 });
 
@@ -162,7 +162,9 @@ export const deeplyNestedTaskChild = task({
 export const deeplyNestedTaskGrandchild = task({
   id: "deeply-nested-task-grandchild",
   run: async ({ message = "test" }: { message?: string }) => {
-    await deeplyNestedTaskGreatGrandchild.triggerAndWait({ message: `${message} - 3` });
+    await deeplyNestedTaskGreatGrandchild.batchTriggerAndWait(
+      Array.from({ length: 100 }, (_, i) => ({ payload: { message: `${message} - ${i}` } }))
+    );
 
     return {
       hello: "world",
@@ -173,6 +175,41 @@ export const deeplyNestedTaskGrandchild = task({
 export const deeplyNestedTaskGreatGrandchild = task({
   id: "deeply-nested-task-great-grandchild",
   run: async ({ message = "test" }: { message?: string }) => {
+    await new Promise((resolve) => setTimeout(resolve, 10000));
+
+    return {
+      hello: "world",
+    };
+  },
+});
+
+export const dependencyCancellationParent = task({
+  id: "dependency-cancellation-parent",
+  run: async ({ message = "test" }: { message?: string }) => {
+    const handle = await dependencyCancellationChild.triggerAndWait({ message: `${message} - 1` });
+
+    return {
+      hello: "world",
+    };
+  },
+});
+
+export const dependencyCancellationChild = task({
+  id: "dependency-cancellation-child",
+  run: async ({ message = "test" }: { message?: string }) => {
+    await dependencyCancellationGrandchild.triggerAndWait({ message: `${message} - 2` });
+
+    return {
+      hello: "world",
+    };
+  },
+});
+
+export const dependencyCancellationGrandchild = task({
+  id: "dependency-cancellation-grandchild",
+  run: async ({ message = "test" }: { message?: string }) => {
+    await wait.for({ seconds: 30 });
+
     return {
       hello: "world",
     };