Fix CodeRabbit suggestions

Author: ericallam

apps/webapp/app/env.server.ts

@@ -472,6 +472,7 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
   RUN_ENGINE_RETRY_WARM_START_THRESHOLD_MS: z.coerce.number().int().default(30_000),
   RUN_ENGINE_PROCESS_WORKER_QUEUE_DEBOUNCE_MS: z.coerce.number().int().default(200),
+  RUN_ENGINE_DEQUEUE_BLOCKING_TIMEOUT_SECONDS: z.coerce.number().int().default(10),
 
   RUN_ENGINE_WORKER_REDIS_HOST: z
     .string()

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

@@ -61,6 +61,7 @@ function createRunEngine() {
       },
       shardCount: env.RUN_ENGINE_RUN_QUEUE_SHARD_COUNT,
       processWorkerQueueDebounceMs: env.RUN_ENGINE_PROCESS_WORKER_QUEUE_DEBOUNCE_MS,
+      dequeueBlockingTimeoutSeconds: env.RUN_ENGINE_DEQUEUE_BLOCKING_TIMEOUT_SECONDS,
     },
     runLock: {
       redis: {

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

@@ -344,21 +344,15 @@ function configurePrismaMetrics(meter: Meter) {
     return { total, busy, free };
   }
 
-  // Register callbacks (one scrape == one DB call) --------------------------
-  totalGauge.addCallback(async (res) => {
-    const { total } = await readPoolCounters();
-    res.observe(total);
-  });
-
-  busyGauge.addCallback(async (res) => {
-    const { busy } = await readPoolCounters();
-    res.observe(busy);
-  });
-
-  freeGauge.addCallback(async (res) => {
-    const { free } = await readPoolCounters();
-    res.observe(free);
-  });
+  meter.addBatchObservableCallback(
+    async (res) => {
+      const { total, busy, free } = await readPoolCounters();
+      res.observe(totalGauge, total);
+      res.observe(busyGauge, busy);
+      res.observe(freeGauge, free);
+    },
+    [totalGauge, busyGauge, freeGauge]
+  );
 }
 
 const SemanticEnvResources = {

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

@@ -125,6 +125,7 @@ export class RunEngine {
       shardCount: options.queue?.shardCount,
       masterQueueConsumersDisabled: options.worker.disabled,
       processWorkerQueueDebounceMs: options.queue?.processWorkerQueueDebounceMs,
+      dequeueBlockingTimeoutSeconds: options.queue?.dequeueBlockingTimeoutSeconds,
       meter: options.meter,
     });
 

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

@@ -79,8 +79,6 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
   private batchSize: number;
   private maxRetries: number;
   private backoff: NonNullable<Required<ReleaseConcurrencyQueueRetryOptions["backoff"]>>;
-  private _lastReleasingsLength: number = 0;
-  private _lastMasterQueueLength: number = 0;
 
   constructor(private readonly options: ReleaseConcurrencyQueueOptions<T>) {
     this.redis = createRedisClient(options.redis);
@@ -129,7 +127,6 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
 
     if (!options.disableConsumers) {
       this.#startConsumers();
-      this.#startMetricsProducer();
       this.#startReleasingsSweeper();
     }
   }
@@ -140,11 +137,13 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
   }
 
   async #updateReleasingsLength(observableResult: ObservableResult<Attributes>) {
-    observableResult.observe(this._lastReleasingsLength);
+    const releasingsLength = await this.redis.zcard(this.#releasingsKey());
+    observableResult.observe(releasingsLength);
   }
 
   async #updateMasterQueueLength(observableResult: ObservableResult<Attributes>) {
-    observableResult.observe(this._lastMasterQueueLength);
+    const masterQueueLength = await this.redis.zcard(this.masterQueuesKey);
+    observableResult.observe(masterQueueLength);
   }
 
   /**
@@ -532,36 +531,6 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
     }
   }
 
-  async #startMetricsProducer() {
-    try {
-      // Produce metrics every 60 seconds, using a tracer span
-      for await (const _ of setInterval(60_000)) {
-        const metrics = await this.getQueueMetrics();
-        this.logger.info("Queue metrics:", { metrics });
-
-        // Update cached values for OpenTelemetry observable gauges
-        this._lastReleasingsLength = await this.redis.zcard(this.#releasingsKey());
-        this._lastMasterQueueLength = await this.redis.zcard(this.masterQueuesKey);
-
-        await startSpan(
-          this.options.tracer,
-          "ReleaseConcurrencyTokenBucketQueue.metrics",
-          async (span) => {},
-          {
-            attributes: {
-              ...flattenAttributes(metrics, "queues"),
-              releasingsLength: this._lastReleasingsLength,
-              masterQueueLength: this._lastMasterQueueLength,
-              forceRecording: true,
-            },
-          }
-        );
-      }
-    } catch (error) {
-      this.logger.error("Error starting metrics producer:", { error });
-    }
-  }
-
   #calculateBackoffScore(item: QueueItemMetadata): string {
     const delay = Math.min(
       this.backoff.maxDelay,

internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -78,7 +78,7 @@ export class DequeueSystem {
         //lock the run so nothing else can modify it
         try {
           const dequeuedRun = await this.$.runLock.lock(
-            "dequeueFromMasterQueue",
+            "dequeueFromWorkerQueue",
             [runId],
             5000,
             async (signal) => {
@@ -135,7 +135,7 @@ export class DequeueSystem {
                   tx: prisma,
                 });
                 this.$.logger.error(
-                  `RunEngine.dequeueFromMasterQueue(): Run is not in a valid state to be dequeued: ${runId}\n ${snapshot.id}:${snapshot.executionStatus}`
+                  `RunEngine.dequeueFromWorkerQueue(): Run is not in a valid state to be dequeued: ${runId}\n ${snapshot.id}:${snapshot.executionStatus}`
                 );
 
                 return;
@@ -192,7 +192,7 @@ export class DequeueSystem {
                 switch (result.code) {
                   case "NO_RUN": {
                     //this should not happen, the run is unrecoverable so we'll ack it
-                    this.$.logger.error("RunEngine.dequeueFromMasterQueue(): No run found", {
+                    this.$.logger.error("RunEngine.dequeueFromWorkerQueue(): No run found", {
                       runId,
                       latestSnapshot: snapshot.id,
                     });
@@ -202,7 +202,7 @@ export class DequeueSystem {
                   case "RUN_ENVIRONMENT_ARCHIVED": {
                     //this happens if the preview branch was archived
                     this.$.logger.warn(
-                      "RunEngine.dequeueFromMasterQueue(): Run environment archived",
+                      "RunEngine.dequeueFromWorkerQueue(): Run environment archived",
                       {
                         runId,
                         latestSnapshot: snapshot.id,
@@ -216,7 +216,7 @@ export class DequeueSystem {
                   case "TASK_NEVER_REGISTERED":
                   case "QUEUE_NOT_FOUND":
                   case "TASK_NOT_IN_LATEST": {
-                    this.$.logger.warn(`RunEngine.dequeueFromMasterQueue(): ${result.code}`, {
+                    this.$.logger.warn(`RunEngine.dequeueFromWorkerQueue(): ${result.code}`, {
                       runId,
                       latestSnapshot: snapshot.id,
                       result,
@@ -234,7 +234,7 @@ export class DequeueSystem {
                   }
                   case "BACKGROUND_WORKER_MISMATCH": {
                     this.$.logger.warn(
-                      "RunEngine.dequeueFromMasterQueue(): Background worker mismatch",
+                      "RunEngine.dequeueFromWorkerQueue(): Background worker mismatch",
                       {
                         runId,
                         latestSnapshot: snapshot.id,
@@ -256,7 +256,7 @@ export class DequeueSystem {
               //check for a valid deployment if it's not a development environment
               if (result.run.runtimeEnvironment.type !== "DEVELOPMENT") {
                 if (!result.deployment || !result.deployment.imageReference) {
-                  this.$.logger.warn("RunEngine.dequeueFromMasterQueue(): No deployment found", {
+                  this.$.logger.warn("RunEngine.dequeueFromWorkerQueue(): No deployment found", {
                     runId,
                     latestSnapshot: snapshot.id,
                     result,
@@ -289,7 +289,7 @@ export class DequeueSystem {
                 const retryConfig = result.task.retryConfig;
 
                 this.$.logger.debug(
-                  "RunEngine.dequeueFromMasterQueue(): maxAttempts not set, using task's retry config",
+                  "RunEngine.dequeueFromWorkerQueue(): maxAttempts not set, using task's retry config",
                   {
                     runId,
                     task: result.task.id,
@@ -300,7 +300,7 @@ export class DequeueSystem {
                 const parsedConfig = RetryOptions.nullable().safeParse(retryConfig);
 
                 if (!parsedConfig.success) {
-                  this.$.logger.error("RunEngine.dequeueFromMasterQueue(): Invalid retry config", {
+                  this.$.logger.error("RunEngine.dequeueFromWorkerQueue(): Invalid retry config", {
                     runId,
                     task: result.task.id,
                     rawRetryConfig: retryConfig,
@@ -373,7 +373,7 @@ export class DequeueSystem {
               });
 
               if (!lockedTaskRun) {
-                this.$.logger.error("RunEngine.dequeueFromMasterQueue(): Failed to lock task run", {
+                this.$.logger.error("RunEngine.dequeueFromWorkerQueue(): Failed to lock task run", {
                   taskRun: result.run.id,
                   taskIdentifier: result.run.taskIdentifier,
                   deployment: result.deployment?.id,
@@ -464,7 +464,7 @@ export class DequeueSystem {
           return dequeuedRun;
         } catch (error) {
           this.$.logger.error(
-            "RunEngine.dequeueFromMasterQueue(): Thrown error while preparing run to be run",
+            "RunEngine.dequeueFromWorkerQueue(): Thrown error while preparing run to be run",
             {
               error,
               runId,
@@ -481,7 +481,7 @@ export class DequeueSystem {
           if (!run) {
             //this isn't ideal because we're not creating a snapshot… but we can't do much else
             this.$.logger.error(
-              "RunEngine.dequeueFromMasterQueue(): Thrown error, then run not found. Nacking.",
+              "RunEngine.dequeueFromWorkerQueue(): Thrown error, then run not found. Nacking.",
               {
                 runId,
                 orgId,
@@ -507,7 +507,7 @@ export class DequeueSystem {
           });
 
           if (!gotRequeued) {
-            this.$.logger.error("RunEngine.dequeueFromMasterQueue(): Failed to requeue run", {
+            this.$.logger.error("RunEngine.dequeueFromWorkerQueue(): Failed to requeue run", {
               runId,
               orgId,
             });
@@ -570,7 +570,7 @@ export class DequeueSystem {
             },
           });
 
-          this.$.logger.debug("RunEngine.dequeueFromMasterQueue(): Pending version", {
+          this.$.logger.debug("RunEngine.dequeueFromWorkerQueue(): Pending version", {
             runId,
             run,
           });

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

@@ -39,6 +39,7 @@ export type RunEngineOptions = {
       FairQueueSelectionStrategyOptions,
       "parentQueueLimit" | "tracer" | "biases" | "reuseSnapshotCount" | "maximumEnvCount"
     >;
+    dequeueBlockingTimeoutSeconds?: number;
   };
   runLock: {
     redis: RedisOptions;

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

@@ -76,6 +76,7 @@ export type RunQueueOptions = {
     disabled?: boolean;
   };
   meter?: Meter;
+  dequeueBlockingTimeoutSeconds?: number;
 };
 
 type DequeuedMessage = {
@@ -510,9 +511,6 @@ export class RunQueue {
     return this.#trace(
       "dequeueMessageFromWorkerQueue",
       async (span) => {
-        // TODO: this is where we read from the worker queue, which is a list of message IDs
-        // We'll perform a BLPOP on the worker list, and then read the message from the message list
-        // We'll then return the message
         const dequeuedMessage = await this.#callDequeueMessageFromWorkerQueue({
           workerQueue,
         });
@@ -1376,8 +1374,7 @@ export class RunQueue {
       workerQueueKey,
       //args
       this.options.redis.keyPrefix ?? "",
-      // TODO: make this configurable
-      String(10)
+      String(this.options.dequeueBlockingTimeoutSeconds ?? 10)
     );
 
     this.abortController.signal.removeEventListener("abort", cleanup);