release concurrency system reliability improvements

apps/webapp/app/components/admin/debugRun.tsx

@@ -66,7 +66,15 @@ function DebugRunContent({ friendlyId }: { friendlyId: string }) {
   );
 }
 
-function DebugRunData({
+function DebugRunData(props: UseDataFunctionReturn<typeof loader>) {
+  if (props.engine === "V1") {
+    return <DebugRunDataEngineV1 {...props} />;
+  }
+
+  return <DebugRunDataEngineV2 {...props} />;
+}
+
+function DebugRunDataEngineV1({
   run,
   queueConcurrencyLimit,
   queueCurrentConcurrency,
@@ -338,3 +346,55 @@ function DebugRunData({
     </Property.Table>
   );
 }
+
+function DebugRunDataEngineV2({
+  run,
+  queueConcurrencyLimit,
+  queueCurrentConcurrency,
+  envConcurrencyLimit,
+  envCurrentConcurrency,
+  keys,
+}: UseDataFunctionReturn<typeof loader>) {
+  return (
+    <Property.Table>
+      <Property.Item>
+        <Property.Label>ID</Property.Label>
+        <Property.Value className="flex items-center gap-2">
+          <ClipboardField value={run.id} variant="tertiary/small" iconButton />
+        </Property.Value>
+      </Property.Item>
+      <Property.Item>
+        <Property.Label>Queue current concurrency</Property.Label>
+        <Property.Value className="flex items-center gap-2">
+          <span>{queueCurrentConcurrency ?? "0"}</span>
+        </Property.Value>
+      </Property.Item>
+      <Property.Item>
+        <Property.Label>Queue concurrency limit</Property.Label>
+        <Property.Value className="flex items-center gap-2">
+          <span>{queueConcurrencyLimit ?? "Not set"}</span>
+        </Property.Value>
+      </Property.Item>
+      <Property.Item>
+        <Property.Label>Env current concurrency</Property.Label>
+        <Property.Value className="flex items-center gap-2">
+          <span>{envCurrentConcurrency ?? "0"}</span>
+        </Property.Value>
+      </Property.Item>
+      <Property.Item>
+        <Property.Label>Env concurrency limit</Property.Label>
+        <Property.Value className="flex items-center gap-2">
+          <span>{envConcurrencyLimit ?? "Not set"}</span>
+        </Property.Value>
+      </Property.Item>
+      {keys.map((key) => (
+        <Property.Item>
+          <Property.Label>{key.label}</Property.Label>
+          <Property.Value className="flex items-center gap-2">
+            <ClipboardField value={key.key} variant="tertiary/small" iconButton />
+          </Property.Value>
+        </Property.Item>
+      ))}
+    </Property.Table>
+  );
+}

apps/webapp/app/env.server.ts

@@ -452,6 +452,10 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_TIMEOUT_PENDING_CANCEL: z.coerce.number().int().default(60_000),
   RUN_ENGINE_TIMEOUT_EXECUTING: z.coerce.number().int().default(60_000),
   RUN_ENGINE_TIMEOUT_EXECUTING_WITH_WAITPOINTS: z.coerce.number().int().default(60_000),
+  RUN_ENGINE_TIMEOUT_SUSPENDED: z.coerce
+    .number()
+    .int()
+    .default(60_000 * 10),
   RUN_ENGINE_DEBUG_WORKER_NOTIFICATIONS: z.coerce.boolean().default(false),
   RUN_ENGINE_PARENT_QUEUE_LIMIT: z.coerce.number().int().default(1000),
   RUN_ENGINE_CONCURRENCY_LIMIT_BIAS: z.coerce.number().default(0.75),
@@ -605,6 +609,11 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_RELEASE_CONCURRENCY_ENABLED: z.string().default("0"),
   RUN_ENGINE_RELEASE_CONCURRENCY_DISABLE_CONSUMERS: z.string().default("0"),
   RUN_ENGINE_RELEASE_CONCURRENCY_MAX_TOKENS_RATIO: z.coerce.number().default(1),
+  RUN_ENGINE_RELEASE_CONCURRENCY_RELEASINGS_MAX_AGE: z.coerce
+    .number()
+    .int()
+    .default(60_000 * 30),
+  RUN_ENGINE_RELEASE_CONCURRENCY_RELEASINGS_POLL_INTERVAL: z.coerce.number().int().default(60_000),
   RUN_ENGINE_RELEASE_CONCURRENCY_MAX_RETRIES: z.coerce.number().int().default(3),
   RUN_ENGINE_RELEASE_CONCURRENCY_CONSUMERS_COUNT: z.coerce.number().int().default(1),
   RUN_ENGINE_RELEASE_CONCURRENCY_POLL_INTERVAL: z.coerce.number().int().default(500),

apps/webapp/app/routes/resources.taskruns.$runParam.debug.ts

@@ -4,6 +4,7 @@ import { z } from "zod";
 import { $replica } from "~/db.server";
 import { requireUserId } from "~/services/session.server";
 import { marqs } from "~/v3/marqs/index.server";
+import { engine } from "~/v3/runEngine.server";
 
 const ParamSchema = z.object({
   runParam: z.string(),
@@ -17,6 +18,7 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
     where: { friendlyId: runParam, project: { organization: { members: { some: { userId } } } } },
     select: {
       id: true,
+      engine: true,
       friendlyId: true,
       queue: true,
       concurrencyKey: true,
@@ -27,6 +29,8 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
           type: true,
           slug: true,
           organizationId: true,
+          project: true,
+          maximumConcurrencyLimit: true,
           organization: {
             select: {
               id: true,
@@ -41,33 +45,132 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
     throw new Response("Not Found", { status: 404 });
   }
 
-  const queueConcurrencyLimit = await marqs.getQueueConcurrencyLimit(
-    run.runtimeEnvironment,
-    run.queue
-  );
-  const envConcurrencyLimit = await marqs.getEnvConcurrencyLimit(run.runtimeEnvironment);
-  const queueCurrentConcurrency = await marqs.currentConcurrencyOfQueue(
-    run.runtimeEnvironment,
-    run.queue,
-    run.concurrencyKey ?? undefined
-  );
-  const envCurrentConcurrency = await marqs.currentConcurrencyOfEnvironment(run.runtimeEnvironment);
-
-  const queueReserveConcurrency = await marqs.reserveConcurrencyOfQueue(
-    run.runtimeEnvironment,
-    run.queue,
-    run.concurrencyKey ?? undefined
-  );
-
-  const envReserveConcurrency = await marqs.reserveConcurrencyOfEnvironment(run.runtimeEnvironment);
-
-  return typedjson({
-    run,
-    queueConcurrencyLimit,
-    envConcurrencyLimit,
-    queueCurrentConcurrency,
-    envCurrentConcurrency,
-    queueReserveConcurrency,
-    envReserveConcurrency,
-  });
+  if (run.engine === "V1") {
+    const queueConcurrencyLimit = await marqs.getQueueConcurrencyLimit(
+      run.runtimeEnvironment,
+      run.queue
+    );
+    const envConcurrencyLimit = await marqs.getEnvConcurrencyLimit(run.runtimeEnvironment);
+    const queueCurrentConcurrency = await marqs.currentConcurrencyOfQueue(
+      run.runtimeEnvironment,
+      run.queue,
+      run.concurrencyKey ?? undefined
+    );
+    const envCurrentConcurrency = await marqs.currentConcurrencyOfEnvironment(
+      run.runtimeEnvironment
+    );
+
+    const queueReserveConcurrency = await marqs.reserveConcurrencyOfQueue(
+      run.runtimeEnvironment,
+      run.queue,
+      run.concurrencyKey ?? undefined
+    );
+
+    const envReserveConcurrency = await marqs.reserveConcurrencyOfEnvironment(
+      run.runtimeEnvironment
+    );
+
+    return typedjson({
+      engine: "V1",
+      run,
+      queueConcurrencyLimit,
+      envConcurrencyLimit,
+      queueCurrentConcurrency,
+      envCurrentConcurrency,
+      queueReserveConcurrency,
+      envReserveConcurrency,
+      keys: [],
+    });
+  } else {
+    const queueConcurrencyLimit = await engine.runQueue.getQueueConcurrencyLimit(
+      run.runtimeEnvironment,
+      run.queue
+    );
+
+    const envConcurrencyLimit = await engine.runQueue.getEnvConcurrencyLimit(
+      run.runtimeEnvironment
+    );
+
+    const queueCurrentConcurrency = await engine.runQueue.currentConcurrencyOfQueue(
+      run.runtimeEnvironment,
+      run.queue,
+      run.concurrencyKey ?? undefined
+    );
+
+    const envCurrentConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment(
+      run.runtimeEnvironment
+    );
+
+    const queueCurrentConcurrencyKey = engine.runQueue.keys.currentConcurrencyKey(
+      run.runtimeEnvironment,
+      run.queue,
+      run.concurrencyKey ?? undefined
+    );
+
+    const envCurrentConcurrencyKey = engine.runQueue.keys.envCurrentConcurrencyKey(
+      run.runtimeEnvironment
+    );
+
+    const queueConcurrencyLimitKey = engine.runQueue.keys.queueConcurrencyLimitKey(
+      run.runtimeEnvironment,
+      run.queue
+    );
+
+    const envConcurrencyLimitKey = engine.runQueue.keys.envConcurrencyLimitKey(
+      run.runtimeEnvironment
+    );
+
+    const releaseConcurrencyBucketKey = `engine:release-concurrency:org:${run.runtimeEnvironment.organizationId}:proj:${run.runtimeEnvironment.project.id}:env:${run.runtimeEnvironment.id}:bucket`;
+    const releaseConcurrencyQueueKey = `engine:release-concurrency:org:${run.runtimeEnvironment.organizationId}:proj:${run.runtimeEnvironment.project.id}:env:${run.runtimeEnvironment.id}:queue`;
+    const releaseConcurrencyMetadataKey = `engine:release-concurrency:org:${run.runtimeEnvironment.organizationId}:proj:${run.runtimeEnvironment.project.id}:env:${run.runtimeEnvironment.id}:metadata`;
+
+    const withPrefix = (key: string) => `engine:runqueue:${key}`;
+
+    const keys = [
+      {
+        label: "Queue current concurrency set",
+        key: withPrefix(queueCurrentConcurrencyKey),
+      },
+      {
+        label: "Env current concurrency set",
+        key: withPrefix(envCurrentConcurrencyKey),
+      },
+      {
+        label: "Queue concurrency limit",
+        key: withPrefix(queueConcurrencyLimitKey),
+      },
+      {
+        label: "Env concurrency limit",
+        key: withPrefix(envConcurrencyLimitKey),
+      },
+      {
+        label: "Release concurrency bucket",
+        key: releaseConcurrencyBucketKey,
+      },
+      {
+        label: "Release concurrency queue",
+        key: releaseConcurrencyQueueKey,
+      },
+      {
+        label: "Release concurrency metadata",
+        key: releaseConcurrencyMetadataKey,
+      },
+      {
+        label: "Release concurrency releasings",
+        key: "engine:release-concurrency:releasings",
+      },
+    ];
+
+    return typedjson({
+      engine: "V2",
+      run,
+      queueConcurrencyLimit,
+      envConcurrencyLimit,
+      queueCurrentConcurrency,
+      envCurrentConcurrency,
+      queueReserveConcurrency: undefined,
+      envReserveConcurrency: undefined,
+      keys,
+    });
+  }
 }

apps/webapp/app/v3/runEngine.server.ts

@@ -76,11 +76,14 @@ function createRunEngine() {
       PENDING_CANCEL: env.RUN_ENGINE_TIMEOUT_PENDING_CANCEL,
       EXECUTING: env.RUN_ENGINE_TIMEOUT_EXECUTING,
       EXECUTING_WITH_WAITPOINTS: env.RUN_ENGINE_TIMEOUT_EXECUTING_WITH_WAITPOINTS,
+      SUSPENDED: env.RUN_ENGINE_TIMEOUT_SUSPENDED,
     },
     releaseConcurrency: {
       disabled: env.RUN_ENGINE_RELEASE_CONCURRENCY_ENABLED === "0",
       disableConsumers: env.RUN_ENGINE_RELEASE_CONCURRENCY_DISABLE_CONSUMERS === "1",
       maxTokensRatio: env.RUN_ENGINE_RELEASE_CONCURRENCY_MAX_TOKENS_RATIO,
+      releasingsMaxAge: env.RUN_ENGINE_RELEASE_CONCURRENCY_RELEASINGS_MAX_AGE,
+      releasingsPollInterval: env.RUN_ENGINE_RELEASE_CONCURRENCY_RELEASINGS_POLL_INTERVAL,
       maxRetries: env.RUN_ENGINE_RELEASE_CONCURRENCY_MAX_RETRIES,
       consumersCount: env.RUN_ENGINE_RELEASE_CONCURRENCY_CONSUMERS_COUNT,
       pollInterval: env.RUN_ENGINE_RELEASE_CONCURRENCY_POLL_INTERVAL,

internal-packages/run-engine/src/engine/index.ts

@@ -182,6 +182,7 @@ export class RunEngine {
       PENDING_CANCEL: 60_000,
       EXECUTING: 60_000,
       EXECUTING_WITH_WAITPOINTS: 60_000,
+      SUSPENDED: 60_000 * 10,
     };
     this.heartbeatTimeouts = {
       ...defaultHeartbeatTimeouts,
@@ -201,6 +202,9 @@ export class RunEngine {
 
     this.releaseConcurrencySystem = new ReleaseConcurrencySystem({
       resources,
+      maxTokensRatio: options.releaseConcurrency?.maxTokensRatio,
+      releasingsMaxAge: options.releaseConcurrency?.releasingsMaxAge,
+      releasingsPollInterval: options.releaseConcurrency?.releasingsPollInterval,
       queueOptions:
         typeof options.releaseConcurrency?.disabled === "boolean" &&
         options.releaseConcurrency.disabled
@@ -223,33 +227,6 @@ export class RunEngine {
               consumersCount: options.releaseConcurrency?.consumersCount ?? 1,
               pollInterval: options.releaseConcurrency?.pollInterval ?? 1000,
               batchSize: options.releaseConcurrency?.batchSize ?? 10,
-              executor: async (descriptor, snapshotId) => {
-                return await this.releaseConcurrencySystem.executeReleaseConcurrencyForSnapshot(
-                  snapshotId
-                );
-              },
-              maxTokens: async (descriptor) => {
-                const environment = await this.prisma.runtimeEnvironment.findFirstOrThrow({
-                  where: { id: descriptor.envId },
-                  select: {
-                    maximumConcurrencyLimit: true,
-                  },
-                });
-
-                return (
-                  environment.maximumConcurrencyLimit *
-                  (options.releaseConcurrency?.maxTokensRatio ?? 1.0)
-                );
-              },
-              keys: {
-                fromDescriptor: (descriptor) =>
-                  `org:${descriptor.orgId}:proj:${descriptor.projectId}:env:${descriptor.envId}`,
-                toDescriptor: (name) => ({
-                  orgId: name.split(":")[1],
-                  projectId: name.split(":")[3],
-                  envId: name.split(":")[5],
-                }),
-              },
               tracer: this.tracer,
             },
     });
@@ -306,6 +283,7 @@ export class RunEngine {
       delayedRunSystem: this.delayedRunSystem,
       machines: this.options.machines,
       retryWarmStartThresholdMs: this.options.retryWarmStartThresholdMs,
+      releaseConcurrencySystem: this.releaseConcurrencySystem,
     });
 
     this.dequeueSystem = new DequeueSystem({
@@ -1297,9 +1275,29 @@ export class RunEngine {
           break;
         }
         case "SUSPENDED": {
-          //todo should we do a periodic check here for whether waitpoints are actually still blocking?
-          //we could at least log some things out if a run has been in this state for a long time
-          throw new NotImplementedError("Not implemented SUSPENDED");
+          const result = await this.waitpointSystem.continueRunIfUnblocked({ runId });
+
+          this.logger.info("handleStalledSnapshot SUSPENDED continueRunIfUnblocked", {
+            runId,
+            result,
+            snapshotId: latestSnapshot.id,
+          });
+
+          switch (result) {
+            case "blocked": {
+              // Reschedule the heartbeat
+              await this.executionSnapshotSystem.restartHeartbeatForRun({
+                runId,
+              });
+              break;
+            }
+            case "unblocked":
+            case "skipped": {
+              break;
+            }
+          }
+
+          break;
         }
         case "PENDING_CANCEL": {
           //if the run is waiting to cancel but the worker hasn't confirmed that,