Move the visibility queue stuff into a graphile job

This commit is contained in:
Eric Allam
2024-04-30 17:29:38 +01:00
parent d1bdd0cf5d
commit e24631e2dd
9 changed files with 183 additions and 60 deletions
+13
View File
@@ -37,6 +37,7 @@ import { TimeoutDeploymentService } from "~/v3/services/timeoutDeployment.server
import { eventRepository } from "~/v3/eventRepository.server";
import { ExecuteTasksWaitingForDeployService } from "~/v3/services/executeTasksWaitingForDeploy";
import { TriggerScheduledTaskService } from "~/v3/services/triggerScheduledTask.server";
import { RequeueTaskRunService } from "~/v3/requeueTaskRun.server";
const workerCatalog = {
indexEndpoint: z.object({
@@ -136,6 +137,9 @@ const workerCatalog = {
"v3.triggerScheduledTask": z.object({
instanceId: z.string(),
}),
"v3.requeueTaskRun": z.object({
runId: z.string(),
}),
};
const executionWorkerCatalog = {
@@ -533,6 +537,15 @@ function getWorkerQueue() {
return await service.call(payload.instanceId);
},
},
"v3.requeueTaskRun": {
priority: 0,
maxAttempts: 3,
handler: async (payload, job) => {
const service = new RequeueTaskRunService();
await service.call(payload.runId);
},
},
},
});
}
+2 -11
View File
@@ -3,26 +3,17 @@ import {
TaskRunError,
TaskRunFailedExecutionResult,
} from "@trigger.dev/core/v3";
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { logger } from "~/services/logger.server";
import { marqs } from "~/v3/marqs/index.server";
import { TaskRunStatus } from "@trigger.dev/database";
import { eventRepository } from "./eventRepository.server";
import { BaseService } from "./services/baseService.server";
import { TaskRunStatus } from "@trigger.dev/database";
const FAILABLE_TASK_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "PENDING", "WAITING_FOR_DEPLOY"];
export class FailedTaskRunService extends BaseService {
public async call({
runFriendlyId,
completion,
env,
}: {
runFriendlyId: string;
completion: TaskRunFailedExecutionResult;
env: AuthenticatedEnvironment;
}) {
public async call(runFriendlyId: string, completion: TaskRunFailedExecutionResult) {
const taskRun = await this._prisma.taskRun.findUnique({
where: { friendlyId: runFriendlyId },
});
@@ -154,11 +154,7 @@ export class DevQueueConsumer {
const service = new FailedTaskRunService();
await service.call({
runFriendlyId: completion.id,
completion,
env: this.env,
});
await service.call(completion.id, completion);
}
/**
@@ -413,7 +409,6 @@ export class DevQueueConsumer {
data: {
lockedAt: new Date(),
lockedById: backgroundTask.id,
status: "EXECUTING",
},
include: {
attempts: {
+50 -25
View File
@@ -20,6 +20,7 @@ import {
MessagePayload,
QueueCapacities,
} from "./types";
import { workerQueue } from "~/services/worker.server";
const tracer = trace.getTracer("marqs");
@@ -258,6 +259,17 @@ export class MarQS {
});
}
await workerQueue.enqueue(
"v3.requeueTaskRun",
{
runId: messageData.messageId,
},
{
runAt: new Date(Date.now() + this.visibilityTimeoutInMs),
jobKey: `requeueTaskRun:${messageData.messageId}`,
}
);
return message;
},
{
@@ -349,6 +361,8 @@ export class MarQS {
[SemanticAttributes.PARENT_QUEUE]: message.parentQueue,
});
workerQueue.dequeue(`requeueTaskRun:${messageId}`);
await this.#callAcknowledgeMessage({
parentQueue: message.parentQueue,
messageKey: this.keys.messageKey(messageId),
@@ -505,16 +519,28 @@ export class MarQS {
// This should increment by the number of seconds, but with a max value of Date.now() + visibilityTimeoutInMs
public async heartbeatMessage(messageId: string, seconds: number = 30) {
// We are still calling this for backwards compatibility, but we should be using the v3.requeueTaskRun job
await this.#callHeartbeatMessage({
visibilityQueue: constants.MESSAGE_VISIBILITY_TIMEOUT_QUEUE,
messageId,
milliseconds: seconds * 1000,
maxVisibilityTimeout: Date.now() + this.visibilityTimeoutInMs,
});
await workerQueue.enqueue(
"v3.requeueTaskRun",
{
runId: messageId,
},
{
runAt: new Date(Date.now() + seconds * 1000),
jobKey: `requeueTaskRun:${messageId}`,
}
);
}
get visibilityTimeoutInMs() {
return this.options.visibilityTimeoutInMs ?? 300000;
return this.options.visibilityTimeoutInMs ?? 300000; // 5 minutes
}
async readMessage(messageId: string) {
@@ -861,7 +887,6 @@ export class MarQS {
const result = await this.redis.dequeueMessage(
messageQueue,
parentQueue,
visibilityQueue,
concurrencyLimitKey,
envConcurrencyLimitKey,
orgConcurrencyLimitKey,
@@ -869,7 +894,6 @@ export class MarQS {
envCurrentConcurrencyKey,
orgCurrentConcurrencyKey,
messageQueue,
String(this.options.visibilityTimeoutInMs ?? 300000), // 5 minutes
String(Date.now()),
String(this.options.defaultEnvConcurrency),
String(this.options.defaultOrgConcurrency)
@@ -995,6 +1019,9 @@ export class MarQS {
);
}
/**
* @deprecated This is being replaced by the v3.requeueTaskRun graphile worker job
*/
#callHeartbeatMessage({
visibilityQueue,
messageId,
@@ -1133,25 +1160,23 @@ end
});
this.redis.defineCommand("dequeueMessage", {
numberOfKeys: 9,
numberOfKeys: 8,
lua: `
-- Keys: childQueue, parentQueue, visibilityQueue, concurrencyLimitKey, envConcurrencyLimitKey, orgConcurrencyLimitKey, currentConcurrencyKey, envCurrentConcurrencyKey, orgCurrentConcurrencyKey
-- Keys: childQueue, parentQueue, concurrencyLimitKey, envConcurrencyLimitKey, orgConcurrencyLimitKey, currentConcurrencyKey, envCurrentConcurrencyKey, orgCurrentConcurrencyKey
local childQueue = KEYS[1]
local parentQueue = KEYS[2]
local visibilityQueue = KEYS[3]
local concurrencyLimitKey = KEYS[4]
local envConcurrencyLimitKey = KEYS[5]
local orgConcurrencyLimitKey = KEYS[6]
local currentConcurrencyKey = KEYS[7]
local envCurrentConcurrencyKey = KEYS[8]
local orgCurrentConcurrencyKey = KEYS[9]
local concurrencyLimitKey = KEYS[3]
local envConcurrencyLimitKey = KEYS[4]
local orgConcurrencyLimitKey = KEYS[5]
local currentConcurrencyKey = KEYS[6]
local envCurrentConcurrencyKey = KEYS[7]
local orgCurrentConcurrencyKey = KEYS[8]
-- Args: childQueueName, visibilityQueue, currentTime, defaultEnvConcurrencyLimit, defaultOrgConcurrencyLimit
-- Args: childQueueName, currentTime, defaultEnvConcurrencyLimit, defaultOrgConcurrencyLimit
local childQueueName = ARGV[1]
local visibilityTimeout = tonumber(ARGV[2])
local currentTime = tonumber(ARGV[3])
local defaultEnvConcurrencyLimit = ARGV[4]
local defaultOrgConcurrencyLimit = ARGV[5]
local currentTime = tonumber(ARGV[2])
local defaultEnvConcurrencyLimit = ARGV[3]
local defaultOrgConcurrencyLimit = ARGV[4]
-- Check current org concurrency against the limit
local orgCurrentConcurrency = tonumber(redis.call('SCARD', orgCurrentConcurrencyKey) or '0')
@@ -1187,11 +1212,9 @@ end
local messageId = messages[1]
local messageScore = tonumber(messages[2])
local timeoutScore = currentTime + visibilityTimeout
-- Move message to timeout queue and update concurrency
redis.call('ZREM', childQueue, messageId)
redis.call('ZADD', visibilityQueue, timeoutScore, messageId)
redis.call('SADD', currentConcurrencyKey, messageId)
redis.call('SADD', envCurrentConcurrencyKey, messageId)
redis.call('SADD', orgCurrentConcurrencyKey, messageId)
@@ -1257,7 +1280,7 @@ else
redis.call('ZADD', parentQueue, earliestMessage[2], messageQueueName)
end
-- Remove the message from the timeout queue
-- Remove the message from the timeout queue (deprecated, will eventually remove this)
redis.call('ZREM', visibilityQueue, messageId)
-- Update the concurrency keys
@@ -1297,7 +1320,7 @@ redis.call('SREM', concurrencyKey, messageId)
redis.call('SREM', envConcurrencyKey, messageId)
redis.call('SREM', orgConcurrencyKey, messageId)
-- Remove the message from the timeout queue
-- Remove the message from the timeout queue (deprecated, will eventually remove this)
redis.call('ZREM', visibilityQueue, messageId)
-- Enqueue the message into the queue
@@ -1325,12 +1348,16 @@ local milliseconds = tonumber(ARGV[2])
local maxVisibilityTimeout = tonumber(ARGV[3])
-- Get the current visibility timeout
local currentVisibilityTimeout = tonumber(redis.call('ZSCORE', visibilityQueue, messageId)) or 0
local zscoreResult = redis.call('ZSCORE', visibilityQueue, messageId)
if currentVisibilityTimeout == 0 then
-- If there's no currentVisibilityTimeout, return and do not execute ZADD
if zscoreResult == false then
return
end
local currentVisibilityTimeout = tonumber(zscoreResult)
-- Calculate the new visibility timeout
local newVisibilityTimeout = math.min(currentVisibilityTimeout + milliseconds * 1000, maxVisibilityTimeout)
@@ -1433,7 +1460,6 @@ declare module "ioredis" {
dequeueMessage(
childQueue: string,
parentQueue: string,
visibilityQueue: string,
concurrencyLimitKey: string,
envConcurrencyLimitKey: string,
orgConcurrencyLimitKey: string,
@@ -1441,7 +1467,6 @@ declare module "ioredis" {
envCurrentConcurrencyKey: string,
orgCurrentConcurrencyKey: string,
childQueueName: string,
visibilityTimeout: string,
currentTime: string,
defaultEnvConcurrencyLimit: string,
defaultOrgConcurrencyLimit: string,
@@ -0,0 +1,81 @@
import { logger } from "~/services/logger.server";
import { marqs } from "~/v3/marqs/index.server";
import assertNever from "assert-never";
import { FailedTaskRunService } from "./failedTaskRun.server";
import { BaseService } from "./services/baseService.server";
export class RequeueTaskRunService extends BaseService {
public async call(runId: string) {
const taskRun = await this._prisma.taskRun.findUnique({
where: { id: runId },
});
if (!taskRun) {
logger.error("[RequeueTaskRunService] Task run not found", {
runId,
});
return;
}
switch (taskRun.status) {
case "PENDING": {
logger.debug("[RequeueTaskRunService] Requeueing task run", { taskRun });
await marqs?.nackMessage(taskRun.id);
break;
}
case "EXECUTING":
case "RETRYING_AFTER_FAILURE": {
logger.debug("[RequeueTaskRunService] Failing task run", { taskRun });
const service = new FailedTaskRunService();
await service.call(taskRun.friendlyId, {
ok: false,
id: taskRun.friendlyId,
retry: undefined,
error: {
type: "INTERNAL_ERROR",
code: "TASK_RUN_HEARTBEAT_TIMEOUT",
message: "Did not receive a heartbeat from the worker in time",
},
});
break;
}
case "WAITING_FOR_DEPLOY": {
logger.debug("[RequeueTaskRunService] Removing task run from queue", { taskRun });
await marqs?.acknowledgeMessage(taskRun.id);
break;
}
case "WAITING_TO_RESUME":
case "PAUSED": {
logger.debug("[RequeueTaskRunService] Requeueing task run", { taskRun });
await marqs?.nackMessage(taskRun.id);
break;
}
case "SYSTEM_FAILURE":
case "INTERRUPTED":
case "CRASHED":
case "COMPLETED_WITH_ERRORS":
case "COMPLETED_SUCCESSFULLY":
case "CANCELED": {
logger.debug("[RequeueTaskRunService] Task run is completed", { taskRun });
await marqs?.acknowledgeMessage(taskRun.id);
break;
}
default: {
assertNever(taskRun.status);
}
}
}
}
@@ -1,9 +1,9 @@
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { BaseService, ServiceValidationError } from "./baseService.server";
import { TaskRunExecution } from "@trigger.dev/core/v3";
import { prisma } from "~/db.server";
import { generateFriendlyId } from "../friendlyIdentifiers";
import { $transaction } from "~/db.server";
import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
import { logger } from "~/services/logger.server";
import { generateFriendlyId } from "../friendlyIdentifiers";
import { BaseService, ServiceValidationError } from "./baseService.server";
export class CreateTaskRunAttemptService extends BaseService {
public async call(
@@ -61,20 +61,37 @@ export class CreateTaskRunAttemptService extends BaseService {
throw new ServiceValidationError("Queue not found", 404);
}
const taskRunAttempt = await prisma.taskRunAttempt.create({
data: {
number: taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1,
friendlyId: generateFriendlyId("attempt"),
taskRunId: taskRun.id,
startedAt: new Date(),
backgroundWorkerId: taskRun.lockedBy.worker.id,
backgroundWorkerTaskId: taskRun.lockedBy.id,
status: "EXECUTING" as const,
queueId: queue.id,
runtimeEnvironmentId: environment.id,
},
const taskRunAttempt = await $transaction(this._prisma, async (tx) => {
const taskRunAttempt = await tx.taskRunAttempt.create({
data: {
number: taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1,
friendlyId: generateFriendlyId("attempt"),
taskRunId: taskRun.id,
startedAt: new Date(),
backgroundWorkerId: taskRun.lockedBy!.worker.id,
backgroundWorkerTaskId: taskRun.lockedBy!.id,
status: "EXECUTING" as const,
queueId: queue.id,
runtimeEnvironmentId: environment.id,
},
});
await tx.taskRun.update({
where: {
id: taskRun.id,
},
data: {
status: "EXECUTING",
},
});
return taskRunAttempt;
});
if (!taskRunAttempt) {
throw new ServiceValidationError("Failed to create task run attempt", 500);
}
const execution: TaskRunExecution = {
task: {
id: taskRun.lockedBy.slug,
@@ -222,7 +222,7 @@ sender.send("TASKS_READY", { tasks: TASK_METADATA }).catch((err) => {
process.title = "trigger-dev-worker";
async function asyncHeartbeat(initialDelayInSeconds: number = 30, intervalInSeconds: number = 5) {
async function asyncHeartbeat(initialDelayInSeconds: number = 30, intervalInSeconds: number = 30) {
async function _doHeartbeat() {
while (true) {
if (_isRunning && _execution) {
+1
View File
@@ -51,6 +51,7 @@ export const TaskRunInternalError = z.object({
"TASK_OUTPUT_ERROR",
"HANDLE_ERROR_ERROR",
"GRACEFUL_EXIT_TIMEOUT",
"TASK_RUN_HEARTBEAT_TIMEOUT",
]),
message: z.string().optional(),
});
@@ -3,7 +3,7 @@ import { logger, task } from "@trigger.dev/sdk/v3";
export const longRunning = task({
id: "long-running",
run: async (payload: { message: string }) => {
logger.info("Long running payloadd", { payload });
logger.info("Long running payloadddd", { payload });
// Wait for 3 minutes
await new Promise((resolve) => setTimeout(resolve, 3 * 60 * 1000));