Add env vars and additional spans

Author: ericallam

apps/webapp/app/env.server.ts

@@ -496,7 +496,10 @@ const EnvironmentSchema = z.object({
   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_MASTER_QUEUE_CONSUMERS_INTERVAL_MS: z.coerce.number().int().default(500),
+  RUN_ENGINE_MASTER_QUEUE_CONSUMERS_INTERVAL_MS: z.coerce.number().int().default(1000),
+  RUN_ENGINE_MASTER_QUEUE_COOLOFF_PERIOD_MS: z.coerce.number().int().default(10_000),
+  RUN_ENGINE_MASTER_QUEUE_COOLOFF_COUNT_THRESHOLD: z.coerce.number().int().default(10),
+  RUN_ENGINE_MASTER_QUEUE_CONSUMER_DEQUEUE_COUNT: z.coerce.number().int().default(10),
   RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_SCHEDULE: z.string().optional(),
   RUN_ENGINE_CONCURRENCY_SWEEPER_PROCESS_MARKED_SCHEDULE: z.string().optional(),
   RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_JITTER_IN_MS: z.coerce.number().int().optional(),

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

@@ -67,6 +67,9 @@ function createRunEngine() {
       dequeueBlockingTimeoutSeconds: env.RUN_ENGINE_DEQUEUE_BLOCKING_TIMEOUT_SECONDS,
       masterQueueConsumersIntervalMs: env.RUN_ENGINE_MASTER_QUEUE_CONSUMERS_INTERVAL_MS,
       masterQueueConsumersDisabled: env.RUN_ENGINE_WORKER_ENABLED === "0",
+      masterQueueCooloffPeriodMs: env.RUN_ENGINE_MASTER_QUEUE_COOLOFF_PERIOD_MS,
+      masterQueueCooloffCountThreshold: env.RUN_ENGINE_MASTER_QUEUE_COOLOFF_COUNT_THRESHOLD,
+      masterQueueConsumerDequeueCount: env.RUN_ENGINE_MASTER_QUEUE_CONSUMER_DEQUEUE_COUNT,
       concurrencySweeper: {
         scanSchedule: env.RUN_ENGINE_CONCURRENCY_SWEEPER_SCAN_SCHEDULE,
         processMarkedSchedule: env.RUN_ENGINE_CONCURRENCY_SWEEPER_PROCESS_MARKED_SCHEDULE,

apps/webapp/app/v3/services/worker/workerGroupTokenService.server.ts

@@ -193,8 +193,10 @@ export class WorkerGroupTokenService extends WithRunEngine {
       return;
     }
 
+    const cacheKey = ["worker-group-token", token, instanceName];
+
     const result = await authenticatedWorkerInstanceCache.authenticatedWorkerInstance.swr(
-      `worker-group-token-${token}`,
+      cacheKey.join("-"),
       async () => {
         const workerGroup = await this.findWorkerGroup({ token });
 

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

@@ -217,7 +217,16 @@ export class RunLocker {
     let lastError: Error | undefined;
 
     for (let attempt = 0; attempt <= maxAttempts; attempt++) {
-      const [error, acquiredLock] = await tryCatch(this.redlock.acquire(sortedResources, duration));
+      const [error, acquiredLock] = await tryCatch(
+        startSpan(this.tracer, "RunLocker.acquireLock", async (span) => {
+          span.setAttributes({
+            resources: joinedResources,
+            attempt,
+            totalWaitTime,
+          });
+          return await this.redlock.acquire(sortedResources, duration);
+        })
+      );
 
       if (!error && acquiredLock) {
         lock = acquiredLock;
@@ -390,7 +399,15 @@ export class RunLocker {
       this.#cleanupExtension(manualContext);
 
       // Release the lock using tryCatch
-      const [releaseError] = await tryCatch(lock.release());
+      const [releaseError] = await tryCatch(
+        startSpan(this.tracer, "RunLocker.releaseLock", async (span) => {
+          span.setAttributes({
+            resources: joinedResources,
+            lockValue: lock.value,
+          });
+          return await lock.release();
+        })
+      );
       if (releaseError) {
         this.logger.warn("[RunLocker] Error releasing lock", {
           error: releaseError,

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

@@ -39,6 +39,9 @@ export type RunEngineOptions = {
     masterQueueConsumersDisabled?: boolean;
     processWorkerQueueDebounceMs?: number;
     masterQueueConsumersIntervalMs?: number;
+    masterQueueCooloffPeriodMs?: number;
+    masterQueueCooloffCountThreshold?: number;
+    masterQueueConsumerDequeueCount?: number;
     workerOptions?: WorkerConcurrencyOptions;
     retryOptions?: RetryOptions;
     defaultEnvConcurrency?: number;

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

@@ -1451,92 +1451,107 @@ export class RunQueue {
     shard: number;
     maxCount: number;
   }): Promise<DequeuedMessage[]> {
-    const queueConcurrencyLimitKey = this.keys.queueConcurrencyLimitKeyFromQueue(messageQueue);
-    const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKeyFromQueue(messageQueue);
-    const envConcurrencyLimitKey = this.keys.envConcurrencyLimitKeyFromQueue(messageQueue);
-    const envConcurrencyLimitBurstFactorKey =
-      this.keys.envConcurrencyLimitBurstFactorKeyFromQueue(messageQueue);
-    const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue);
-    const messageKeyPrefix = this.keys.messageKeyPrefixFromQueue(messageQueue);
-    const envQueueKey = this.keys.envQueueKeyFromQueue(messageQueue);
-    const masterQueueKey = this.keys.masterQueueKeyForShard(shard);
-
-    this.logger.debug("#callDequeueMessagesFromQueue", {
-      messageQueue,
-      queueConcurrencyLimitKey,
-      envConcurrencyLimitKey,
-      envConcurrencyLimitBurstFactorKey,
-      queueCurrentConcurrencyKey,
-      envCurrentConcurrencyKey,
-      messageKeyPrefix,
-      envQueueKey,
-      masterQueueKey,
-      shard,
-      maxCount,
-    });
+    return this.#trace("callDequeueMessagesFromQueue", async (span) => {
+      span.setAttributes({
+        messageQueue,
+        shard,
+        maxCount,
+      });
 
-    const result = await this.redis.dequeueMessagesFromQueue(
-      //keys
-      messageQueue,
-      queueConcurrencyLimitKey,
-      envConcurrencyLimitKey,
-      envConcurrencyLimitBurstFactorKey,
-      queueCurrentConcurrencyKey,
-      envCurrentConcurrencyKey,
-      messageKeyPrefix,
-      envQueueKey,
-      masterQueueKey,
-      //args
-      messageQueue,
-      String(Date.now()),
-      String(this.options.defaultEnvConcurrency),
-      String(this.options.defaultEnvConcurrencyBurstFactor ?? 1),
-      this.options.redis.keyPrefix ?? "",
-      String(maxCount)
-    );
+      const queueConcurrencyLimitKey = this.keys.queueConcurrencyLimitKeyFromQueue(messageQueue);
+      const queueCurrentConcurrencyKey =
+        this.keys.queueCurrentConcurrencyKeyFromQueue(messageQueue);
+      const envConcurrencyLimitKey = this.keys.envConcurrencyLimitKeyFromQueue(messageQueue);
+      const envConcurrencyLimitBurstFactorKey =
+        this.keys.envConcurrencyLimitBurstFactorKeyFromQueue(messageQueue);
+      const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(messageQueue);
+      const messageKeyPrefix = this.keys.messageKeyPrefixFromQueue(messageQueue);
+      const envQueueKey = this.keys.envQueueKeyFromQueue(messageQueue);
+      const masterQueueKey = this.keys.masterQueueKeyForShard(shard);
 
-    if (!result) {
-      return [];
-    }
+      this.logger.debug("#callDequeueMessagesFromQueue", {
+        messageQueue,
+        queueConcurrencyLimitKey,
+        envConcurrencyLimitKey,
+        envConcurrencyLimitBurstFactorKey,
+        queueCurrentConcurrencyKey,
+        envCurrentConcurrencyKey,
+        messageKeyPrefix,
+        envQueueKey,
+        masterQueueKey,
+        shard,
+        maxCount,
+      });
 
-    this.logger.debug("dequeueMessagesFromQueue raw result", {
-      result,
-      service: this.name,
-    });
+      const result = await this.redis.dequeueMessagesFromQueue(
+        //keys
+        messageQueue,
+        queueConcurrencyLimitKey,
+        envConcurrencyLimitKey,
+        envConcurrencyLimitBurstFactorKey,
+        queueCurrentConcurrencyKey,
+        envCurrentConcurrencyKey,
+        messageKeyPrefix,
+        envQueueKey,
+        masterQueueKey,
+        //args
+        messageQueue,
+        String(Date.now()),
+        String(this.options.defaultEnvConcurrency),
+        String(this.options.defaultEnvConcurrencyBurstFactor ?? 1),
+        this.options.redis.keyPrefix ?? "",
+        String(maxCount)
+      );
 
-    const messages = [];
-    for (let i = 0; i < result.length; i += 3) {
-      const messageId = result[i];
-      const messageScore = result[i + 1];
-      const rawMessage = result[i + 2];
+      if (!result) {
+        span.setAttribute("message_count", 0);
 
-      //read message
-      const parsedMessage = OutputPayload.safeParse(JSON.parse(rawMessage));
-      if (!parsedMessage.success) {
-        this.logger.error(`[${this.name}] Failed to parse message`, {
+        return [];
+      }
+
+      this.logger.debug("dequeueMessagesFromQueue raw result", {
+        result,
+        service: this.name,
+      });
+
+      const messages = [];
+      for (let i = 0; i < result.length; i += 3) {
+        const messageId = result[i];
+        const messageScore = result[i + 1];
+        const rawMessage = result[i + 2];
+
+        //read message
+        const parsedMessage = OutputPayload.safeParse(JSON.parse(rawMessage));
+        if (!parsedMessage.success) {
+          this.logger.error(`[${this.name}] Failed to parse message`, {
+            messageId,
+            error: parsedMessage.error,
+            service: this.name,
+          });
+
+          continue;
+        }
+
+        const message = parsedMessage.data;
+
+        messages.push({
           messageId,
-          error: parsedMessage.error,
-          service: this.name,
+          messageScore,
+          message,
         });
-
-        continue;
       }
 
-      const message = parsedMessage.data;
-
-      messages.push({
-        messageId,
-        messageScore,
-        message,
+      this.logger.debug("dequeueMessagesFromQueue parsed result", {
+        messages,
+        service: this.name,
       });
-    }
 
-    this.logger.debug("dequeueMessagesFromQueue parsed result", {
-      messages,
-      service: this.name,
-    });
+      const filteredMessages = messages.filter(Boolean) as DequeuedMessage[];
+
+      span.setAttribute("message_count", filteredMessages.length);
 
-    return messages.filter(Boolean) as DequeuedMessage[];
+      return filteredMessages;
+    });
   }
 
   async #callDequeueMessageFromWorkerQueue({
@@ -1569,7 +1584,16 @@ export class RunQueue {
 
       this.abortController.signal.addEventListener("abort", cleanup);
 
-      const result = await blockingClient.blpop(workerQueueKey, blockingPopTimeoutSeconds);
+      const result = await this.#trace("popMessageFromWorkerQueue", async (span) => {
+        span.setAttributes({
+          workerQueue,
+          workerQueueKey,
+          blockingPopTimeoutSeconds,
+          blocking: true,
+        });
+
+        return await blockingClient.blpop(workerQueueKey, blockingPopTimeoutSeconds);
+      });
 
       this.abortController.signal.removeEventListener("abort", cleanup);
 
@@ -1607,7 +1631,15 @@ export class RunQueue {
 
       const [, messageKey] = result;
 
-      const workerQueueLength = await this.redis.llen(workerQueueKey);
+      const workerQueueLength = await this.#trace("getWorkerQueueLength", async (span) => {
+        span.setAttributes({
+          workerQueue,
+          workerQueueKey,
+        });
+
+        return await this.redis.llen(workerQueueKey);
+      });
+
       const message = await this.#dequeueMessageFromKey(messageKey);
 
       if (!message) {
@@ -1626,7 +1658,15 @@ export class RunQueue {
         workerQueueKey,
       });
 
-      const result = await this.redis.dequeueMessageFromWorkerQueueNonBlocking(workerQueueKey);
+      const result = await this.#trace("popMessageFromWorkerQueue", async (span) => {
+        span.setAttributes({
+          workerQueue,
+          workerQueueKey,
+          blocking: false,
+        });
+
+        return await this.redis.dequeueMessageFromWorkerQueueNonBlocking(workerQueueKey);
+      });
 
       if (!result) {
         return;
@@ -2070,27 +2110,42 @@ export class RunQueue {
   }
 
   async #dequeueMessageFromKey(messageKey: string) {
-    const rawMessage = await this.redis.dequeueMessageFromKey(
-      messageKey,
-      this.options.redis.keyPrefix ?? ""
-    );
+    return this.#trace("dequeueMessageFromKey", async (span) => {
+      span.setAttributes({
+        messageKey,
+      });
 
-    if (!rawMessage) {
-      return;
-    }
+      const rawMessage = await this.redis.dequeueMessageFromKey(
+        messageKey,
+        this.options.redis.keyPrefix ?? ""
+      );
 
-    const [error, message] = parseRawMessage(rawMessage);
+      if (!rawMessage) {
+        span.setAttribute("result", "NO_MESSAGE");
 
-    if (error) {
-      this.logger.error(`[${this.name}] Failed to parse message`, {
-        messageKey,
-        error,
-        service: this.name,
-        message: message ?? rawMessage,
-      });
-    }
+        return;
+      }
 
-    return message;
+      const [error, message] = parseRawMessage(rawMessage);
+
+      if (error) {
+        this.logger.error(`[${this.name}] Failed to parse message`, {
+          messageKey,
+          error,
+          service: this.name,
+          message: message ?? rawMessage,
+        });
+      }
+
+      if (message) {
+        span.setAttribute("result", "SUCCESS");
+        span.setAttribute("messageId", message.runId);
+      } else {
+        span.setAttribute("result", "NO_MESSAGE");
+      }
+
+      return message;
+    });
   }
 
   #registerCommands() {