chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,37 @@
|
||||
export class AsyncWorker {
|
||||
private running = false;
|
||||
private timeout?: NodeJS.Timeout;
|
||||
|
||||
constructor(
|
||||
private readonly fn: () => Promise<void>,
|
||||
private readonly interval: number
|
||||
) {}
|
||||
|
||||
start() {
|
||||
if (this.running) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.running = true;
|
||||
|
||||
this.#run();
|
||||
}
|
||||
|
||||
stop() {
|
||||
this.running = false;
|
||||
}
|
||||
|
||||
async #run() {
|
||||
if (!this.running) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await this.fn();
|
||||
} catch (e) {
|
||||
console.error(e);
|
||||
}
|
||||
|
||||
this.timeout = setTimeout(this.#run.bind(this), this.interval);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
import type { Logger } from "@trigger.dev/core/logger";
|
||||
import type { Redis } from "ioredis";
|
||||
import { prisma } from "~/db.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import type { MarQS } from "./index.server";
|
||||
import { marqs as marqsv3 } from "./index.server";
|
||||
import { env } from "~/env.server";
|
||||
|
||||
export type MarqsConcurrencyMonitorOptions = {
|
||||
dryRun?: boolean;
|
||||
abortSignal?: AbortSignal;
|
||||
};
|
||||
|
||||
export interface MarqsConcurrencyResolveCompletedRunsCallback {
|
||||
(candidateRunIds: string[]): Promise<Array<{ id: string }>>;
|
||||
}
|
||||
|
||||
export class MarqsConcurrencyMonitor {
|
||||
private _logger: Logger;
|
||||
|
||||
constructor(
|
||||
private marqs: MarQS,
|
||||
private callback: MarqsConcurrencyResolveCompletedRunsCallback,
|
||||
private options: MarqsConcurrencyMonitorOptions = {}
|
||||
) {
|
||||
this._logger = logger.child({
|
||||
component: "marqs",
|
||||
operation: "concurrencyMonitor",
|
||||
dryRun: this.dryRun,
|
||||
marqs: marqs.name,
|
||||
});
|
||||
}
|
||||
|
||||
get dryRun() {
|
||||
return typeof this.options.dryRun === "boolean" ? this.options.dryRun : false;
|
||||
}
|
||||
|
||||
get keys() {
|
||||
return this.marqs.keys;
|
||||
}
|
||||
|
||||
get signal() {
|
||||
return this.options.abortSignal;
|
||||
}
|
||||
|
||||
public async call() {
|
||||
this._logger.debug("[MarqsConcurrencyMonitor] Initiating monitoring");
|
||||
|
||||
const stats = {
|
||||
streamCallbacks: 0,
|
||||
processedKeys: 0,
|
||||
};
|
||||
|
||||
const { stream, redis } = this.marqs.queueConcurrencyScanStream(
|
||||
10,
|
||||
() => {
|
||||
this._logger.debug("[MarqsConcurrencyMonitor] stream closed", {
|
||||
stats,
|
||||
});
|
||||
},
|
||||
(error) => {
|
||||
this._logger.debug("[MarqsConcurrencyMonitor] stream error", {
|
||||
stats,
|
||||
error: {
|
||||
name: error.name,
|
||||
message: error.message,
|
||||
stack: error.stack,
|
||||
},
|
||||
});
|
||||
}
|
||||
);
|
||||
|
||||
stream.on("data", async (keys) => {
|
||||
stream.pause();
|
||||
|
||||
if (this.signal?.aborted) {
|
||||
stream.destroy();
|
||||
return;
|
||||
}
|
||||
|
||||
stats.streamCallbacks++;
|
||||
|
||||
const uniqueKeys = Array.from(new Set<string>(keys));
|
||||
|
||||
if (uniqueKeys.length === 0) {
|
||||
stream.resume();
|
||||
return;
|
||||
}
|
||||
|
||||
this._logger.debug("[MarqsConcurrencyMonitor] correcting queues concurrency", {
|
||||
keys: uniqueKeys,
|
||||
});
|
||||
|
||||
stats.processedKeys += uniqueKeys.length;
|
||||
|
||||
await Promise.allSettled(uniqueKeys.map((key) => this.#processKey(key, redis))).finally(
|
||||
() => {
|
||||
stream.resume();
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
async #processKey(key: string, redis: Redis) {
|
||||
key = this.keys.stripKeyPrefix(key);
|
||||
const envKey = this.keys.envCurrentConcurrencyKeyFromQueue(key);
|
||||
|
||||
let runIds: string[] = [];
|
||||
|
||||
try {
|
||||
// Next, we need to get all the items from the key, and any parent keys (org, env, queue) using sunion.
|
||||
runIds = await redis.sunion(envKey, key);
|
||||
} catch (e) {
|
||||
this._logger.error("[MarqsConcurrencyMonitor] error during sunion", {
|
||||
key,
|
||||
envKey,
|
||||
runIds,
|
||||
error: e,
|
||||
});
|
||||
}
|
||||
|
||||
if (runIds.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const perfNow = performance.now();
|
||||
|
||||
const completeRuns = await this.callback(runIds);
|
||||
|
||||
const durationMs = performance.now() - perfNow;
|
||||
|
||||
const completedRunIds = completeRuns.map((run) => run.id);
|
||||
|
||||
if (completedRunIds.length === 0) {
|
||||
this._logger.debug("[MarqsConcurrencyMonitor] no completed runs found", {
|
||||
key,
|
||||
envKey,
|
||||
runIds,
|
||||
durationMs,
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
this._logger.debug("[MarqsConcurrencyMonitor] removing completed runs from queue", {
|
||||
key,
|
||||
envKey,
|
||||
completedRunIds,
|
||||
durationMs,
|
||||
});
|
||||
|
||||
if (this.dryRun) {
|
||||
return;
|
||||
}
|
||||
|
||||
const pipeline = redis.pipeline();
|
||||
|
||||
pipeline.srem(key, ...completedRunIds);
|
||||
pipeline.srem(envKey, ...completedRunIds);
|
||||
|
||||
try {
|
||||
await pipeline.exec();
|
||||
} catch (e) {
|
||||
this._logger.error("[MarqsConcurrencyMonitor] error removing completed runs from queue", {
|
||||
key,
|
||||
envKey,
|
||||
completedRunIds,
|
||||
error: e,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
static async initiateV3Monitoring(abortSignal?: AbortSignal) {
|
||||
if (!marqsv3) {
|
||||
return;
|
||||
}
|
||||
|
||||
const instance = new MarqsConcurrencyMonitor(
|
||||
marqsv3,
|
||||
(runIds) =>
|
||||
prisma.taskRun.findMany({
|
||||
select: { id: true },
|
||||
where: {
|
||||
id: {
|
||||
in: runIds,
|
||||
},
|
||||
status: {
|
||||
in: [
|
||||
"CANCELED",
|
||||
"COMPLETED_SUCCESSFULLY",
|
||||
"COMPLETED_WITH_ERRORS",
|
||||
"CRASHED",
|
||||
"SYSTEM_FAILURE",
|
||||
"INTERRUPTED",
|
||||
],
|
||||
},
|
||||
},
|
||||
}),
|
||||
{ dryRun: env.V3_MARQS_CONCURRENCY_MONITOR_ENABLED === "0", abortSignal }
|
||||
);
|
||||
|
||||
await instance.call();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
export const MARQS_RESUME_PRIORITY_TIMESTAMP_OFFSET = 31_556_952 * 1000; // 1 year
|
||||
export const MARQS_RETRY_PRIORITY_TIMESTAMP_OFFSET = 15_778_476 * 1000; // 6 months
|
||||
export const MARQS_DELAYED_REQUEUE_THRESHOLD_IN_MS = 500;
|
||||
export const MARQS_SCHEDULED_REQUEUE_AVAILABLE_AT_THRESHOLD_IN_MS = 500;
|
||||
@@ -0,0 +1,45 @@
|
||||
import { z } from "zod";
|
||||
import { singleton } from "~/utils/singleton";
|
||||
import type { ZodSubscriber } from "../utils/zodPubSub.server";
|
||||
import { ZodPubSub } from "../utils/zodPubSub.server";
|
||||
import { env } from "~/env.server";
|
||||
import { Gauge } from "prom-client";
|
||||
import { metricsRegister } from "~/metrics.server";
|
||||
|
||||
const messageCatalog = {
|
||||
CANCEL_ATTEMPT: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
backgroundWorkerId: z.string(),
|
||||
attemptId: z.string(),
|
||||
taskRunId: z.string(),
|
||||
}),
|
||||
};
|
||||
|
||||
export type DevSubscriber = ZodSubscriber<typeof messageCatalog>;
|
||||
|
||||
export const devPubSub = singleton("devPubSub", initializeDevPubSub);
|
||||
|
||||
function initializeDevPubSub() {
|
||||
const pubSub = new ZodPubSub({
|
||||
redis: {
|
||||
port: env.PUBSUB_REDIS_PORT,
|
||||
host: env.PUBSUB_REDIS_HOST,
|
||||
username: env.PUBSUB_REDIS_USERNAME,
|
||||
password: env.PUBSUB_REDIS_PASSWORD,
|
||||
tlsDisabled: env.PUBSUB_REDIS_TLS_DISABLED === "true",
|
||||
clusterMode: env.PUBSUB_REDIS_CLUSTER_MODE_ENABLED === "1",
|
||||
},
|
||||
schema: messageCatalog,
|
||||
});
|
||||
|
||||
new Gauge({
|
||||
name: "dev_pub_sub_subscribers",
|
||||
help: "Number of dev pub sub subscribers",
|
||||
collect() {
|
||||
this.set(pubSub.subscriberCount);
|
||||
},
|
||||
registers: [metricsRegister],
|
||||
});
|
||||
|
||||
return pubSub;
|
||||
}
|
||||
@@ -0,0 +1,623 @@
|
||||
import type { Context, Span } from "@opentelemetry/api";
|
||||
import { ROOT_CONTEXT, SpanKind, context, trace } from "@opentelemetry/api";
|
||||
import type {
|
||||
V3TaskRunExecution,
|
||||
TaskRunExecutionLazyAttemptPayload,
|
||||
TaskRunExecutionResult,
|
||||
TaskRunFailedExecutionResult,
|
||||
serverWebsocketMessages,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { getMaxDuration } from "@trigger.dev/core/v3/isomorphic";
|
||||
import type { ZodMessageSender } from "@trigger.dev/core/v3/zodMessageHandler";
|
||||
import type { BackgroundWorker, BackgroundWorkerTask } from "@trigger.dev/database";
|
||||
import { z } from "zod";
|
||||
import { prisma } from "~/db.server";
|
||||
import { createNewSession, disconnectSession } from "~/models/runtimeEnvironment.server";
|
||||
import { findQueueInEnvironment, sanitizeQueueName } from "~/models/taskQueue.server";
|
||||
import type { RedisClient } from "~/redis.server";
|
||||
import { createRedisClient } from "~/redis.server";
|
||||
import type { AuthenticatedEnvironment } from "~/services/apiAuth.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { marqs } from "~/v3/marqs/index.server";
|
||||
import { resolveVariablesForEnvironment } from "../environmentVariables/environmentVariablesRepository.server";
|
||||
import { FailedTaskRunService } from "../failedTaskRun.server";
|
||||
import { CancelDevSessionRunsService } from "../services/cancelDevSessionRuns.server";
|
||||
import { CompleteAttemptService } from "../services/completeAttempt.server";
|
||||
import { attributesFromAuthenticatedEnv, tracer } from "../tracer.server";
|
||||
import type { DevSubscriber } from "./devPubSub.server";
|
||||
import { devPubSub } from "./devPubSub.server";
|
||||
|
||||
const MessageBody = z.discriminatedUnion("type", [
|
||||
z.object({
|
||||
type: z.literal("EXECUTE"),
|
||||
taskIdentifier: z.string(),
|
||||
}),
|
||||
]);
|
||||
|
||||
type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTask[] };
|
||||
|
||||
export type DevQueueConsumerOptions = {
|
||||
maximumItemsPerTrace?: number;
|
||||
traceTimeoutSeconds?: number;
|
||||
ipAddress?: string;
|
||||
};
|
||||
|
||||
export class DevQueueConsumer {
|
||||
private _backgroundWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
|
||||
private _backgroundWorkerSubscriber: Map<string, DevSubscriber> = new Map();
|
||||
private _deprecatedWorkers: Map<string, BackgroundWorkerWithTasks> = new Map();
|
||||
private _enabled = false;
|
||||
private _maximumItemsPerTrace: number;
|
||||
private _traceTimeoutSeconds: number;
|
||||
private _perTraceCountdown: number | undefined;
|
||||
private _lastNewTrace: Date | undefined;
|
||||
private _currentSpanContext: Context | undefined;
|
||||
private _taskFailures: number = 0;
|
||||
private _taskSuccesses: number = 0;
|
||||
private _currentSpan: Span | undefined;
|
||||
private _endSpanInNextIteration = false;
|
||||
private _inProgressRuns: Map<string, string> = new Map(); // Keys are task run friendly IDs, values are TaskRun internal ids/queue message ids
|
||||
private _connectionLostAt?: Date;
|
||||
private _redisClient: RedisClient;
|
||||
|
||||
constructor(
|
||||
public id: string,
|
||||
public env: AuthenticatedEnvironment,
|
||||
private _sender: ZodMessageSender<typeof serverWebsocketMessages>,
|
||||
private _options: DevQueueConsumerOptions = {}
|
||||
) {
|
||||
this._traceTimeoutSeconds = _options.traceTimeoutSeconds ?? 60;
|
||||
this._maximumItemsPerTrace = _options.maximumItemsPerTrace ?? 1_000;
|
||||
this._redisClient = createRedisClient("tr:devQueueConsumer", {
|
||||
keyPrefix: "tr:devQueueConsumer:",
|
||||
...devPubSub.redisOptions,
|
||||
});
|
||||
}
|
||||
|
||||
// This method is called when a background worker is deprecated and will no longer be used unless a run is locked to it
|
||||
public async deprecateBackgroundWorker(id: string) {
|
||||
const backgroundWorker = this._backgroundWorkers.get(id);
|
||||
|
||||
if (!backgroundWorker) {
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug("[DevQueueConsumer] Deprecating background worker", {
|
||||
backgroundWorker: backgroundWorker.id,
|
||||
env: this.env.id,
|
||||
});
|
||||
|
||||
this._deprecatedWorkers.set(id, backgroundWorker);
|
||||
this._backgroundWorkers.delete(id);
|
||||
}
|
||||
|
||||
public async registerBackgroundWorker(id: string, inProgressRuns: string[] = []) {
|
||||
const backgroundWorker = await prisma.backgroundWorker.findFirst({
|
||||
where: { friendlyId: id, runtimeEnvironmentId: this.env.id },
|
||||
include: {
|
||||
tasks: true,
|
||||
},
|
||||
});
|
||||
|
||||
if (!backgroundWorker) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (this._backgroundWorkers.has(backgroundWorker.id)) {
|
||||
return;
|
||||
}
|
||||
|
||||
this._backgroundWorkers.set(backgroundWorker.id, backgroundWorker);
|
||||
|
||||
logger.debug("[DevQueueConsumer] Registered background worker", {
|
||||
backgroundWorker: backgroundWorker.id,
|
||||
inProgressRuns,
|
||||
env: this.env.id,
|
||||
});
|
||||
|
||||
const subscriber = await devPubSub.subscribe(`backgroundWorker:${backgroundWorker.id}:*`);
|
||||
|
||||
subscriber.on("CANCEL_ATTEMPT", async (message) => {
|
||||
await this._sender.send("BACKGROUND_WORKER_MESSAGE", {
|
||||
backgroundWorkerId: backgroundWorker.friendlyId,
|
||||
data: {
|
||||
type: "CANCEL_ATTEMPT",
|
||||
taskAttemptId: message.attemptId,
|
||||
taskRunId: message.taskRunId,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
this._backgroundWorkerSubscriber.set(backgroundWorker.id, subscriber);
|
||||
|
||||
for (const runId of inProgressRuns) {
|
||||
this._inProgressRuns.set(runId, runId);
|
||||
}
|
||||
|
||||
// Start reading from the queue if we haven't already
|
||||
await this.#enable();
|
||||
}
|
||||
|
||||
public async taskAttemptCompleted(
|
||||
workerId: string,
|
||||
completion: TaskRunExecutionResult,
|
||||
execution: V3TaskRunExecution
|
||||
) {
|
||||
if (completion.ok) {
|
||||
this._taskSuccesses++;
|
||||
} else {
|
||||
this._taskFailures++;
|
||||
}
|
||||
|
||||
logger.debug("[DevQueueConsumer] taskAttemptCompleted()", {
|
||||
taskRunCompletion: completion,
|
||||
execution,
|
||||
env: this.env.id,
|
||||
});
|
||||
|
||||
const service = new CompleteAttemptService();
|
||||
const result = await service.call({ completion, execution, env: this.env });
|
||||
|
||||
if (result === "COMPLETED") {
|
||||
this._inProgressRuns.delete(execution.run.id);
|
||||
}
|
||||
}
|
||||
|
||||
public async taskRunFailed(workerId: string, completion: TaskRunFailedExecutionResult) {
|
||||
this._taskFailures++;
|
||||
|
||||
logger.debug("[DevQueueConsumer] taskRunFailed()", { completion, env: this.env.id });
|
||||
|
||||
this._inProgressRuns.delete(completion.id);
|
||||
|
||||
const service = new FailedTaskRunService();
|
||||
|
||||
await service.call(completion.id, completion);
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Use `taskRunHeartbeat` instead
|
||||
*/
|
||||
public async taskHeartbeat(workerId: string, id: string) {
|
||||
logger.debug("[DevQueueConsumer] taskHeartbeat()", { id });
|
||||
|
||||
const taskRunAttempt = await prisma.taskRunAttempt.findFirst({
|
||||
where: { friendlyId: id },
|
||||
});
|
||||
|
||||
if (!taskRunAttempt) {
|
||||
return;
|
||||
}
|
||||
|
||||
await marqs?.heartbeatMessage(taskRunAttempt.taskRunId);
|
||||
}
|
||||
|
||||
public async taskRunHeartbeat(workerId: string, id: string) {
|
||||
logger.debug("[DevQueueConsumer] taskRunHeartbeat()", { id });
|
||||
|
||||
await marqs?.heartbeatMessage(id);
|
||||
}
|
||||
|
||||
public async stop(reason: string = "CLI disconnected") {
|
||||
if (!this._enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
logger.debug("[DevQueueConsumer] Stopping dev queue consumer", { env: this.env });
|
||||
|
||||
this._enabled = false;
|
||||
|
||||
// Create the session
|
||||
const session = await disconnectSession(this.env.id);
|
||||
|
||||
const runIds = Array.from(this._inProgressRuns.values());
|
||||
this._inProgressRuns.clear();
|
||||
|
||||
if (runIds.length > 0) {
|
||||
await CancelDevSessionRunsService.enqueue(
|
||||
{
|
||||
runIds,
|
||||
cancelledAt: new Date(),
|
||||
reason,
|
||||
cancelledSessionId: session?.id,
|
||||
},
|
||||
new Date(Date.now() + 1000 * 10) // 10 seconds from now
|
||||
);
|
||||
}
|
||||
|
||||
// We need to unsubscribe from the background worker channels
|
||||
for (const [id, subscriber] of this._backgroundWorkerSubscriber) {
|
||||
logger.debug("Unsubscribing from background worker channel", { id });
|
||||
|
||||
await subscriber.stopListening();
|
||||
this._backgroundWorkerSubscriber.delete(id);
|
||||
|
||||
logger.debug("Unsubscribed from background worker channel", { id });
|
||||
}
|
||||
|
||||
// We need to end the current span
|
||||
if (this._currentSpan) {
|
||||
this._currentSpan.end();
|
||||
}
|
||||
}
|
||||
|
||||
async #enable() {
|
||||
if (this._enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
await this._redisClient.set(`connection:${this.env.id}`, this.id, "EX", 60 * 60 * 24); // 24 hours
|
||||
|
||||
this._enabled = true;
|
||||
// Create the session
|
||||
await createNewSession(this.env, this._options.ipAddress ?? "unknown");
|
||||
|
||||
this._perTraceCountdown = this._options.maximumItemsPerTrace;
|
||||
this._lastNewTrace = new Date();
|
||||
this._taskFailures = 0;
|
||||
this._taskSuccesses = 0;
|
||||
|
||||
this.#doWork().finally(() => {});
|
||||
}
|
||||
|
||||
async #doWork() {
|
||||
if (!this._enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
const canSendMessage = await this._sender.validateCanSendMessage();
|
||||
|
||||
if (!canSendMessage) {
|
||||
this._connectionLostAt ??= new Date();
|
||||
|
||||
if (Date.now() - this._connectionLostAt.getTime() > 60 * 1000) {
|
||||
logger.debug("Connection lost for more than 60 seconds, stopping the consumer", {
|
||||
env: this.env,
|
||||
});
|
||||
|
||||
await this.stop("Connection lost for more than 60 seconds");
|
||||
return;
|
||||
}
|
||||
|
||||
setTimeout(() => this.#doWork(), 1000);
|
||||
return;
|
||||
}
|
||||
|
||||
this._connectionLostAt = undefined;
|
||||
|
||||
const currentConnection = await this._redisClient.get(`connection:${this.env.id}`);
|
||||
|
||||
if (currentConnection && currentConnection !== this.id) {
|
||||
logger.debug("Another connection is active, stopping the consumer", {
|
||||
currentConnection,
|
||||
env: this.env,
|
||||
});
|
||||
|
||||
await this.stop("Another connection is active");
|
||||
return;
|
||||
}
|
||||
|
||||
// Check if the trace has expired
|
||||
if (
|
||||
this._perTraceCountdown === 0 ||
|
||||
Date.now() - this._lastNewTrace!.getTime() > this._traceTimeoutSeconds * 1000 ||
|
||||
this._currentSpanContext === undefined ||
|
||||
this._endSpanInNextIteration
|
||||
) {
|
||||
if (this._currentSpan) {
|
||||
this._currentSpan.setAttribute("tasks.period.failures", this._taskFailures);
|
||||
this._currentSpan.setAttribute("tasks.period.successes", this._taskSuccesses);
|
||||
|
||||
logger.debug("Ending DevQueueConsumer.doWork() trace", {
|
||||
isRecording: this._currentSpan.isRecording(),
|
||||
});
|
||||
|
||||
this._currentSpan.end();
|
||||
}
|
||||
|
||||
// Create a new trace
|
||||
this._currentSpan = tracer.startSpan(
|
||||
"DevQueueConsumer.doWork()",
|
||||
{
|
||||
kind: SpanKind.CONSUMER,
|
||||
attributes: {
|
||||
...attributesFromAuthenticatedEnv(this.env),
|
||||
},
|
||||
},
|
||||
ROOT_CONTEXT
|
||||
);
|
||||
|
||||
// Get the span trace context
|
||||
this._currentSpanContext = trace.setSpan(ROOT_CONTEXT, this._currentSpan);
|
||||
|
||||
this._perTraceCountdown = this._options.maximumItemsPerTrace;
|
||||
this._lastNewTrace = new Date();
|
||||
this._taskFailures = 0;
|
||||
this._taskSuccesses = 0;
|
||||
this._endSpanInNextIteration = false;
|
||||
}
|
||||
|
||||
return context.with(this._currentSpanContext ?? ROOT_CONTEXT, async () => {
|
||||
await this.#doWorkInternal();
|
||||
this._perTraceCountdown = this._perTraceCountdown! - 1;
|
||||
});
|
||||
}
|
||||
|
||||
async #doWorkInternal() {
|
||||
// Attempt to dequeue a message from the environment's queue
|
||||
// If no message is available, reschedule the worker to run again in 1 second
|
||||
// If a message is available, find the BackgroundWorkerTask that matches the message's taskIdentifier
|
||||
// If no matching task is found, nack the message and reschedule the worker to run again in 1 second
|
||||
// If the matching task is found, create the task attempt and lock the task run, then send the task run to the client
|
||||
// Store the message as a processing message
|
||||
// If the websocket connection disconnects before the task run is completed, nack the message
|
||||
// When the task run completes, ack the message
|
||||
// Using a heartbeat mechanism, if the client keeps responding with a heartbeat, we'll keep the message processing and increase the visibility timeout.
|
||||
|
||||
const message = await marqs?.dequeueMessageInEnv(this.env);
|
||||
|
||||
if (!message) {
|
||||
setTimeout(() => this.#doWork(), 1000);
|
||||
return;
|
||||
}
|
||||
|
||||
const dequeuedStart = Date.now();
|
||||
|
||||
const messageBody = MessageBody.safeParse(message.data);
|
||||
|
||||
if (!messageBody.success) {
|
||||
logger.error("Failed to parse message", {
|
||||
queueMessage: message.data,
|
||||
error: messageBody.error,
|
||||
env: this.env,
|
||||
});
|
||||
|
||||
await marqs?.acknowledgeMessage(
|
||||
message.messageId,
|
||||
"Failed to parse message.data with MessageBody schema in DevQueueConsumer"
|
||||
);
|
||||
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const existingTaskRun = await prisma.taskRun.findFirst({
|
||||
where: {
|
||||
id: message.messageId,
|
||||
},
|
||||
});
|
||||
|
||||
if (!existingTaskRun) {
|
||||
logger.debug("Failed to find existing task run, acking", {
|
||||
messageId: message.messageId,
|
||||
});
|
||||
|
||||
await marqs?.acknowledgeMessage(
|
||||
message.messageId,
|
||||
"Failed to find task run in DevQueueConsumer"
|
||||
);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const backgroundWorker = existingTaskRun.lockedToVersionId
|
||||
? (this._deprecatedWorkers.get(existingTaskRun.lockedToVersionId) ??
|
||||
this._backgroundWorkers.get(existingTaskRun.lockedToVersionId))
|
||||
: this.#getLatestBackgroundWorker();
|
||||
|
||||
if (!backgroundWorker) {
|
||||
logger.debug("Failed to find background worker, acking", {
|
||||
messageId: message.messageId,
|
||||
lockedToVersionId: existingTaskRun.lockedToVersionId,
|
||||
deprecatedWorkers: Array.from(this._deprecatedWorkers.keys()),
|
||||
backgroundWorkers: Array.from(this._backgroundWorkers.keys()),
|
||||
latestWorker: this.#getLatestBackgroundWorker(),
|
||||
});
|
||||
|
||||
await marqs?.acknowledgeMessage(
|
||||
message.messageId,
|
||||
"Failed to find background worker in DevQueueConsumer"
|
||||
);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const backgroundTask = backgroundWorker.tasks.find(
|
||||
(task) => task.slug === existingTaskRun.taskIdentifier
|
||||
);
|
||||
|
||||
if (!backgroundTask) {
|
||||
logger.warn("No matching background task found for task run", {
|
||||
taskRun: existingTaskRun.id,
|
||||
taskIdentifier: existingTaskRun.taskIdentifier,
|
||||
backgroundWorker: backgroundWorker.id,
|
||||
taskSlugs: backgroundWorker.tasks.map((task) => task.slug),
|
||||
});
|
||||
|
||||
await marqs?.acknowledgeMessage(
|
||||
message.messageId,
|
||||
"No matching background task found in DevQueueConsumer"
|
||||
);
|
||||
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const lockedAt = new Date();
|
||||
const startedAt = existingTaskRun.startedAt ?? new Date();
|
||||
|
||||
const lockedTaskRun = await prisma.taskRun.update({
|
||||
where: {
|
||||
id: message.messageId,
|
||||
},
|
||||
data: {
|
||||
lockedAt,
|
||||
lockedById: backgroundTask.id,
|
||||
status: "EXECUTING",
|
||||
lockedToVersionId: backgroundWorker.id,
|
||||
taskVersion: backgroundWorker.version,
|
||||
sdkVersion: backgroundWorker.sdkVersion,
|
||||
cliVersion: backgroundWorker.cliVersion,
|
||||
startedAt,
|
||||
maxDurationInSeconds: getMaxDuration(
|
||||
existingTaskRun.maxDurationInSeconds,
|
||||
backgroundTask.maxDurationInSeconds
|
||||
),
|
||||
},
|
||||
});
|
||||
|
||||
if (!lockedTaskRun) {
|
||||
logger.warn("Failed to lock task run", {
|
||||
taskRun: existingTaskRun.id,
|
||||
taskIdentifier: existingTaskRun.taskIdentifier,
|
||||
backgroundWorker: backgroundWorker.id,
|
||||
messageId: message.messageId,
|
||||
});
|
||||
|
||||
await marqs?.acknowledgeMessage(
|
||||
message.messageId,
|
||||
"Failed to lock task run in DevQueueConsumer"
|
||||
);
|
||||
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const queue = await findQueueInEnvironment(
|
||||
lockedTaskRun.queue,
|
||||
this.env.id,
|
||||
backgroundTask.id,
|
||||
backgroundTask
|
||||
);
|
||||
|
||||
if (!queue) {
|
||||
logger.debug("[DevQueueConsumer] Failed to find queue", {
|
||||
queueName: lockedTaskRun.queue,
|
||||
sanitizedName: sanitizeQueueName(lockedTaskRun.queue),
|
||||
taskRun: lockedTaskRun.id,
|
||||
messageId: message.messageId,
|
||||
});
|
||||
|
||||
await marqs?.nackMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 1000);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!this._enabled) {
|
||||
logger.debug("Dev queue consumer is disabled", { env: this.env, queueMessage: message });
|
||||
|
||||
await marqs?.nackMessage(message.messageId);
|
||||
return;
|
||||
}
|
||||
|
||||
const variables = await resolveVariablesForEnvironment(this.env);
|
||||
|
||||
if (backgroundWorker.supportsLazyAttempts) {
|
||||
const payload: TaskRunExecutionLazyAttemptPayload = {
|
||||
traceContext: lockedTaskRun.traceContext as Record<string, unknown>,
|
||||
environment: variables.reduce((acc: Record<string, string>, curr) => {
|
||||
acc[curr.key] = curr.value;
|
||||
return acc;
|
||||
}, {}),
|
||||
runId: lockedTaskRun.friendlyId,
|
||||
messageId: lockedTaskRun.id,
|
||||
isTest: lockedTaskRun.isTest,
|
||||
isReplay: !!lockedTaskRun.replayedFromTaskRunFriendlyId,
|
||||
metrics: [
|
||||
{
|
||||
name: "start",
|
||||
event: "dequeue",
|
||||
timestamp: dequeuedStart,
|
||||
duration: Date.now() - dequeuedStart,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
try {
|
||||
await this._sender.send("BACKGROUND_WORKER_MESSAGE", {
|
||||
backgroundWorkerId: backgroundWorker.friendlyId,
|
||||
data: {
|
||||
type: "EXECUTE_RUN_LAZY_ATTEMPT",
|
||||
payload,
|
||||
},
|
||||
});
|
||||
|
||||
logger.debug("Executing the run", {
|
||||
messageId: message.messageId,
|
||||
});
|
||||
|
||||
this._inProgressRuns.set(lockedTaskRun.friendlyId, message.messageId);
|
||||
} catch (e) {
|
||||
if (e instanceof Error) {
|
||||
this._currentSpan?.recordException(e);
|
||||
} else {
|
||||
this._currentSpan?.recordException(new Error(String(e)));
|
||||
}
|
||||
|
||||
this._endSpanInNextIteration = true;
|
||||
|
||||
// We now need to unlock the task run and delete the task run attempt
|
||||
await prisma.$transaction([
|
||||
prisma.taskRun.update({
|
||||
where: {
|
||||
id: lockedTaskRun.id,
|
||||
},
|
||||
data: {
|
||||
lockedAt: null,
|
||||
lockedById: null,
|
||||
status: "PENDING",
|
||||
startedAt: existingTaskRun.startedAt,
|
||||
},
|
||||
}),
|
||||
]);
|
||||
|
||||
this._inProgressRuns.delete(lockedTaskRun.friendlyId);
|
||||
|
||||
// Finally we need to nack the message so it can be retried
|
||||
await marqs?.nackMessage(message.messageId);
|
||||
} finally {
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
}
|
||||
} else {
|
||||
logger.debug("We no longer support non-lazy attempts, aborting this run", {
|
||||
messageId: message.messageId,
|
||||
backgroundWorker,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(
|
||||
message.messageId,
|
||||
"Non-lazy attempts are no longer supported in DevQueueConsumer"
|
||||
);
|
||||
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
}
|
||||
}
|
||||
|
||||
// Get the latest background worker based on the version.
|
||||
// Versions are in the format of 20240101.1 and 20240101.2, or even 20240101.10, 20240101.11, etc.
|
||||
#getLatestBackgroundWorker() {
|
||||
const workers = Array.from(this._backgroundWorkers.values());
|
||||
|
||||
if (workers.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
return workers.reduce((acc, curr) => {
|
||||
const accParts = acc.version.split(".").map(Number);
|
||||
const currParts = curr.version.split(".").map(Number);
|
||||
|
||||
// Compare the major part
|
||||
if (accParts[0] < currParts[0]) {
|
||||
return curr;
|
||||
} else if (accParts[0] > currParts[0]) {
|
||||
return acc;
|
||||
}
|
||||
|
||||
// Compare the minor part (assuming all versions have two parts)
|
||||
if (accParts[1] < currParts[1]) {
|
||||
return curr;
|
||||
} else {
|
||||
return acc;
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,598 @@
|
||||
import type { Cache as UnkeyCache } from "@unkey/cache";
|
||||
import { createCache, DefaultStatefulContext, Namespace } from "@unkey/cache";
|
||||
import { createLRUMemoryStore } from "@internal/cache";
|
||||
import { randomUUID } from "crypto";
|
||||
import type { Redis } from "ioredis";
|
||||
import type { EnvQueues, MarQSFairDequeueStrategy, MarQSKeyProducer } from "./types";
|
||||
import seedrandom from "seedrandom";
|
||||
import type { Tracer } from "@opentelemetry/api";
|
||||
import { startSpan } from "../tracing.server";
|
||||
|
||||
export type FairDequeuingStrategyBiases = {
|
||||
/**
|
||||
* How much to bias towards environments with higher concurrency limits
|
||||
* 0 = no bias, 1 = full bias based on limit differences
|
||||
*/
|
||||
concurrencyLimitBias: number;
|
||||
|
||||
/**
|
||||
* How much to bias towards environments with more available capacity
|
||||
* 0 = no bias, 1 = full bias based on available capacity
|
||||
*/
|
||||
availableCapacityBias: number;
|
||||
|
||||
/**
|
||||
* Controls randomization of queue ordering within environments
|
||||
* 0 = strict age-based ordering (oldest first)
|
||||
* 1 = completely random ordering
|
||||
* Values between 0-1 blend between age-based and random ordering
|
||||
*/
|
||||
queueAgeRandomization: number;
|
||||
};
|
||||
|
||||
export type FairDequeuingStrategyOptions = {
|
||||
redis: Redis;
|
||||
keys: MarQSKeyProducer;
|
||||
defaultEnvConcurrency: number;
|
||||
parentQueueLimit: number;
|
||||
tracer: Tracer;
|
||||
seed?: string;
|
||||
/**
|
||||
* Configure biasing for environment shuffling
|
||||
* If not provided, no biasing will be applied (completely random shuffling)
|
||||
*/
|
||||
biases?: FairDequeuingStrategyBiases;
|
||||
reuseSnapshotCount?: number;
|
||||
maximumEnvCount?: number;
|
||||
/**
|
||||
* Maximum number of queues to process per environment
|
||||
* If not provided, all queues in an environment will be processed
|
||||
*/
|
||||
maximumQueuePerEnvCount?: number;
|
||||
};
|
||||
|
||||
type FairQueueConcurrency = {
|
||||
current: number;
|
||||
limit: number;
|
||||
reserve: number;
|
||||
};
|
||||
|
||||
type FairQueue = { id: string; age: number; org: string; env: string };
|
||||
|
||||
type FairQueueSnapshot = {
|
||||
id: string;
|
||||
envs: Record<string, { concurrency: FairQueueConcurrency }>;
|
||||
queues: Array<FairQueue>;
|
||||
};
|
||||
|
||||
type WeightedEnv = {
|
||||
envId: string;
|
||||
weight: number;
|
||||
};
|
||||
|
||||
type WeightedQueue = {
|
||||
queue: FairQueue;
|
||||
weight: number;
|
||||
};
|
||||
|
||||
const emptyFairQueueSnapshot: FairQueueSnapshot = {
|
||||
id: "empty",
|
||||
envs: {},
|
||||
queues: [],
|
||||
};
|
||||
|
||||
const defaultBiases: FairDequeuingStrategyBiases = {
|
||||
concurrencyLimitBias: 0,
|
||||
availableCapacityBias: 0,
|
||||
queueAgeRandomization: 0, // Default to completely age-based ordering
|
||||
};
|
||||
|
||||
export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
|
||||
private _cache: UnkeyCache<{
|
||||
concurrencyLimit: number;
|
||||
}>;
|
||||
|
||||
private _rng: seedrandom.PRNG;
|
||||
private _reusedSnapshotForConsumer: Map<
|
||||
string,
|
||||
{ snapshot: FairQueueSnapshot; reuseCount: number }
|
||||
> = new Map();
|
||||
|
||||
constructor(private options: FairDequeuingStrategyOptions) {
|
||||
const ctx = new DefaultStatefulContext();
|
||||
const memory = createLRUMemoryStore(500);
|
||||
|
||||
this._cache = createCache({
|
||||
concurrencyLimit: new Namespace<number>(ctx, {
|
||||
stores: [memory],
|
||||
fresh: 60_000, // The time in milliseconds that a value is considered fresh. Cache hits within this time will return the cached value.
|
||||
stale: 180_000, // The time in milliseconds that a value is considered stale. Cache hits within this time will return the cached value and trigger a background refresh.
|
||||
}),
|
||||
});
|
||||
|
||||
this._rng = seedrandom(options.seed);
|
||||
}
|
||||
|
||||
async distributeFairQueuesFromParentQueue(
|
||||
parentQueue: string,
|
||||
consumerId: string
|
||||
): Promise<Array<EnvQueues>> {
|
||||
return await startSpan(
|
||||
this.options.tracer,
|
||||
"distributeFairQueuesFromParentQueue",
|
||||
async (span) => {
|
||||
span.setAttribute("consumer_id", consumerId);
|
||||
span.setAttribute("parent_queue", parentQueue);
|
||||
|
||||
const snapshot = await this.#createQueueSnapshot(parentQueue, consumerId);
|
||||
|
||||
span.setAttributes({
|
||||
snapshot_env_count: Object.keys(snapshot.envs).length,
|
||||
snapshot_queue_count: snapshot.queues.length,
|
||||
});
|
||||
|
||||
const queues = snapshot.queues;
|
||||
|
||||
if (queues.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const envQueues = this.#shuffleQueuesByEnv(snapshot);
|
||||
|
||||
span.setAttribute(
|
||||
"shuffled_queue_count",
|
||||
envQueues.reduce((sum, env) => sum + env.queues.length, 0)
|
||||
);
|
||||
|
||||
if (envQueues[0]?.queues[0]) {
|
||||
span.setAttribute("winning_env", envQueues[0].envId);
|
||||
span.setAttribute(
|
||||
"winning_org",
|
||||
this.options.keys.orgIdFromQueue(envQueues[0].queues[0])
|
||||
);
|
||||
}
|
||||
|
||||
return envQueues;
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#shuffleQueuesByEnv(snapshot: FairQueueSnapshot): Array<EnvQueues> {
|
||||
const envs = Object.keys(snapshot.envs);
|
||||
const biases = this.options.biases ?? defaultBiases;
|
||||
|
||||
if (biases.concurrencyLimitBias === 0 && biases.availableCapacityBias === 0) {
|
||||
const shuffledEnvs = this.#shuffle(envs);
|
||||
return this.#orderQueuesByEnvs(shuffledEnvs, snapshot);
|
||||
}
|
||||
|
||||
// Find the maximum concurrency limit for normalization
|
||||
const maxLimit = Math.max(...envs.map((envId) => snapshot.envs[envId].concurrency.limit));
|
||||
|
||||
// Calculate weights for each environment
|
||||
const weightedEnvs: WeightedEnv[] = envs.map((envId) => {
|
||||
const env = snapshot.envs[envId];
|
||||
|
||||
// Start with base weight of 1
|
||||
let weight = 1;
|
||||
|
||||
// Add normalized concurrency limit bias if configured
|
||||
if (biases.concurrencyLimitBias > 0) {
|
||||
const normalizedLimit = env.concurrency.limit / maxLimit;
|
||||
// Square or cube the bias to make it more pronounced at higher values
|
||||
weight *= 1 + Math.pow(normalizedLimit * biases.concurrencyLimitBias, 2);
|
||||
}
|
||||
|
||||
// Add available capacity bias if configured
|
||||
if (biases.availableCapacityBias > 0) {
|
||||
const usedCapacityPercentage = env.concurrency.current / env.concurrency.limit;
|
||||
const availableCapacityBonus = 1 - usedCapacityPercentage;
|
||||
// Square or cube the bias to make it more pronounced at higher values
|
||||
weight *= 1 + Math.pow(availableCapacityBonus * biases.availableCapacityBias, 2);
|
||||
}
|
||||
|
||||
return { envId, weight };
|
||||
});
|
||||
|
||||
const shuffledEnvs = this.#weightedShuffle(weightedEnvs);
|
||||
return this.#orderQueuesByEnvs(shuffledEnvs, snapshot);
|
||||
}
|
||||
|
||||
#weightedShuffle(weightedItems: WeightedEnv[]): string[] {
|
||||
const totalWeight = weightedItems.reduce((sum, item) => sum + item.weight, 0);
|
||||
const result: string[] = [];
|
||||
const items = [...weightedItems];
|
||||
|
||||
while (items.length > 0) {
|
||||
let random = this._rng() * totalWeight;
|
||||
let index = 0;
|
||||
|
||||
// Find item based on weighted random selection
|
||||
while (random > 0 && index < items.length) {
|
||||
random -= items[index].weight;
|
||||
index++;
|
||||
}
|
||||
index = Math.max(0, index - 1);
|
||||
|
||||
// Add selected item to result and remove from items
|
||||
result.push(items[index].envId);
|
||||
items.splice(index, 1);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
#orderQueuesByEnvs(envs: string[], snapshot: FairQueueSnapshot): Array<EnvQueues> {
|
||||
const queuesByEnv = snapshot.queues.reduce(
|
||||
(acc, queue) => {
|
||||
if (!acc[queue.env]) {
|
||||
acc[queue.env] = [];
|
||||
}
|
||||
acc[queue.env].push(queue);
|
||||
return acc;
|
||||
},
|
||||
{} as Record<string, Array<FairQueue>>
|
||||
);
|
||||
|
||||
return envs.reduce((acc, envId) => {
|
||||
if (queuesByEnv[envId]) {
|
||||
// Get ordered queues for this env
|
||||
const orderedQueues = this.#weightedRandomQueueOrder(queuesByEnv[envId]);
|
||||
|
||||
// Apply queue limit if maximumQueuePerEnvCount is set
|
||||
const limitedQueues = this.options.maximumQueuePerEnvCount
|
||||
? orderedQueues.slice(0, this.options.maximumQueuePerEnvCount)
|
||||
: orderedQueues;
|
||||
|
||||
// Only add the env if it has queues
|
||||
if (limitedQueues.length > 0) {
|
||||
acc.push({
|
||||
envId,
|
||||
queues: limitedQueues.map((queue) => queue.id),
|
||||
});
|
||||
}
|
||||
}
|
||||
return acc;
|
||||
}, [] as Array<EnvQueues>);
|
||||
}
|
||||
|
||||
#weightedRandomQueueOrder(queues: FairQueue[]): FairQueue[] {
|
||||
if (queues.length <= 1) return queues;
|
||||
|
||||
const biases = this.options.biases ?? defaultBiases;
|
||||
|
||||
// When queueAgeRandomization is 0, use strict age-based ordering
|
||||
if (biases.queueAgeRandomization === 0) {
|
||||
return [...queues].sort((a, b) => b.age - a.age);
|
||||
}
|
||||
|
||||
// Find the maximum age for normalization
|
||||
const maxAge = Math.max(...queues.map((q) => q.age));
|
||||
|
||||
// Calculate weights for each queue
|
||||
const weightedQueues: WeightedQueue[] = queues.map((queue) => {
|
||||
// Normalize age to be between 0 and 1
|
||||
const normalizedAge = queue.age / maxAge;
|
||||
|
||||
// Calculate weight: combine base weight with configurable age influence
|
||||
const baseWeight = 1;
|
||||
const weight = baseWeight + normalizedAge * biases.queueAgeRandomization;
|
||||
|
||||
return { queue, weight };
|
||||
});
|
||||
|
||||
// Perform weighted random selection for ordering
|
||||
const result: FairQueue[] = [];
|
||||
let remainingQueues = [...weightedQueues];
|
||||
let totalWeight = remainingQueues.reduce((sum, wq) => sum + wq.weight, 0);
|
||||
|
||||
while (remainingQueues.length > 0) {
|
||||
let random = this._rng() * totalWeight;
|
||||
let index = 0;
|
||||
|
||||
// Find queue based on weighted random selection
|
||||
while (random > 0 && index < remainingQueues.length) {
|
||||
random -= remainingQueues[index].weight;
|
||||
index++;
|
||||
}
|
||||
index = Math.max(0, index - 1);
|
||||
|
||||
// Add selected queue to result and remove from remaining
|
||||
result.push(remainingQueues[index].queue);
|
||||
totalWeight -= remainingQueues[index].weight;
|
||||
remainingQueues.splice(index, 1);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
#shuffle<T>(array: Array<T>): Array<T> {
|
||||
let currentIndex = array.length;
|
||||
let temporaryValue;
|
||||
let randomIndex;
|
||||
|
||||
const newArray = [...array];
|
||||
|
||||
while (currentIndex !== 0) {
|
||||
randomIndex = Math.floor(this._rng() * currentIndex);
|
||||
currentIndex -= 1;
|
||||
|
||||
temporaryValue = newArray[currentIndex];
|
||||
newArray[currentIndex] = newArray[randomIndex];
|
||||
newArray[randomIndex] = temporaryValue;
|
||||
}
|
||||
|
||||
return newArray;
|
||||
}
|
||||
|
||||
async #createQueueSnapshot(parentQueue: string, consumerId: string): Promise<FairQueueSnapshot> {
|
||||
return await startSpan(this.options.tracer, "createQueueSnapshot", async (span) => {
|
||||
span.setAttribute("consumer_id", consumerId);
|
||||
span.setAttribute("parent_queue", parentQueue);
|
||||
|
||||
if (
|
||||
typeof this.options.reuseSnapshotCount === "number" &&
|
||||
this.options.reuseSnapshotCount > 0
|
||||
) {
|
||||
const key = `${parentQueue}:${consumerId}`;
|
||||
const reusedSnapshot = this._reusedSnapshotForConsumer.get(key);
|
||||
|
||||
if (reusedSnapshot) {
|
||||
if (reusedSnapshot.reuseCount < this.options.reuseSnapshotCount) {
|
||||
span.setAttribute("reused_snapshot", true);
|
||||
|
||||
this._reusedSnapshotForConsumer.set(key, {
|
||||
snapshot: reusedSnapshot.snapshot,
|
||||
reuseCount: reusedSnapshot.reuseCount + 1,
|
||||
});
|
||||
|
||||
return reusedSnapshot.snapshot;
|
||||
} else {
|
||||
this._reusedSnapshotForConsumer.delete(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
span.setAttribute("reused_snapshot", false);
|
||||
|
||||
const now = Date.now();
|
||||
|
||||
let queues = await this.#allChildQueuesByScore(parentQueue, consumerId, now);
|
||||
|
||||
span.setAttribute("parent_queue_count", queues.length);
|
||||
|
||||
if (queues.length === 0) {
|
||||
return emptyFairQueueSnapshot;
|
||||
}
|
||||
|
||||
// Apply env selection if maximumEnvCount is specified
|
||||
let selectedEnvIds: Set<string>;
|
||||
if (this.options.maximumEnvCount && this.options.maximumEnvCount > 0) {
|
||||
selectedEnvIds = this.#selectTopEnvs(queues, this.options.maximumEnvCount);
|
||||
// Filter queues to only include selected envs
|
||||
queues = queues.filter((queue) => selectedEnvIds.has(queue.env));
|
||||
|
||||
span.setAttribute("selected_env_count", selectedEnvIds.size);
|
||||
}
|
||||
|
||||
span.setAttribute("selected_queue_count", queues.length);
|
||||
|
||||
const envIds = new Set<string>();
|
||||
|
||||
for (const queue of queues) {
|
||||
envIds.add(queue.env);
|
||||
}
|
||||
|
||||
const envs = await Promise.all(
|
||||
Array.from(envIds).map(async (envId) => {
|
||||
return { id: envId, concurrency: await this.#getEnvConcurrency(envId) };
|
||||
})
|
||||
);
|
||||
|
||||
const envsAtFullConcurrency = envs.filter(
|
||||
(env) => env.concurrency.current >= env.concurrency.limit + env.concurrency.reserve
|
||||
);
|
||||
|
||||
const envIdsAtFullConcurrency = new Set(envsAtFullConcurrency.map((env) => env.id));
|
||||
|
||||
const envsSnapshot = envs.reduce(
|
||||
(acc, env) => {
|
||||
if (!envIdsAtFullConcurrency.has(env.id)) {
|
||||
acc[env.id] = env;
|
||||
}
|
||||
return acc;
|
||||
},
|
||||
{} as Record<string, { concurrency: FairQueueConcurrency }>
|
||||
);
|
||||
|
||||
span.setAttributes({
|
||||
env_count: envs.length,
|
||||
envs_at_full_concurrency_count: envsAtFullConcurrency.length,
|
||||
});
|
||||
|
||||
const queuesSnapshot = queues.filter((queue) => !envIdsAtFullConcurrency.has(queue.env));
|
||||
|
||||
const snapshot = {
|
||||
id: randomUUID(),
|
||||
envs: envsSnapshot,
|
||||
queues: queuesSnapshot,
|
||||
};
|
||||
|
||||
if (
|
||||
typeof this.options.reuseSnapshotCount === "number" &&
|
||||
this.options.reuseSnapshotCount > 0
|
||||
) {
|
||||
this._reusedSnapshotForConsumer.set(`${parentQueue}:${consumerId}`, {
|
||||
snapshot,
|
||||
reuseCount: 0,
|
||||
});
|
||||
}
|
||||
|
||||
return snapshot;
|
||||
});
|
||||
}
|
||||
|
||||
#selectTopEnvs(queues: FairQueue[], maximumEnvCount: number): Set<string> {
|
||||
// Group queues by env
|
||||
const queuesByEnv = queues.reduce(
|
||||
(acc, queue) => {
|
||||
if (!acc[queue.env]) {
|
||||
acc[queue.env] = [];
|
||||
}
|
||||
acc[queue.env].push(queue);
|
||||
return acc;
|
||||
},
|
||||
{} as Record<string, FairQueue[]>
|
||||
);
|
||||
|
||||
// Calculate average age for each env
|
||||
const envAverageAges = Object.entries(queuesByEnv).map(([envId, envQueues]) => {
|
||||
const averageAge = envQueues.reduce((sum, q) => sum + q.age, 0) / envQueues.length;
|
||||
return { envId, averageAge };
|
||||
});
|
||||
|
||||
// Perform weighted shuffle based on average ages
|
||||
const maxAge = Math.max(...envAverageAges.map((e) => e.averageAge));
|
||||
const weightedEnvs = envAverageAges.map((env) => ({
|
||||
envId: env.envId,
|
||||
weight: env.averageAge / maxAge, // Normalize weights
|
||||
}));
|
||||
|
||||
// Select top N envs using weighted shuffle
|
||||
const selectedEnvs = new Set<string>();
|
||||
let remainingEnvs = [...weightedEnvs];
|
||||
let totalWeight = remainingEnvs.reduce((sum, env) => sum + env.weight, 0);
|
||||
|
||||
while (selectedEnvs.size < maximumEnvCount && remainingEnvs.length > 0) {
|
||||
let random = this._rng() * totalWeight;
|
||||
let index = 0;
|
||||
|
||||
while (random > 0 && index < remainingEnvs.length) {
|
||||
random -= remainingEnvs[index].weight;
|
||||
index++;
|
||||
}
|
||||
index = Math.max(0, index - 1);
|
||||
|
||||
selectedEnvs.add(remainingEnvs[index].envId);
|
||||
totalWeight -= remainingEnvs[index].weight;
|
||||
remainingEnvs.splice(index, 1);
|
||||
}
|
||||
|
||||
return selectedEnvs;
|
||||
}
|
||||
|
||||
async #getEnvConcurrency(envId: string): Promise<FairQueueConcurrency> {
|
||||
return await startSpan(this.options.tracer, "getEnvConcurrency", async (span) => {
|
||||
span.setAttribute("env_id", envId);
|
||||
|
||||
const [currentValue, limitValue, reserveValue] = await Promise.all([
|
||||
this.#getEnvCurrentConcurrency(envId),
|
||||
this.#getEnvConcurrencyLimit(envId),
|
||||
this.#getEnvReserveConcurrency(envId),
|
||||
]);
|
||||
|
||||
span.setAttribute("current_value", currentValue);
|
||||
span.setAttribute("limit_value", limitValue);
|
||||
span.setAttribute("reserve_value", reserveValue);
|
||||
|
||||
return { current: currentValue, limit: limitValue, reserve: reserveValue };
|
||||
});
|
||||
}
|
||||
|
||||
async #allChildQueuesByScore(
|
||||
parentQueue: string,
|
||||
consumerId: string,
|
||||
now: number
|
||||
): Promise<Array<FairQueue>> {
|
||||
return await startSpan(this.options.tracer, "allChildQueuesByScore", async (span) => {
|
||||
span.setAttribute("consumer_id", consumerId);
|
||||
span.setAttribute("parent_queue", parentQueue);
|
||||
|
||||
const valuesWithScores = await this.options.redis.zrangebyscore(
|
||||
parentQueue,
|
||||
"-inf",
|
||||
now,
|
||||
"WITHSCORES",
|
||||
"LIMIT",
|
||||
0,
|
||||
this.options.parentQueueLimit
|
||||
);
|
||||
|
||||
const result: Array<FairQueue> = [];
|
||||
|
||||
for (let i = 0; i < valuesWithScores.length; i += 2) {
|
||||
result.push({
|
||||
id: valuesWithScores[i],
|
||||
age: now - Number(valuesWithScores[i + 1]),
|
||||
env: this.options.keys.envIdFromQueue(valuesWithScores[i]),
|
||||
org: this.options.keys.orgIdFromQueue(valuesWithScores[i]),
|
||||
});
|
||||
}
|
||||
|
||||
span.setAttribute("queue_count", result.length);
|
||||
|
||||
if (result.length === this.options.parentQueueLimit) {
|
||||
span.setAttribute("parent_queue_limit_reached", true);
|
||||
}
|
||||
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
async #getEnvConcurrencyLimit(envId: string) {
|
||||
return await startSpan(this.options.tracer, "getEnvConcurrencyLimit", async (span) => {
|
||||
span.setAttribute("env_id", envId);
|
||||
|
||||
const key = this.options.keys.envConcurrencyLimitKey(envId);
|
||||
|
||||
const result = await this._cache.concurrencyLimit.swr(key, async () => {
|
||||
const value = await this.options.redis.get(key);
|
||||
|
||||
if (!value) {
|
||||
return this.options.defaultEnvConcurrency;
|
||||
}
|
||||
|
||||
return Number(value);
|
||||
});
|
||||
|
||||
return result.val ?? this.options.defaultEnvConcurrency;
|
||||
});
|
||||
}
|
||||
|
||||
async #getEnvCurrentConcurrency(envId: string) {
|
||||
return await startSpan(this.options.tracer, "getEnvCurrentConcurrency", async (span) => {
|
||||
span.setAttribute("env_id", envId);
|
||||
|
||||
const key = this.options.keys.envCurrentConcurrencyKey(envId);
|
||||
|
||||
const result = await this.options.redis.scard(key);
|
||||
|
||||
span.setAttribute("current_value", result);
|
||||
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
async #getEnvReserveConcurrency(envId: string) {
|
||||
return await startSpan(this.options.tracer, "getEnvReserveConcurrency", async (span) => {
|
||||
span.setAttribute("env_id", envId);
|
||||
|
||||
const key = this.options.keys.envReserveConcurrencyKey(envId);
|
||||
|
||||
const result = await this.options.redis.scard(key);
|
||||
|
||||
span.setAttribute("current_value", result);
|
||||
|
||||
return result;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export class NoopFairDequeuingStrategy implements MarQSFairDequeueStrategy {
|
||||
async distributeFairQueuesFromParentQueue(
|
||||
parentQueue: string,
|
||||
consumerId: string
|
||||
): Promise<Array<EnvQueues>> {
|
||||
return [];
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,287 @@
|
||||
import type { MarQSKeyProducer, MarQSKeyProducerEnv, QueueDescriptor } from "./types";
|
||||
|
||||
const constants = {
|
||||
SHARED_QUEUE: "sharedQueue",
|
||||
SHARED_WORKER_QUEUE: "sharedWorkerQueue",
|
||||
CURRENT_CONCURRENCY_PART: "currentConcurrency",
|
||||
CONCURRENCY_LIMIT_PART: "concurrency",
|
||||
DISABLED_CONCURRENCY_LIMIT_PART: "disabledConcurrency",
|
||||
ENV_PART: "env",
|
||||
ORG_PART: "org",
|
||||
QUEUE_PART: "queue",
|
||||
CONCURRENCY_KEY_PART: "ck",
|
||||
MESSAGE_PART: "message",
|
||||
RESERVE_CONCURRENCY_PART: "reserveConcurrency",
|
||||
} as const;
|
||||
|
||||
const ORG_REGEX = /org:([^:]+):/;
|
||||
const ENV_REGEX = /env:([^:]+):/;
|
||||
const QUEUE_REGEX = /queue:([^:]+)(?::|$)/;
|
||||
const CONCURRENCY_KEY_REGEX = /ck:([^:]+)(?::|$)/;
|
||||
|
||||
export class MarQSShortKeyProducer implements MarQSKeyProducer {
|
||||
constructor(private _prefix: string) {}
|
||||
|
||||
sharedQueueScanPattern() {
|
||||
return `${this._prefix}*${constants.SHARED_QUEUE}`;
|
||||
}
|
||||
|
||||
queueCurrentConcurrencyScanPattern() {
|
||||
return `${this._prefix}${constants.ORG_PART}:*:${constants.ENV_PART}:*:queue:*:${constants.CURRENT_CONCURRENCY_PART}`;
|
||||
}
|
||||
|
||||
stripKeyPrefix(key: string): string {
|
||||
if (key.startsWith(this._prefix)) {
|
||||
return key.slice(this._prefix.length);
|
||||
}
|
||||
|
||||
return key;
|
||||
}
|
||||
|
||||
queueConcurrencyLimitKey(env: MarQSKeyProducerEnv, queue: string) {
|
||||
return [this.queueKey(env, queue), constants.CONCURRENCY_LIMIT_PART].join(":");
|
||||
}
|
||||
|
||||
envConcurrencyLimitKey(envId: string): string;
|
||||
envConcurrencyLimitKey(env: MarQSKeyProducerEnv): string;
|
||||
envConcurrencyLimitKey(envOrId: MarQSKeyProducerEnv | string): string {
|
||||
return [
|
||||
this.envKeySection(typeof envOrId === "string" ? envOrId : envOrId.id),
|
||||
constants.CONCURRENCY_LIMIT_PART,
|
||||
].join(":");
|
||||
}
|
||||
|
||||
queueKey(orgId: string, envId: string, queue: string, concurrencyKey?: string): string;
|
||||
queueKey(env: MarQSKeyProducerEnv, queue: string, concurrencyKey?: string): string;
|
||||
queueKey(
|
||||
envOrOrgId: MarQSKeyProducerEnv | string,
|
||||
queueOrEnvId: string,
|
||||
queueOrConcurrencyKey: string,
|
||||
concurrencyKeyOrPriority?: string | number
|
||||
): string {
|
||||
if (typeof envOrOrgId === "string") {
|
||||
return [
|
||||
this.orgKeySection(envOrOrgId),
|
||||
this.envKeySection(queueOrEnvId),
|
||||
this.queueSection(queueOrConcurrencyKey),
|
||||
]
|
||||
.concat(
|
||||
typeof concurrencyKeyOrPriority === "string"
|
||||
? this.concurrencyKeySection(concurrencyKeyOrPriority)
|
||||
: []
|
||||
)
|
||||
.join(":");
|
||||
} else {
|
||||
return [
|
||||
this.orgKeySection(envOrOrgId.organizationId),
|
||||
this.envKeySection(envOrOrgId.id),
|
||||
this.queueSection(queueOrEnvId),
|
||||
]
|
||||
.concat(queueOrConcurrencyKey ? this.concurrencyKeySection(queueOrConcurrencyKey) : [])
|
||||
.join(":");
|
||||
}
|
||||
}
|
||||
|
||||
queueKeyFromQueue(queue: string): string {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return this.queueKey(
|
||||
descriptor.organization,
|
||||
descriptor.environment,
|
||||
descriptor.name,
|
||||
descriptor.concurrencyKey
|
||||
);
|
||||
}
|
||||
|
||||
envSharedQueueKey(env: MarQSKeyProducerEnv) {
|
||||
if (env.type === "DEVELOPMENT") {
|
||||
return [
|
||||
this.orgKeySection(env.organizationId),
|
||||
this.envKeySection(env.id),
|
||||
constants.SHARED_QUEUE,
|
||||
].join(":");
|
||||
}
|
||||
|
||||
return this.sharedQueueKey();
|
||||
}
|
||||
|
||||
sharedQueueKey(): string {
|
||||
return constants.SHARED_QUEUE;
|
||||
}
|
||||
|
||||
sharedWorkerQueueKey(): string {
|
||||
return constants.SHARED_WORKER_QUEUE;
|
||||
}
|
||||
|
||||
queueConcurrencyLimitKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return this.queueConcurrencyLimitKeyFromDescriptor(descriptor);
|
||||
}
|
||||
|
||||
queueCurrentConcurrencyKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
return this.currentConcurrencyKeyFromDescriptor(descriptor);
|
||||
}
|
||||
|
||||
queueReserveConcurrencyKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return this.queueReserveConcurrencyKeyFromDescriptor(descriptor);
|
||||
}
|
||||
|
||||
queueCurrentConcurrencyKey(
|
||||
env: MarQSKeyProducerEnv,
|
||||
queue: string,
|
||||
concurrencyKey?: string
|
||||
): string {
|
||||
return [this.queueKey(env, queue, concurrencyKey), constants.CURRENT_CONCURRENCY_PART].join(
|
||||
":"
|
||||
);
|
||||
}
|
||||
|
||||
envConcurrencyLimitKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return `${constants.ENV_PART}:${descriptor.environment}:${constants.CONCURRENCY_LIMIT_PART}`;
|
||||
}
|
||||
|
||||
envCurrentConcurrencyKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return `${constants.ENV_PART}:${descriptor.environment}:${constants.CURRENT_CONCURRENCY_PART}`;
|
||||
}
|
||||
|
||||
envReserveConcurrencyKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return this.envReserveConcurrencyKey(descriptor.environment);
|
||||
}
|
||||
|
||||
envReserveConcurrencyKey(envId: string): string {
|
||||
return `${constants.ENV_PART}:${this.shortId(envId)}:${constants.RESERVE_CONCURRENCY_PART}`;
|
||||
}
|
||||
|
||||
envCurrentConcurrencyKey(envId: string): string;
|
||||
envCurrentConcurrencyKey(env: MarQSKeyProducerEnv): string;
|
||||
envCurrentConcurrencyKey(envOrId: MarQSKeyProducerEnv | string): string {
|
||||
return [
|
||||
this.envKeySection(typeof envOrId === "string" ? envOrId : envOrId.id),
|
||||
constants.CURRENT_CONCURRENCY_PART,
|
||||
].join(":");
|
||||
}
|
||||
|
||||
envQueueKeyFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return `${constants.ENV_PART}:${descriptor.environment}:${constants.QUEUE_PART}`;
|
||||
}
|
||||
|
||||
envQueueKey(env: MarQSKeyProducerEnv): string {
|
||||
return [constants.ENV_PART, this.shortId(env.id), constants.QUEUE_PART].join(":");
|
||||
}
|
||||
|
||||
messageKey(messageId: string) {
|
||||
return `${constants.MESSAGE_PART}:${messageId}`;
|
||||
}
|
||||
|
||||
nackCounterKey(messageId: string): string {
|
||||
return `${constants.MESSAGE_PART}:${messageId}:nacks`;
|
||||
}
|
||||
|
||||
orgIdFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return descriptor.organization;
|
||||
}
|
||||
|
||||
envIdFromQueue(queue: string) {
|
||||
const descriptor = this.queueDescriptorFromQueue(queue);
|
||||
|
||||
return descriptor.environment;
|
||||
}
|
||||
|
||||
queueDescriptorFromQueue(queue: string): QueueDescriptor {
|
||||
const match = queue.match(QUEUE_REGEX);
|
||||
|
||||
if (!match) {
|
||||
throw new Error(`Invalid queue: ${queue}, no queue name found`);
|
||||
}
|
||||
|
||||
const [, queueName] = match;
|
||||
|
||||
const envMatch = queue.match(ENV_REGEX);
|
||||
|
||||
if (!envMatch) {
|
||||
throw new Error(`Invalid queue: ${queue}, no environment found`);
|
||||
}
|
||||
|
||||
const [, envId] = envMatch;
|
||||
|
||||
const orgMatch = queue.match(ORG_REGEX);
|
||||
|
||||
if (!orgMatch) {
|
||||
throw new Error(`Invalid queue: ${queue}, no organization found`);
|
||||
}
|
||||
|
||||
const [, orgId] = orgMatch;
|
||||
|
||||
const concurrencyKeyMatch = queue.match(CONCURRENCY_KEY_REGEX);
|
||||
|
||||
const concurrencyKey = concurrencyKeyMatch ? concurrencyKeyMatch[1] : undefined;
|
||||
|
||||
return {
|
||||
name: queueName,
|
||||
environment: envId,
|
||||
organization: orgId,
|
||||
concurrencyKey,
|
||||
};
|
||||
}
|
||||
|
||||
private shortId(id: string) {
|
||||
// Return the last 12 characters of the id
|
||||
return id.slice(-12);
|
||||
}
|
||||
|
||||
private envKeySection(envId: string) {
|
||||
return `${constants.ENV_PART}:${this.shortId(envId)}`;
|
||||
}
|
||||
|
||||
private orgKeySection(orgId: string) {
|
||||
return `${constants.ORG_PART}:${this.shortId(orgId)}`;
|
||||
}
|
||||
|
||||
private queueSection(queue: string) {
|
||||
return `${constants.QUEUE_PART}:${queue}`;
|
||||
}
|
||||
|
||||
private concurrencyKeySection(concurrencyKey: string) {
|
||||
return `${constants.CONCURRENCY_KEY_PART}:${concurrencyKey}`;
|
||||
}
|
||||
|
||||
private currentConcurrencyKeyFromDescriptor(descriptor: QueueDescriptor) {
|
||||
return [
|
||||
this.queueKey(
|
||||
descriptor.organization,
|
||||
descriptor.environment,
|
||||
descriptor.name,
|
||||
descriptor.concurrencyKey
|
||||
),
|
||||
constants.CURRENT_CONCURRENCY_PART,
|
||||
].join(":");
|
||||
}
|
||||
|
||||
private queueReserveConcurrencyKeyFromDescriptor(descriptor: QueueDescriptor) {
|
||||
return [
|
||||
this.queueKey(descriptor.organization, descriptor.environment, descriptor.name),
|
||||
constants.RESERVE_CONCURRENCY_PART,
|
||||
].join(":");
|
||||
}
|
||||
|
||||
private queueConcurrencyLimitKeyFromDescriptor(descriptor: QueueDescriptor) {
|
||||
return [
|
||||
this.queueKey(descriptor.organization, descriptor.environment, descriptor.name),
|
||||
constants.CONCURRENCY_LIMIT_PART,
|
||||
].join(":");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,111 @@
|
||||
import type { RuntimeEnvironmentType } from "@trigger.dev/database";
|
||||
import { z } from "zod";
|
||||
|
||||
export type QueueRange = { offset: number; count: number };
|
||||
|
||||
export type QueueDescriptor = {
|
||||
organization: string;
|
||||
environment: string;
|
||||
name: string;
|
||||
concurrencyKey?: string;
|
||||
};
|
||||
|
||||
export type MarQSKeyProducerEnv = {
|
||||
id: string;
|
||||
organizationId: string;
|
||||
type: RuntimeEnvironmentType;
|
||||
};
|
||||
|
||||
export interface MarQSKeyProducer {
|
||||
queueConcurrencyLimitKey(env: MarQSKeyProducerEnv, queue: string): string;
|
||||
|
||||
envConcurrencyLimitKey(envId: string): string;
|
||||
envConcurrencyLimitKey(env: MarQSKeyProducerEnv): string;
|
||||
|
||||
envCurrentConcurrencyKey(envId: string): string;
|
||||
envCurrentConcurrencyKey(env: MarQSKeyProducerEnv): string;
|
||||
|
||||
envReserveConcurrencyKey(envId: string): string;
|
||||
|
||||
queueKey(orgId: string, envId: string, queue: string, concurrencyKey?: string): string;
|
||||
queueKey(env: MarQSKeyProducerEnv, queue: string, concurrencyKey?: string): string;
|
||||
|
||||
queueKeyFromQueue(queue: string): string;
|
||||
|
||||
envQueueKey(env: MarQSKeyProducerEnv): string;
|
||||
envSharedQueueKey(env: MarQSKeyProducerEnv): string;
|
||||
sharedQueueKey(): string;
|
||||
sharedQueueScanPattern(): string;
|
||||
sharedWorkerQueueKey(): string;
|
||||
queueCurrentConcurrencyScanPattern(): string;
|
||||
queueConcurrencyLimitKeyFromQueue(queue: string): string;
|
||||
queueCurrentConcurrencyKeyFromQueue(queue: string): string;
|
||||
queueCurrentConcurrencyKey(
|
||||
env: MarQSKeyProducerEnv,
|
||||
queue: string,
|
||||
concurrencyKey?: string
|
||||
): string;
|
||||
envConcurrencyLimitKeyFromQueue(queue: string): string;
|
||||
envCurrentConcurrencyKeyFromQueue(queue: string): string;
|
||||
envReserveConcurrencyKeyFromQueue(queue: string): string;
|
||||
envQueueKeyFromQueue(queue: string): string;
|
||||
messageKey(messageId: string): string;
|
||||
nackCounterKey(messageId: string): string;
|
||||
stripKeyPrefix(key: string): string;
|
||||
orgIdFromQueue(queue: string): string;
|
||||
envIdFromQueue(queue: string): string;
|
||||
|
||||
queueReserveConcurrencyKeyFromQueue(queue: string): string;
|
||||
queueDescriptorFromQueue(queue: string): QueueDescriptor;
|
||||
}
|
||||
|
||||
export type EnvQueues = {
|
||||
envId: string;
|
||||
queues: string[];
|
||||
};
|
||||
|
||||
const MarQSPriorityLevel = z.enum(["resume", "retry"]);
|
||||
|
||||
export type MarQSPriorityLevel = z.infer<typeof MarQSPriorityLevel>;
|
||||
|
||||
export interface MarQSFairDequeueStrategy {
|
||||
distributeFairQueuesFromParentQueue(
|
||||
parentQueue: string,
|
||||
consumerId: string
|
||||
): Promise<Array<EnvQueues>>;
|
||||
}
|
||||
|
||||
export const MessagePayload = z.object({
|
||||
version: z.literal("1"),
|
||||
data: z.record(z.unknown()),
|
||||
queue: z.string(),
|
||||
messageId: z.string(),
|
||||
timestamp: z.number(),
|
||||
parentQueue: z.string(),
|
||||
concurrencyKey: z.string().optional(),
|
||||
priority: MarQSPriorityLevel.optional(),
|
||||
availableAt: z.number().optional(),
|
||||
enqueueMethod: z.enum(["enqueue", "requeue", "replace"]).default("enqueue"),
|
||||
});
|
||||
|
||||
export type MessagePayload = z.infer<typeof MessagePayload>;
|
||||
|
||||
export interface MessageQueueSubscriber {
|
||||
messageEnqueued(message: MessagePayload): Promise<void>;
|
||||
messageDequeued(message: MessagePayload): Promise<void>;
|
||||
messageAcked(message: MessagePayload): Promise<void>;
|
||||
messageNacked(message: MessagePayload): Promise<void>;
|
||||
messageReplaced(message: MessagePayload): Promise<void>;
|
||||
messageRequeued(message: MessagePayload): Promise<void>;
|
||||
}
|
||||
|
||||
export interface VisibilityTimeoutStrategy {
|
||||
startHeartbeat(messageId: string, timeoutInMs: number): Promise<void>;
|
||||
heartbeat(messageId: string, timeoutInMs: number): Promise<void>;
|
||||
cancelHeartbeat(messageId: string): Promise<void>;
|
||||
}
|
||||
|
||||
export type EnqueueMessageReserveConcurrencyOptions = {
|
||||
messageId: string;
|
||||
recursiveQueue: boolean;
|
||||
};
|
||||
@@ -0,0 +1,39 @@
|
||||
import { legacyRunEngineWorker } from "../legacyRunEngineWorker.server";
|
||||
import { TaskRunHeartbeatFailedService } from "../taskRunHeartbeatFailed.server";
|
||||
import type { VisibilityTimeoutStrategy } from "./types";
|
||||
|
||||
export class V3GraphileVisibilityTimeout implements VisibilityTimeoutStrategy {
|
||||
async startHeartbeat(messageId: string, timeoutInMs: number): Promise<void> {
|
||||
await TaskRunHeartbeatFailedService.enqueue(messageId, new Date(Date.now() + timeoutInMs));
|
||||
}
|
||||
|
||||
async heartbeat(messageId: string, timeoutInMs: number): Promise<void> {
|
||||
await TaskRunHeartbeatFailedService.enqueue(messageId, new Date(Date.now() + timeoutInMs));
|
||||
}
|
||||
|
||||
async cancelHeartbeat(messageId: string): Promise<void> {
|
||||
await TaskRunHeartbeatFailedService.dequeue(messageId);
|
||||
}
|
||||
}
|
||||
|
||||
export class V3LegacyRunEngineWorkerVisibilityTimeout implements VisibilityTimeoutStrategy {
|
||||
async startHeartbeat(messageId: string, timeoutInMs: number): Promise<void> {
|
||||
await legacyRunEngineWorker.enqueue({
|
||||
id: `heartbeat:${messageId}`,
|
||||
job: "runHeartbeat",
|
||||
payload: { runId: messageId },
|
||||
availableAt: new Date(Date.now() + timeoutInMs),
|
||||
});
|
||||
}
|
||||
|
||||
async heartbeat(messageId: string, timeoutInMs: number): Promise<void> {
|
||||
await legacyRunEngineWorker.reschedule(
|
||||
`heartbeat:${messageId}`,
|
||||
new Date(Date.now() + timeoutInMs)
|
||||
);
|
||||
}
|
||||
|
||||
async cancelHeartbeat(messageId: string): Promise<void> {
|
||||
await legacyRunEngineWorker.ack(`heartbeat:${messageId}`);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user