fix a few more resume issues

This commit is contained in:
nicktrn
2024-03-24 22:16:01 +00:00
parent 9c586759a6
commit 26439105a9
3 changed files with 99 additions and 112 deletions
@@ -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 {
@@ -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<typeof CoordinatorToPlatformMessages, "READY_FOR_RESUME">
): Promise<void> {
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",
},
},
},
+30 -38
View File
@@ -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;