From 26439105a9872f3a28dd21d4a9ec6fbbe1ae57a4 Mon Sep 17 00:00:00 2001 From: nicktrn <55853254+nicktrn@users.noreply.github.com> Date: Sun, 24 Mar 2024 22:16:01 +0000 Subject: [PATCH] fix a few more resume issues --- .../v3/marqs/sharedQueueConsumer.server.ts | 124 +++++++++--------- .../app/v3/services/resumeAttempt.server.ts | 19 +-- .../cli-v3/src/workers/prod/entry-point.ts | 68 +++++----- 3 files changed, 99 insertions(+), 112 deletions(-) diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 5bb4b1ef0..2929343fe 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -138,11 +138,6 @@ export class SharedQueueConsumer { logger.debug("Stopping shared queue consumer"); this._enabled = false; - - // TODO: think about automatic prod cancellation - - // We need to cancel all the in progress task run attempts and ack the messages so they will stop processing - // await this.#cancelInProgressAttempts(reason); } async #cancelInProgressAttempts(reason: string) { @@ -249,7 +244,7 @@ export class SharedQueueConsumer { const message = await marqs?.dequeueMessageInSharedQueue(); if (!message) { - setTimeout(() => this.#doWork(), this._options.nextTickInterval); + this.#doMoreWork(this._options.nextTickInterval); return; } @@ -272,8 +267,8 @@ export class SharedQueueConsumer { queueMessage: message.data, envId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + this.#ackAndDoMoreWork(message.messageId); return; } @@ -286,9 +281,7 @@ export class SharedQueueConsumer { env: environment, }); - await marqs?.acknowledgeMessage(message.messageId); - - setTimeout(() => this.#doWork(), this._options.interval); + this.#ackAndDoMoreWork(message.messageId); return; } @@ -305,8 +298,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + this.#ackAndDoMoreWork(message.messageId); return; } @@ -333,8 +326,8 @@ export class SharedQueueConsumer { status: existingTaskRun.status, retryingFromCheckpoint, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + this.#ackAndDoMoreWork(message.messageId); return; } @@ -345,8 +338,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + this.#ackAndDoMoreWork(message.messageId); return; } @@ -356,8 +349,8 @@ export class SharedQueueConsumer { messageId: message.messageId, deployment: deployment.id, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + this.#ackAndDoMoreWork(message.messageId); return; } @@ -374,9 +367,7 @@ export class SharedQueueConsumer { taskSlugs: deployment.worker.tasks.map((task) => task.slug), }); - await marqs?.acknowledgeMessage(message.messageId); - - setTimeout(() => this.#doWork(), this._options.interval); + this.#ackAndDoMoreWork(message.messageId); return; } @@ -412,9 +403,7 @@ export class SharedQueueConsumer { messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - - setTimeout(() => this.#doWork(), this._options.interval); + this.#ackAndDoMoreWork(message.messageId); return; } @@ -428,8 +417,7 @@ export class SharedQueueConsumer { }); if (!queue) { - await marqs?.nackMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.nextTickInterval); + await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval); return; } @@ -468,11 +456,10 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); + + await this.#ackAndDoMoreWork(message.messageId); return; } - - return; } else if (isRetry) { socketIo.coordinatorNamespace.emit("READY_FOR_RETRY", { version: "v1", @@ -520,11 +507,10 @@ export class SharedQueueConsumer { }), ]); - // Finally we need to nack the message so it can be retried - await marqs?.nackMessage(message.messageId); - } finally { - setTimeout(() => this.#doWork(), this._options.interval); + await this.#nackAndDoMoreWork(message.messageId); + return; } + break; } // Resume after dependency completed with no remaining retries @@ -542,7 +528,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); + + await this.#ackAndDoMoreWork(message.messageId); return; } } catch (e) { @@ -554,12 +541,12 @@ export class SharedQueueConsumer { this._endSpanInNextIteration = true; - // Finally we need to nack the message so it can be retried - await marqs?.nackMessage(message.messageId); + await this.#nackAndDoMoreWork(message.messageId); return; - } finally { - setTimeout(() => this.#doWork(), this._options.interval); } + + this.#doMoreWork(); + return; } if (messageBody.data.completedAttemptIds.length < 1) { @@ -567,8 +554,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + await this.#ackAndDoMoreWork(message.messageId); return; } @@ -583,8 +570,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + await this.#ackAndDoMoreWork(message.messageId); return; } @@ -607,8 +594,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + await this.#ackAndDoMoreWork(message.messageId); return; } @@ -622,8 +609,7 @@ export class SharedQueueConsumer { }); if (!queue) { - await marqs?.nackMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.nextTickInterval); + await this.#nackAndDoMoreWork(message.messageId, this._options.nextTickInterval); return; } @@ -655,16 +641,15 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + + await this.#ackAndDoMoreWork(message.messageId); return; } const completion = await this._tasks.getCompletionPayloadFromAttempt(completedAttempt.id); if (!completion) { - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + await this.#ackAndDoMoreWork(message.messageId); return; } @@ -675,8 +660,7 @@ export class SharedQueueConsumer { ); if (!executionPayload) { - await marqs?.acknowledgeMessage(message.messageId); - setTimeout(() => this.#doWork(), this._options.interval); + await this.#ackAndDoMoreWork(message.messageId); return; } @@ -701,11 +685,10 @@ export class SharedQueueConsumer { this._endSpanInNextIteration = true; - // Finally we need to nack the message so it can be retried - await marqs?.nackMessage(message.messageId); - } finally { - setTimeout(() => this.#doWork(), this._options.interval); + await this.#nackAndDoMoreWork(message.messageId); + return; } + break; } // Resume after duration-based wait @@ -722,7 +705,8 @@ export class SharedQueueConsumer { queueMessage: message.data, messageId: message.messageId, }); - await marqs?.acknowledgeMessage(message.messageId); + + await this.#ackAndDoMoreWork(message.messageId); return; } } catch (e) { @@ -734,19 +718,35 @@ export class SharedQueueConsumer { this._endSpanInNextIteration = true; - // Finally we need to nack the message so it can be retried - await marqs?.nackMessage(message.messageId); - } finally { - setTimeout(() => this.#doWork(), this._options.interval); + await this.#nackAndDoMoreWork(message.messageId); + return; } + break; } } + + this.#doMoreWork(); + return; } #envIdFromQueue(queueName: string) { return queueName.split(":")[1]; } + + #doMoreWork(intervalInMs = this._options.interval) { + setTimeout(() => this.#doWork(), intervalInMs); + } + + async #ackAndDoMoreWork(messageId: string, intervalInMs?: number) { + await marqs?.acknowledgeMessage(messageId); + this.#doMoreWork(intervalInMs); + } + + async #nackAndDoMoreWork(messageId: string, intervalInMs?: number) { + await marqs?.nackMessage(messageId); + this.#doMoreWork(intervalInMs); + } } class SharedQueueTasks { diff --git a/apps/webapp/app/v3/services/resumeAttempt.server.ts b/apps/webapp/app/v3/services/resumeAttempt.server.ts index cbecf804d..001f0e520 100644 --- a/apps/webapp/app/v3/services/resumeAttempt.server.ts +++ b/apps/webapp/app/v3/services/resumeAttempt.server.ts @@ -4,25 +4,20 @@ import { TaskRunExecution, TaskRunExecutionResult, } from "@trigger.dev/core/v3"; -import { $transaction, PrismaClient, prisma } from "~/db.server"; +import { $transaction } from "~/db.server"; import { logger } from "~/services/logger.server"; import { marqs } from "../marqs.server"; import { socketIo } from "../handleSocketIo.server"; import { sharedQueueTasks } from "../marqs/sharedQueueConsumer.server"; +import { BaseService } from "./baseService.server"; -export class ResumeAttemptService { - #prismaClient: PrismaClient; - - constructor(prismaClient: PrismaClient = prisma) { - this.#prismaClient = prismaClient; - } - +export class ResumeAttemptService extends BaseService { public async call( params: InferSocketMessageSchema ): Promise { logger.debug(`ResumeAttemptService.call()`, params); - await $transaction(this.#prismaClient, async (tx) => { + await $transaction(this._prisma, async (tx) => { const attempt = await tx.taskRunAttempt.findUnique({ where: { friendlyId: params.attemptFriendlyId, @@ -146,7 +141,7 @@ export class ResumeAttemptService { const executions: TaskRunExecution[] = []; for (const completedAttemptId of completedAttemptIds) { - const completedAttempt = await prisma.taskRunAttempt.findUnique({ + const completedAttempt = await tx.taskRunAttempt.findUnique({ where: { id: completedAttemptId, taskRun: { @@ -200,7 +195,7 @@ export class ResumeAttemptService { executions.push(executionPayload.execution); } - await prisma.taskRunAttempt.update({ + const updated = await tx.taskRunAttempt.update({ where: { id: attempt.id, }, @@ -209,7 +204,7 @@ export class ResumeAttemptService { taskRun: { update: { data: { - status: "EXECUTING", + status: attempt.number > 1 ? "RETRYING_AFTER_FAILURE" : "EXECUTING", }, }, }, diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index f9eac4c3f..e13dfc237 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -97,7 +97,7 @@ class ProdWorker { } ); - this.#prepareForCheckpoint("WAIT_FOR_DURATION", willCheckpointAndRestore); + this.#prepareForWait("WAIT_FOR_DURATION", willCheckpointAndRestore); }); this.#backgroundWorker.onWaitForTask.attach(async (message) => { @@ -115,7 +115,7 @@ class ProdWorker { } ); - this.#prepareForCheckpoint("WAIT_FOR_TASK", willCheckpointAndRestore); + this.#prepareForWait("WAIT_FOR_TASK", willCheckpointAndRestore); }); this.#backgroundWorker.onWaitForBatch.attach(async (message) => { @@ -133,7 +133,7 @@ class ProdWorker { } ); - this.#prepareForCheckpoint("WAIT_FOR_BATCH", willCheckpointAndRestore); + this.#prepareForWait("WAIT_FOR_BATCH", willCheckpointAndRestore); }); this.#httpPort = port; @@ -169,8 +169,8 @@ class ProdWorker { } } - #prepareForCheckpoint(reason: WaitReason, willCheckpointAndRestore: boolean) { - logger.log(reason, { willCheckpointAndRestore }); + #prepareForWait(reason: WaitReason, willCheckpointAndRestore: boolean) { + logger.log(`prepare for ${reason}`, { willCheckpointAndRestore }); this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore }); @@ -178,8 +178,27 @@ class ProdWorker { this.paused = true; this.nextResumeAfter = reason; } + } - if (this.runningInKubernetes) { + async #prepareForRetry(willCheckpointAndRestore: boolean, shouldExit: boolean) { + logger.log("prepare for retry", { willCheckpointAndRestore, shouldExit }); + + // Graceful shutdown on final attempt + if (shouldExit) { + if (willCheckpointAndRestore) { + logger.log("WARNING: Will checkpoint but also requested exit. This won't end well."); + } + + await this.#backgroundWorker.close(); + process.exit(0); + } + + this.executing = false; + this.attemptFriendlyId = undefined; + + if (willCheckpointAndRestore) { + this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); + this.#coordinatorSocket.close(); return; } } @@ -267,35 +286,7 @@ class ProdWorker { logger.log("completion acknowledged", { willCheckpointAndRestore, shouldExit }); - // Graceful shutdown on final attempt - if (shouldExit) { - if (willCheckpointAndRestore) { - logger.log("WARNING: Will checkpoint but also requested exit. This won't end well."); - } - - await this.#backgroundWorker.close(); - process.exit(0); - } - - this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); - - // Give coordinator a chance to request exit - // TODO: Consider disabling automatic reconnect instead, and re-enabling it on postStart hook - await new Promise((resolve) => { - // The timeout duration is below the minimum wait duration that triggers a checkpoint - // We are unlikely to restore before this, so this will have resolved on restore - setTimeout(resolve, 5_000); - }); - - // Remove executing state as late as possible to prevent further execution until we've - this.executing = false; - this.attemptFriendlyId = undefined; - - if (willCheckpointAndRestore) { - this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); - this.#coordinatorSocket.close(); - return; - } + this.#prepareForRetry(willCheckpointAndRestore, shouldExit); }, REQUEST_ATTEMPT_CANCELLATION: async (message) => { if (!this.executing) { @@ -389,7 +380,7 @@ class ProdWorker { } await setTimeout(200); - process.exit(1); + process.exit(1); } } @@ -410,7 +401,9 @@ class ProdWorker { }); this.#backgroundWorker.waitCompletedNotification(); + this.paused = false; + this.nextResumeAfter = undefined; return; } @@ -546,7 +539,6 @@ class ProdWorker { break; } } - logger.log("postStart", { url: req.url }); return reply.text("postStart ok"); } @@ -580,7 +572,7 @@ class ProdWorker { this.#httpPort = getRandomPortNumber(); await setTimeout(100); - this.start(); + this.start(); }); return httpServer;