Add queue size limit guard on triggering tasks

ericallam · ericallam · commit 62a0a2874838 · 2024-10-02T17:59:58.000-07:00
diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts
@@ -213,6 +213,9 @@ const EnvironmentSchema = z.object({
   TASK_PAYLOAD_OFFLOAD_THRESHOLD: z.coerce.number().int().default(524_288), // 512KB
   TASK_PAYLOAD_MAXIMUM_SIZE: z.coerce.number().int().default(3_145_728), // 3MB
   TASK_RUN_METADATA_MAXIMUM_SIZE: z.coerce.number().int().default(4_096), // 4KB
+
+  MAXIMUM_DEV_QUEUE_SIZE: z.coerce.number().int().optional(),
+  MAXIMUM_DEPLOYED_QUEUE_SIZE: z.coerce.number().int().optional(),
 });
 
 export type Environment = z.infer<typeof EnvironmentSchema>;
diff --git a/apps/webapp/app/v3/marqs/index.server.ts b/apps/webapp/app/v3/marqs/index.server.ts
@@ -139,6 +139,10 @@ export class MarQS {
     return this.redis.zcard(this.keys.queueKey(env, queue, concurrencyKey));
   }
 
+  public async lengthOfEnvQueue(env: AuthenticatedEnvironment) {
+    return this.redis.zcard(this.keys.envQueueKey(env));
+  }
+
   public async oldestMessageInQueue(
     env: AuthenticatedEnvironment,
     queue: string,
diff --git a/apps/webapp/app/v3/marqs/marqsKeyProducer.server.ts b/apps/webapp/app/v3/marqs/marqsKeyProducer.server.ts
@@ -134,6 +134,10 @@ export class MarQSShortKeyProducer implements MarQSKeyProducer {
     return `${constants.ENV_PART}:${envId}:${constants.QUEUE_PART}`;
   }
 
+  envQueueKey(env: AuthenticatedEnvironment): string {
+    return [constants.ENV_PART, this.shortId(env.id), constants.QUEUE_PART].join(":");
+  }
+
   messageKey(messageId: string) {
     return `${constants.MESSAGE_PART}:${messageId}`;
   }
diff --git a/apps/webapp/app/v3/marqs/types.ts b/apps/webapp/app/v3/marqs/types.ts
@@ -26,6 +26,7 @@ export interface MarQSKeyProducer {
   envConcurrencyLimitKey(env: AuthenticatedEnvironment): string;
   orgConcurrencyLimitKey(env: AuthenticatedEnvironment): string;
   queueKey(env: AuthenticatedEnvironment, queue: string, concurrencyKey?: string): string;
+  envQueueKey(env: AuthenticatedEnvironment): string;
   envSharedQueueKey(env: AuthenticatedEnvironment): string;
   sharedQueueKey(): string;
   sharedQueueScanPattern(): string;
diff --git a/apps/webapp/app/v3/queueSizeLimits.server.ts b/apps/webapp/app/v3/queueSizeLimits.server.ts
@@ -0,0 +1,40 @@
+import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
+import { env } from "~/env.server";
+import { MarQS } from "./marqs/index.server";
+
+export type QueueSizeGuardResult = {
+  isWithinLimits: boolean;
+  maximumSize?: number;
+  queueSize?: number;
+};
+
+export async function guardQueueSizeLimitsForEnv(
+  environment: AuthenticatedEnvironment,
+  marqs?: MarQS
+): Promise<QueueSizeGuardResult> {
+  const maximumSize = getMaximumSizeForEnvironment(environment);
+
+  if (typeof maximumSize === "undefined") {
+    return { isWithinLimits: true };
+  }
+
+  if (!marqs) {
+    return { isWithinLimits: true, maximumSize };
+  }
+
+  const queueSize = await marqs.lengthOfEnvQueue(environment);
+
+  return {
+    isWithinLimits: queueSize < maximumSize,
+    maximumSize,
+    queueSize,
+  };
+}
+
+function getMaximumSizeForEnvironment(environment: AuthenticatedEnvironment): number | undefined {
+  if (environment.type === "DEVELOPMENT") {
+    return environment.organization.maximumDevQueueSize ?? env.MAXIMUM_DEV_QUEUE_SIZE;
+  } else {
+    return environment.organization.maximumDeployedQueueSize ?? env.MAXIMUM_DEPLOYED_QUEUE_SIZE;
+  }
+}
diff --git a/apps/webapp/app/v3/services/triggerTask.server.ts b/apps/webapp/app/v3/services/triggerTask.server.ts
@@ -22,6 +22,7 @@ import { createTag, MAX_TAGS_PER_RUN } from "~/models/taskRunTag.server";
 import { findCurrentWorkerFromEnvironment } from "../models/workerDeployment.server";
 import { handleMetadataPacket } from "~/utils/packets";
 import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
+import { guardQueueSizeLimitsForEnv } from "../queueSizeLimits.server";
 
 export type TriggerTaskServiceOptions = {
   idempotencyKey?: string;
@@ -82,6 +83,24 @@ export class TriggerTaskService extends BaseService {
         }
       }
 
+      const queueSizeGuard = await guardQueueSizeLimitsForEnv(environment, marqs);
+
+      logger.debug("Queue size guard result", {
+        queueSizeGuard,
+        environment: {
+          id: environment.id,
+          type: environment.type,
+          organization: environment.organization,
+          project: environment.project,
+        },
+      });
+
+      if (!queueSizeGuard.isWithinLimits) {
+        throw new ServiceValidationError(
+          `Cannot trigger ${taskId} as the queue size limit for this environment has been reached. The maximum size is ${queueSizeGuard.maximumSize}`
+        );
+      }
+
       if (
         body.options?.tags &&
         typeof body.options.tags !== "string" &&
diff --git a/packages/database/prisma/migrations/20241003002757_add_max_queue_sizes_to_org/migration.sql b/packages/database/prisma/migrations/20241003002757_add_max_queue_sizes_to_org/migration.sql
@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "Organization" ADD COLUMN     "maximumDeployedQueueSize" INTEGER,
+ADD COLUMN     "maximumDevQueueSize" INTEGER;
diff --git a/packages/database/prisma/schema.prisma b/packages/database/prisma/schema.prisma
@@ -114,6 +114,9 @@ model Organization {
   /// This is deprecated and will be removed in the future
   maximumSchedulesLimit          Int @default(5)
 
+  maximumDevQueueSize      Int?
+  maximumDeployedQueueSize Int?
+
   createdAt DateTime  @default(now())
   updatedAt DateTime  @updatedAt
   deletedAt DateTime?
diff --git a/references/v3-catalog/src/trigger/simple.ts b/references/v3-catalog/src/trigger/simple.ts
@@ -2,6 +2,7 @@ import "server-only";
 import { logger, SubtaskUnwrapError, task, tasks, wait } from "@trigger.dev/sdk/v3";
 import { traceAsync } from "@/telemetry.js";
 import { HeaderGenerator } from "header-generator";
+import { setTimeout as setTimeoutP } from "node:timers/promises";
 
 let headerGenerator = new HeaderGenerator({
   browsers: [{ name: "firefox", minVersion: 90 }, { name: "chrome", minVersion: 110 }, "safari"],
@@ -215,3 +216,22 @@ export const retryTask = task({
     throw new Error("This task will always fail");
   },
 });
+
+export const maximumQueueDepthParent = task({
+  id: "maximum-queue-depth-parent",
+  run: async (payload: any) => {
+    await maximumQueueDepthChild.trigger({});
+    await maximumQueueDepthChild.trigger({});
+    await maximumQueueDepthChild.trigger({});
+  },
+});
+
+export const maximumQueueDepthChild = task({
+  id: "maximum-queue-depth-child",
+  queue: {
+    concurrencyLimit: 1,
+  },
+  run: async (payload: any) => {
+    await setTimeoutP(10_000);
+  },
+});
diff --git a/references/v3-catalog/trigger.config.ts b/references/v3-catalog/trigger.config.ts
@@ -18,7 +18,7 @@ export default defineConfig({
   instrumentations: [new OpenAIInstrumentation()],
   additionalFiles: ["wrangler/wrangler.toml"],
   retries: {
-    enabledInDev: true,
+    enabledInDev: false,
     default: {
       maxAttempts: 10,
       minTimeoutInMs: 5_000,

Original file line number	Diff line number	Diff line change
`@@ -134,6 +134,10 @@ export class MarQSShortKeyProducer implements MarQSKeyProducer {`
`134`	`134`	return `${constants.ENV_PART}:${envId}:${constants.QUEUE_PART}`;
`135`	`135`	`}`
`136`	`136`
	`137`	`+ envQueueKey(env: AuthenticatedEnvironment): string {`
	`138`	`+ return [constants.ENV_PART, this.shortId(env.id), constants.QUEUE_PART].join(":");`
	`139`	`+ }`
	`140`	`+`
`137`	`141`	`messageKey(messageId: string) {`
`138`	`142`	return `${constants.MESSAGE_PART}:${messageId}`;
`139`	`143`	`}`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+-- AlterTable`
	`2`	`+ALTER TABLE "Organization" ADD COLUMN "maximumDeployedQueueSize" INTEGER,`
	`3`	`+ADD COLUMN "maximumDevQueueSize" INTEGER;`