Add release concurrency metrics sending to otel

Author: ericallam

.vscode/launch.json

@@ -138,7 +138,7 @@
       "type": "node-terminal",
       "request": "launch",
       "name": "Debug RunEngine tests",
-      "command": "pnpm run test ./src/engine/tests/releaseConcurrencyTokenBucketQueue.test.ts -t 'Should provide metrics about queues via getQueueMetrics'",
+      "command": "pnpm run test ./src/engine/tests/releaseConcurrencyTokenBucketQueue.test.ts -t 'Should retrieve metrics for all queues via getQueueMetrics'",
       "cwd": "${workspaceFolder}/internal-packages/run-engine",
       "sourceMaps": true
     },

internal-packages/redis/src/index.ts

@@ -8,7 +8,7 @@ const defaultOptions: Partial<RedisOptions> = {
     const delay = Math.min(times * 50, 1000);
     return delay;
   },
-  maxRetriesPerRequest: process.env.GITHUB_ACTIONS ? 50 : process.env.VITEST ? 1 : 20,
+  maxRetriesPerRequest: process.env.GITHUB_ACTIONS ? 50 : process.env.VITEST ? 5 : 20,
 };
 
 const logger = new Logger("Redis", "debug");

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

@@ -1,8 +1,9 @@
 import { Callback, createRedisClient, Redis, Result, type RedisOptions } from "@internal/redis";
-import { Tracer } from "@internal/tracing";
+import { startSpan, Tracer } from "@internal/tracing";
 import { Logger } from "@trigger.dev/core/logger";
-import { setInterval } from "node:timers/promises";
 import { z } from "zod";
+import { setInterval } from "node:timers/promises";
+import { flattenAttributes } from "@trigger.dev/core/v3";
 
 export type ReleaseConcurrencyQueueRetryOptions = {
   maxRetries?: number;
@@ -81,6 +82,7 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
 
     if (!options.disableConsumers) {
       this.#startConsumers();
+      this.#startMetricsProducer();
     }
   }
 
@@ -397,6 +399,30 @@ 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 });
+
+        await startSpan(
+          this.options.tracer,
+          "ReleaseConcurrencyTokenBucketQueue.metrics",
+          async (span) => {},
+          {
+            attributes: {
+              ...flattenAttributes(metrics, "queues"),
+              forceRecording: true,
+            },
+          }
+        );
+      }
+    } catch (error) {
+      this.logger.error("Error starting metrics producer:", { error });
+    }
+  }
+
   #calculateBackoffScore(item: QueueItemMetadata): string {
     const delay = Math.min(
       this.backoff.maxDelay,
@@ -405,6 +431,137 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
     return String(Date.now() + delay);
   }
 
+  async getQueueMetrics(): Promise<
+    Array<{ releaseQueue: string; currentTokens: number; queueLength: number }>
+  > {
+    const streamRedis = this.redis.duplicate();
+    const queuePattern = `${this.keyPrefix}*:queue`;
+    const stream = streamRedis.scanStream({
+      match: queuePattern,
+      type: "zset",
+      count: 100,
+    });
+
+    let resolvePromise: (
+      value: Array<{ releaseQueue: string; currentTokens: number; queueLength: number }>
+    ) => void;
+    let rejectPromise: (reason?: any) => void;
+
+    const promise = new Promise<
+      Array<{ releaseQueue: string; currentTokens: number; queueLength: number }>
+    >((resolve, reject) => {
+      resolvePromise = resolve;
+      rejectPromise = reject;
+    });
+
+    const metrics: Map<
+      string,
+      { releaseQueue: string; currentTokens: number; queueLength: number }
+    > = new Map();
+
+    async function getMetricsForKeys(queueKeys: string[]) {
+      if (queueKeys.length === 0) {
+        return [];
+      }
+
+      const pipeline = streamRedis.pipeline();
+
+      queueKeys.forEach((queueKey) => {
+        const releaseQueue = queueKey
+          .replace(":queue", "")
+          .replace(streamRedis.options.keyPrefix ?? "", "");
+        const bucketKey = `${releaseQueue}:bucket`;
+
+        pipeline.get(bucketKey);
+        pipeline.zcard(`${releaseQueue}:queue`);
+      });
+
+      const result = await pipeline.exec();
+
+      if (!result) {
+        return [];
+      }
+
+      const results = result.map(([resultError, queueLengthOrCurrentTokens]) => {
+        if (resultError) {
+          return null;
+        }
+
+        return queueLengthOrCurrentTokens ? Number(queueLengthOrCurrentTokens) : 0;
+      });
+
+      // Now zip the results with the queue keys
+      const zippedResults = queueKeys.map((queueKey, index) => {
+        const releaseQueue = queueKey
+          .replace(":queue", "")
+          .replace(streamRedis.options.keyPrefix ?? "", "");
+
+        // Current tokens are at indexes 0, 2, 4, 6, etc.
+        // Queue length are at indexes 1, 3, 5, 7, etc.
+
+        const currentTokens = results[index * 2];
+        const queueLength = results[index * 2 + 1];
+
+        if (typeof currentTokens !== "number" || typeof queueLength !== "number") {
+          return null;
+        }
+
+        return {
+          releaseQueue,
+          currentTokens: currentTokens,
+          queueLength: queueLength,
+        };
+      });
+
+      return zippedResults.filter((result) => result !== null);
+    }
+
+    stream.on("end", () => {
+      streamRedis.quit();
+      resolvePromise(Array.from(metrics.values()));
+    });
+
+    stream.on("error", (error) => {
+      this.logger.error("Error getting queue metrics:", { error });
+
+      stream.pause();
+      streamRedis.quit();
+      rejectPromise(error);
+    });
+
+    stream.on("data", async (keys) => {
+      stream.pause();
+
+      const uniqueKeys = Array.from(new Set<string>(keys));
+
+      if (uniqueKeys.length === 0) {
+        stream.resume();
+        return;
+      }
+
+      const unresolvedKeys = uniqueKeys.filter((key) => !metrics.has(key));
+
+      if (unresolvedKeys.length === 0) {
+        stream.resume();
+        return;
+      }
+
+      this.logger.debug("Fetching queue metrics for keys", { keys: uniqueKeys });
+
+      await getMetricsForKeys(unresolvedKeys).then((results) => {
+        results.forEach((result) => {
+          if (result) {
+            metrics.set(result.releaseQueue, result);
+          }
+        });
+
+        stream.resume();
+      });
+    });
+
+    return promise;
+  }
+
   #registerCommands() {
     this.redis.defineCommand("consumeToken", {
       numberOfKeys: 4,
@@ -424,7 +581,7 @@ local currentTokens = tonumber(redis.call("GET", bucketKey) or maxTokens)
 
 -- If we have enough tokens, then consume them
 if currentTokens >= 1 then
-  newCurrentTokens = currentTokens - 1
+  local newCurrentTokens = currentTokens - 1
 
   redis.call("SET", bucketKey, newCurrentTokens)
   redis.call("ZREM", queueKey, releaserId)

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

@@ -1,7 +1,6 @@
 import { redisTest, StartedRedisContainer } from "@internal/testcontainers";
 import { ReleaseConcurrencyTokenBucketQueue } from "../releaseConcurrencyTokenBucketQueue.js";
 import { setTimeout } from "node:timers/promises";
-import { createRedisClient, Redis } from "@internal/redis";
 
 type TestQueueDescriptor = {
   name: string;
@@ -680,4 +679,38 @@ describe("ReleaseConcurrencyQueue", () => {
       await queue.quit();
     }
   });
+
+  redisTest(
+    "Should retrieve metrics for all queues via getQueueMetrics",
+    async ({ redisContainer }) => {
+      const { queue } = createReleaseConcurrencyQueue(redisContainer, 1);
+
+      // Set up multiple queues with different states
+      await queue.attemptToRelease({ name: "metrics-queue1" }, "run1"); // Consume 1 token from queue1
+
+      // Add more items to queue1 that will be queued due to no tokens
+      await queue.attemptToRelease({ name: "metrics-queue1" }, "run2"); // This will be queued
+      await queue.attemptToRelease({ name: "metrics-queue1" }, "run3"); // This will be queued
+      await queue.attemptToRelease({ name: "metrics-queue1" }, "run4"); // This will be queued
+
+      const metrics = await queue.getQueueMetrics();
+
+      expect(metrics).toHaveLength(1);
+      expect(metrics[0].releaseQueue).toBe("metrics-queue1");
+      expect(metrics[0].currentTokens).toBe(0);
+      expect(metrics[0].queueLength).toBe(3);
+
+      // Now add 10 items to 100 different queues
+      for (let i = 0; i < 100; i++) {
+        for (let j = 0; j < 10; j++) {
+          await queue.attemptToRelease({ name: `metrics-queue2-${i}` }, `run${i}-${j}`);
+        }
+      }
+
+      const metrics2 = await queue.getQueueMetrics();
+      expect(metrics2.length).toBeGreaterThan(90);
+
+      await queue.quit();
+    }
+  );
 });