fix(run-engine): keep parked ck credit when pass 2 repairs a variant

Pass 2 repairs a variant that reached ckIndex with no ckVtime entry, which
during a mixed rollout is one a flag-off instance re-enqueued after a flag-on
instance had parked its tag. Both of its routes repaired at the floor and threw
the parked credit away, so a variant could reset its clock by draining and
coming back, which is the starvation the arrival rule exists to stop.

The discovery batch registered at the floor and left the parked tag sitting in
ckVtimeIdle. The immediate serve was worse: it took the floor as its tag,
charged a quantum, drained, and re-parked the lower value over the credit the
variant had earned, so the loss outlived the rollout instead of ending with it.

Both now reclaim max(floor, parked), gated on the current score being at the
floor so a live advanced tag can't be rewound. The gated batch already did this
and its copy of the rule moves into a shared local, since a third copy of it
was the alternative. The extra read only fires where there was no entry to
begin with, so a serving call is unchanged.

The invariant comment above these two sites said floor registration is safe
here because everything they see is established work that lost its tag. That
answers whether there is anything to stack behind, and says nothing about
whether there is credit to reclaim.

Reported by Devin on #4367.

Author: 1stvamp

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

@@ -5608,8 +5608,18 @@ local function tryServe(ckQueueName, mayRaiseFloor, knownRegistered)
 
           -- Advance this variant's virtual time (weight hook: fixed 1 today)
           local weight = 1
-          local tag = tonumber(redis.call('ZSCORE', ckVtimeKey, ckQueueName) or floor)
-          if tag < floor then tag = floor end
+          local storedTag = redis.call('ZSCORE', ckVtimeKey, ckQueueName)
+          local tag = floor
+          if storedTag then
+            tag = tonumber(storedTag)
+            if tag < floor then tag = floor end
+          else
+            -- No entry means pass 2 reached a variant that lost its tag, so this serve is a
+            -- repair. It may still have credit parked from an earlier drain, and serving it
+            -- at the floor would hand that credit out a second time.
+            local parkedTag = redis.call('ZSCORE', ckVtimeIdleKey, ckQueueName)
+            if parkedTag and tonumber(parkedTag) > floor then tag = tonumber(parkedTag) end
+          end
           -- Pass 1 only: it walks in ascending tag order, so anything it has not visited
           -- sits above this. Pass 2 goes by message age, so its tag says nothing about the
           -- entries it skipped and must not move the floor over them.
@@ -5726,7 +5736,25 @@ local ckQueues = redis.call('ZRANGEBYSCORE', ckIndexKey, '-inf', tostring(curren
 -- because everything these two see is established work that lost its tag, which a tenant
 -- cannot mint into while the flag is on. A new enqueue-side path has to stack instead, or
 -- fresh keys start entering at the floor again.
+-- Only a member this call put at the floor may take its parked tag back: a batched ZADD
+-- reports how many it added and not which, so a live advanced tag would otherwise be
+-- rewound. Both of pass 2's registration routes need it, since a variant that drained
+-- under the flag and was re-enqueued by a flag-off instance arrives here with credit
+-- parked and no entry.
+local function restoreParkedTags(names)
+  for _, ckQueueName in ipairs(names) do
+    local cur = redis.call('ZSCORE', ckVtimeKey, ckQueueName)
+    if cur and tonumber(cur) <= floor then
+      local parked = redis.call('ZSCORE', ckVtimeIdleKey, ckQueueName)
+      if parked and tonumber(parked) > floor then
+        redis.call('ZADD', ckVtimeKey, 'XX', parked, ckQueueName)
+      end
+    end
+  end
+end
+
 local discovered = nil
+local discoveredNames = nil
 for _, ckQueueName in ipairs(ckQueues) do
   if not attempted[ckQueueName] then
     if dequeuedCount < actualMaxCount then
@@ -5736,14 +5764,17 @@ for _, ckQueueName in ipairs(ckQueues) do
       -- call however many variants it registers. Skipping attempted matters:
       -- tryServe GCs a drained variant out of both indexes, and re-adding it here
       -- would resurrect a ckVtime entry with no ckIndex member.
-      if discovered == nil then discovered = {ckVtimeKey, 'NX'} end
+      if discovered == nil then discovered = {ckVtimeKey, 'NX'}; discoveredNames = {} end
       table.insert(discovered, tostring(floor))
       table.insert(discovered, ckQueueName)
+      table.insert(discoveredNames, ckQueueName)
     end
   end
 end
 if discovered ~= nil then
-  redis.call('ZADD', unpack(discovered))
+  if redis.call('ZADD', unpack(discovered)) > 0 then
+    restoreParkedTags(discoveredNames)
+  end
 end
 
 -- One variadic ZADD NX settles the whole batch: it registers the genuinely unregistered
@@ -5760,18 +5791,7 @@ if gatedPending ~= nil then
     table.insert(gatedArgs, ckQueueName)
   end
   if redis.call('ZADD', unpack(gatedArgs)) > 0 then
-    for _, ckQueueName in ipairs(gatedPending) do
-      -- The batch ZADD reports how many it added, not which, and gatedPending can hold an
-      -- already-registered variant when the pass-1 scan was truncated. At the floor is the
-      -- discriminator: anything above it has spent credit that must not be rewound.
-      local gateCur = redis.call('ZSCORE', ckVtimeKey, ckQueueName)
-      if gateCur and tonumber(gateCur) <= floor then
-        local gateIdle = redis.call('ZSCORE', ckVtimeIdleKey, ckQueueName)
-        if gateIdle and tonumber(gateIdle) > floor then
-          redis.call('ZADD', ckVtimeKey, 'XX', gateIdle, ckQueueName)
-        end
-      end
-    end
+    restoreParkedTags(gatedPending)
     redis.call('EXPIRE', ckVtimeKey, stateTtl)
     -- This is the one write path that touches ckVtime without going through the
     -- floor-persist block below, which only runs when the call served something. Left

internal-packages/run-engine/src/run-queue/tests/ckVtimeMixedRollout.test.ts

@@ -0,0 +1,157 @@
+import { redisTest } from "@internal/testcontainers";
+import { trace } from "@internal/tracing";
+import { Logger } from "@trigger.dev/core/logger";
+import { Decimal } from "@trigger.dev/database";
+import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js";
+import { RunQueue } from "../index.js";
+import { RunQueueFullKeyProducer } from "../keyProducer.js";
+import type { InputPayload } from "../types.js";
+
+// Pass 2 repairs a variant that reached ckIndex with no ckVtime entry, which during a
+// mixed rollout is one a flag-off instance re-enqueued after a flag-on instance parked its
+// tag. Both of its routes used to repair at the floor and drop the parked credit: the
+// discovery batch registered there, and an immediate serve took the floor as its tag,
+// advanced it, and re-parked the lower value over the credit the variant had earned.
+//
+// The floor is held at 0 by a second variant throughout, so a restored tag can't be
+// mistaken for it.
+
+const testOptions = {
+  name: "rq",
+  tracer: trace.getTracer("rq"),
+  workers: 1,
+  defaultEnvConcurrency: 25,
+  logger: new Logger("RunQueue", "error"),
+  retryOptions: {
+    maxAttempts: 5,
+    factor: 1.1,
+    minTimeoutInMs: 100,
+    maxTimeoutInMs: 1_000,
+    randomize: true,
+  },
+  keys: new RunQueueFullKeyProducer(),
+};
+
+const authenticatedEnvDev = {
+  id: "e1234",
+  type: "DEVELOPMENT" as const,
+  maximumConcurrencyLimit: 10,
+  concurrencyLimitBurstFactor: new Decimal(2.0),
+  project: { id: "p1234" },
+  organization: { id: "o1234" },
+};
+
+const QUEUE = "task/my-task";
+
+function createQueue(
+  redisContainer: any,
+  opts: { vtime?: boolean; maxAttempts?: number } = {}
+): any {
+  const redis = {
+    keyPrefix: "runqueue:test:",
+    host: redisContainer.getHost(),
+    port: redisContainer.getPort(),
+  };
+  return new RunQueue({
+    ...testOptions,
+    masterQueueConsumersDisabled: true,
+    workerOptions: { disabled: true },
+    ...(opts.maxAttempts
+      ? { retryOptions: { ...testOptions.retryOptions, maxAttempts: opts.maxAttempts } }
+      : {}),
+    ...(opts.vtime === false ? {} : { ckVirtualTimeScheduling: { enabled: true } }),
+    queueSelectionStrategy: new FairQueueSelectionStrategy({ redis, keys: testOptions.keys }),
+    redis,
+  } as any) as any;
+}
+
+function makeMessage(overrides: Partial<InputPayload> = {}): InputPayload {
+  return {
+    runId: "r1",
+    taskIdentifier: QUEUE,
+    orgId: "o1234",
+    projectId: "p1234",
+    environmentId: "e1234",
+    environmentType: "DEVELOPMENT",
+    queue: QUEUE,
+    timestamp: Date.now(),
+    attempt: 0,
+    ...overrides,
+  };
+}
+
+const variantName = (ck: string) => testOptions.keys.queueKey(authenticatedEnvDev, QUEUE, ck);
+const shardFor = () => testOptions.keys.masterQueueShardForEnvironment(authenticatedEnvDev.id, 2);
+
+describe("CK vtime: pass-2 repair keeps parked credit", () => {
+  const setup = async (redisContainer: any) => {
+    const on = createQueue(redisContainer);
+    const off = createQueue(redisContainer, { vtime: false });
+    const t0 = Date.now() - 100_000;
+
+    await on.enqueueMessage({
+      env: authenticatedEnvDev,
+      message: makeMessage({ runId: "h0", concurrencyKey: "heavy", timestamp: t0 }),
+      workerQueue: authenticatedEnvDev.id,
+      skipDequeueProcessing: true,
+    });
+    await on.enqueueMessage({
+      env: authenticatedEnvDev,
+      message: makeMessage({ runId: "l0", concurrencyKey: "light", timestamp: t0 + 1 }),
+      workerQueue: authenticatedEnvDev.id,
+      skipDequeueProcessing: true,
+    });
+
+    const heavy = variantName("heavy");
+    const ckVtimeKey = testOptions.keys.ckVtimeKeyFromQueue(heavy);
+    const ckVtimeIdleKey = testOptions.keys.ckVtimeIdleKeyFromQueue(heavy);
+
+    // Serving heavy's only message drains it, so the dequeue parks its tag.
+    await on.testDequeueFromMasterQueue(shardFor(), authenticatedEnvDev.id, 1);
+    expect(await on.redis.zscore(ckVtimeKey, heavy)).toBeNull();
+    // Lifted clear of the floor so a restore is unambiguous.
+    await on.redis.zadd(ckVtimeIdleKey, 5, heavy);
+    await on.acknowledgeMessage(authenticatedEnvDev.organization.id, "h0", {
+      skipDequeueProcessing: true,
+    });
+
+    // The rollout step: an instance with the flag off re-enqueues, so heavy reaches
+    // ckIndex without a ckVtime entry and pass 2 is the only thing that can see it.
+    await off.enqueueMessage({
+      env: authenticatedEnvDev,
+      message: makeMessage({ runId: "h1", concurrencyKey: "heavy", timestamp: t0 + 2 }),
+      workerQueue: authenticatedEnvDev.id,
+      skipDequeueProcessing: true,
+    });
+    expect(await on.redis.zscore(ckVtimeKey, heavy)).toBeNull();
+
+    return { on, off, heavy, ckVtimeKey, ckVtimeIdleKey };
+  };
+
+  redisTest("the discovery batch registers it at its parked tag", async ({ redisContainer }) => {
+    const { on, off, heavy, ckVtimeKey } = await setup(redisContainer);
+    try {
+      // maxCount 1, so pass 1 fills the batch on light and pass 2 only discovers.
+      await on.testDequeueFromMasterQueue(shardFor(), authenticatedEnvDev.id, 1);
+      expect(await on.redis.zscore(ckVtimeKey, heavy)).toBe("5");
+    } finally {
+      await off.quit();
+      await on.quit();
+    }
+  });
+
+  redisTest("an immediate serve charges from its parked tag", async ({ redisContainer }) => {
+    const { on, off, heavy, ckVtimeKey, ckVtimeIdleKey } = await setup(redisContainer);
+    try {
+      // maxCount 2 leaves pass 2 a slot, so heavy is served rather than just registered.
+      await on.testDequeueFromMasterQueue(shardFor(), authenticatedEnvDev.id, 2);
+      // Served from 5 and charged one quantum, then drained again and re-parked at 6.
+      // Taking the floor as its tag would have re-parked 1 over the credit it earned.
+      expect(await on.redis.zscore(ckVtimeKey, heavy)).toBeNull();
+      expect(await on.redis.zscore(ckVtimeIdleKey, heavy)).toBe("6");
+    } finally {
+      await off.quit();
+      await on.quit();
+    }
+  });
+});