fix(run-engine): gate counter INCR on ZADD/SADD and 24h TTL the counter

ericallam · ericallam · commit ad7d627ba8a4 · 2026-05-12T09:44:45.000+01:00
diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts
@@ -2577,7 +2577,8 @@ export class RunQueue {
         messageQueue,
         JSON.stringify(message),
         String(messageScore),
-        ckWildcardName
+        ckWildcardName,
+        this.options.redis.keyPrefix ?? ""
       );
     } else {
       await this.redis.nackMessage(
@@ -3461,7 +3462,9 @@ end
 -- Slow path: normal enqueue
 redis.call('SET', messageKey, messageData)
 
--- Lazy-init lengthCounter from existing ckIndex variants (once per base queue, ever).
+-- Lazy-init lengthCounter from existing ckIndex variants (once per base queue per 24h).
+-- The 24h TTL means the counter periodically re-anchors to truth, bounding any drift
+-- that accumulated during rolling-deploy overlap windows.
 -- Run BEFORE the ZADD so we capture pre-state; the subsequent INCR accounts for the new message.
 -- The counter tracks ONLY CK-variant messages — the read path adds ZCARD(base) separately,
 -- so the base zset is intentionally excluded here.
@@ -3471,12 +3474,16 @@ if redis.call('EXISTS', lengthCounterKey) == 0 then
   for _, v in ipairs(variants) do
     total = total + tonumber(redis.call('ZCARD', keyPrefix .. v) or '0')
   end
-  redis.call('SET', lengthCounterKey, total)
+  redis.call('SET', lengthCounterKey, total, 'EX', 86400)
 end
 
-redis.call('ZADD', queueKey, messageScore, messageId)
+-- INCR is gated on ZADD returning 1 (new entry). A duplicate enqueue (same messageId
+-- already in the variant zset) returns 0 and must not bump the counter.
+local added = redis.call('ZADD', queueKey, messageScore, messageId)
 redis.call('ZADD', envQueueKey, messageScore, messageId)
-redis.call('INCR', lengthCounterKey)
+if added == 1 then
+  redis.call('INCR', lengthCounterKey)
+end
 
 -- Rebalance CK index
 local earliest = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES')
@@ -3571,23 +3578,24 @@ end
 -- Slow path: normal enqueue
 redis.call('SET', messageKey, messageData)
 
--- Lazy-init lengthCounter from existing ckIndex variants (once per base queue, ever).
--- Run BEFORE the ZADD so we capture pre-state; the subsequent INCR accounts for the new message.
--- The counter tracks ONLY CK-variant messages — the read path adds ZCARD(base) separately,
--- so the base zset is intentionally excluded here.
+-- Lazy-init lengthCounter from existing ckIndex variants (once per base queue per 24h).
+-- See enqueueMessageCkTracked for the TTL rationale.
 if redis.call('EXISTS', lengthCounterKey) == 0 then
   local total = 0
   local variants = redis.call('ZRANGE', ckIndexKey, 0, -1)
   for _, v in ipairs(variants) do
     total = total + tonumber(redis.call('ZCARD', keyPrefix .. v) or '0')
   end
-  redis.call('SET', lengthCounterKey, total)
+  redis.call('SET', lengthCounterKey, total, 'EX', 86400)
 end
 
-redis.call('ZADD', queueKey, messageScore, messageId)
+-- INCR is gated on ZADD returning 1 (new entry).
+local added = redis.call('ZADD', queueKey, messageScore, messageId)
 redis.call('ZADD', envQueueKey, messageScore, messageId)
 redis.call('ZADD', ttlQueueKey, ttlScore, ttlMember)
-redis.call('INCR', lengthCounterKey)
+if added == 1 then
+  redis.call('INCR', lengthCounterKey)
+end
 
 -- Rebalance CK index
 local earliest = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES')
@@ -3692,7 +3700,7 @@ for i, member in ipairs(expiredMembers) do
       redis.call('SREM', envDequeuedKey, runId)
 
       -- Rebalance CK index if this is a CK queue
-      local ckMatch = string.match(rawQueueKey, "^(.+):ck:.+$")
+      local ckMatch = string.match(rawQueueKey, "(.-):ck:")
       if ckMatch then
         local ckIndexKey = keyPrefix .. ckMatch .. ":ckIndex"
         local earliest = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES')
@@ -3797,7 +3805,7 @@ for i, member in ipairs(expiredMembers) do
       redis.call('SREM', envDequeuedKey, runId)
 
       -- Rebalance CK index AND update counters if this is a CK queue
-      local ckMatch = string.match(rawQueueKey, "^(.+):ck:.+$")
+      local ckMatch = string.match(rawQueueKey, "(.-):ck:")
       if ckMatch then
         local lengthCounterKey = keyPrefix .. ckMatch .. ":lengthCounter"
         local runningCounterKey = keyPrefix .. ckMatch .. ":runningCounter"
@@ -4325,11 +4333,17 @@ local queue = messageData.queue
 local queueCurrentDequeuedKey = keyPrefix .. queue .. ':currentDequeued'
 local envCurrentDequeuedKey = keyPrefix .. string.match(queue, "(.+):queue:") .. ':currentDequeued'
 
--- If CK, bump the runningCounter for the base queue (lazy init from ckIndex).
+-- SADD first so we know if this dequeue is new (return 1) or a duplicate (return 0).
+-- INCR runningCounter is gated on the new-membership result so re-dequeues don't inflate.
+local addedDeq = redis.call('SADD', queueCurrentDequeuedKey, messageData.runId)
+redis.call('SADD', envCurrentDequeuedKey, messageData.runId)
+
+-- If CK + new addition, bump the runningCounter for the base queue (lazy init from ckIndex).
 -- The counter tracks ONLY CK-variant currentDequeued — the read path adds the base
--- SCARD separately, so we exclude the base currentDequeued here.
+-- SCARD separately, so we exclude the base currentDequeued here. 24h TTL on init
+-- bounds any drift from rolling-deploy v1/v2 overlap.
 local baseQueue = string.match(queue, "(.-):ck:")
-if baseQueue then
+if baseQueue and addedDeq == 1 then
   local runningCounterKey = keyPrefix .. baseQueue .. ':runningCounter'
   if redis.call('EXISTS', runningCounterKey) == 0 then
     local ckIndexKey = keyPrefix .. baseQueue .. ':ckIndex'
@@ -4338,14 +4352,14 @@ if baseQueue then
     for _, v in ipairs(variants) do
       total = total + tonumber(redis.call('SCARD', keyPrefix .. v .. ':currentDequeued') or '0')
     end
-    redis.call('SET', runningCounterKey, total)
+    -- Subtract 1 because we already SADD'd our own runId above, which the per-variant
+    -- SCARDs now reflect. Without this, the lazy-init seed double-counts the new member
+    -- relative to what the subsequent INCR will add.
+    redis.call('SET', runningCounterKey, math.max(0, total - 1), 'EX', 86400)
   end
   redis.call('INCR', runningCounterKey)
 end
 
-redis.call('SADD', queueCurrentDequeuedKey, messageData.runId)
-redis.call('SADD', envCurrentDequeuedKey, messageData.runId)
-
 return message
       `,
     });
@@ -4749,6 +4763,8 @@ local messageQueueName = ARGV[2]
 local messageData = ARGV[3]
 local messageScore = tonumber(ARGV[4])
 local ckWildcardName = ARGV[5]
+-- keyPrefix for prepending to variant names stored as values in ckIndex (lazy-init only)
+local keyPrefix = ARGV[6]
 
 local function decrFloored(key)
   if tonumber(redis.call('GET', key) or '0') > 0 then
@@ -4766,6 +4782,19 @@ if removedFromDequeued == 1 then
   decrFloored(runningCounterKey)
 end
 
+-- Lazy-init lengthCounter if missing (e.g. expired via 24h TTL). nack re-queues a
+-- message, which means lengthCounter must be present before we INCR. Without this,
+-- a nack after counter expiry would create the counter at 1 and stay drifted until
+-- next reset.
+if redis.call('EXISTS', lengthCounterKey) == 0 then
+  local total = 0
+  local variants = redis.call('ZRANGE', ckIndexKey, 0, -1)
+  for _, v in ipairs(variants) do
+    total = total + tonumber(redis.call('ZCARD', keyPrefix .. v) or '0')
+  end
+  redis.call('SET', lengthCounterKey, total, 'EX', 86400)
+end
+
 -- ZADD back to the variant zset. INCR lengthCounter only if it's a new entry.
 local added = redis.call('ZADD', messageQueueKey, messageScore, messageId)
 redis.call('ZADD', envQueueKey, messageScore, messageId)
@@ -5511,6 +5540,7 @@ declare module "@internal/redis" {
       messageData: string,
       messageScore: string,
       ckWildcardName: string,
+      keyPrefix: string,
       callback?: Callback<void>
     ): Result<void, Context>;
 
diff --git a/internal-packages/run-engine/src/run-queue/tests/ckCounters.test.ts b/internal-packages/run-engine/src/run-queue/tests/ckCounters.test.ts
@@ -2,7 +2,6 @@ import { redisTest } from "@internal/testcontainers";
 import { trace } from "@internal/tracing";
 import { Logger } from "@trigger.dev/core/logger";
 import { describe } from "vitest";
-import { setTimeout } from "node:timers/promises";
 import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js";
 import { RunQueue } from "../index.js";
 import { RunQueueFullKeyProducer } from "../keyProducer.js";
@@ -327,4 +326,105 @@ describe("CK base-queue counters", () => {
       }
     }
   );
+
+  redisTest(
+    "lengthCounter has 24h TTL after lazy-init",
+    async ({ redisContainer }) => {
+      const queue = createQueue(redisContainer);
+      try {
+        await queue.enqueueMessage({
+          env: authenticatedEnvDev,
+          message: makeMessage({ runId: "r1", concurrencyKey: "ck-a" }),
+          workerQueue: authenticatedEnvDev.id,
+          skipDequeueProcessing: true,
+        });
+
+        const counterKey = testOptions.keys.queueLengthCounterKey(
+          authenticatedEnvDev,
+          "task/my-task"
+        );
+        const ttl = await queue.redis.ttl(counterKey);
+        // Expect roughly 86400; allow slack for test scheduling.
+        expect(ttl).toBeGreaterThan(86000);
+        expect(ttl).toBeLessThanOrEqual(86400);
+      } finally {
+        await queue.quit();
+      }
+    }
+  );
+
+  redisTest(
+    "duplicate CK enqueue (same runId) does not inflate lengthCounter",
+    async ({ redisContainer }) => {
+      const queue = createQueue(redisContainer);
+      try {
+        const msg = makeMessage({ runId: "r1", concurrencyKey: "ck-a" });
+
+        // First enqueue: counter goes 0 -> 1
+        await queue.enqueueMessage({
+          env: authenticatedEnvDev,
+          message: msg,
+          workerQueue: authenticatedEnvDev.id,
+          skipDequeueProcessing: true,
+        });
+
+        // Same runId again: ZADD returns 0 (already in zset), counter must stay at 1
+        await queue.enqueueMessage({
+          env: authenticatedEnvDev,
+          message: msg,
+          workerQueue: authenticatedEnvDev.id,
+          skipDequeueProcessing: true,
+        });
+
+        expect(await queue.lengthOfQueue(authenticatedEnvDev, msg.queue)).toBe(1);
+      } finally {
+        await queue.quit();
+      }
+    }
+  );
+
+  redisTest(
+    "nack lazy-inits lengthCounter when it expired",
+    async ({ redisContainer }) => {
+      const queue = createQueue(redisContainer);
+      try {
+        const msg = makeMessage({ runId: "r1", concurrencyKey: "ck-a" });
+        // Seed three messages on the CK variant so the lazy-init has a non-trivial floor.
+        for (let i = 0; i < 3; i++) {
+          await queue.enqueueMessage({
+            env: authenticatedEnvDev,
+            message: makeMessage({ runId: `seed-${i}`, concurrencyKey: "ck-a" }),
+            workerQueue: authenticatedEnvDev.id,
+            skipDequeueProcessing: true,
+          });
+        }
+
+        // Simulate counter expiry (the 24h TTL kicked in).
+        const counterKey = testOptions.keys.queueLengthCounterKey(
+          authenticatedEnvDev,
+          "task/my-task"
+        );
+        await queue.redis.del(counterKey);
+        expect(await queue.redis.exists(counterKey)).toBe(0);
+
+        // Dequeue one to currentConcurrency so we have something to nack back.
+        const shard = testOptions.keys.masterQueueShardForEnvironment(msg.environmentId, 2);
+        await queue.testDequeueFromMasterQueue(shard, msg.environmentId, 1);
+
+        // Nack a CK message. nackMessageCkTracked should lazy-init the counter
+        // (find 2 already in zset + 1 we're re-queuing) rather than starting from 1.
+        await queue.nackMessage({
+          orgId: msg.orgId,
+          messageId: "seed-0",
+          skipDequeueProcessing: true,
+        });
+
+        // 3 originals, 1 was dequeued (still re-queued by nack), counter should now reflect all 3.
+        const observed = await queue.lengthOfQueue(authenticatedEnvDev, msg.queue);
+        expect(observed).toBe(3);
+      } finally {
+        await queue.quit();
+      }
+    }
+  );
 });
