diff --git a/internal-packages/run-engine/src/engine/systems/debounceSystem.ts b/internal-packages/run-engine/src/engine/systems/debounceSystem.ts index f749c0768..a43a6b664 100644 --- a/internal-packages/run-engine/src/engine/systems/debounceSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/debounceSystem.ts @@ -121,6 +121,22 @@ end return { 0, value } `, }); + + // Atomically sets runId only if current value equals expected pending claim. + // This prevents the TOCTOU race condition where between GET (check claim) and SET (register), + // another server could claim and register a different run, which would get overwritten. + // Returns 1 if set succeeded, 0 if claim mismatch (lost the claim). + this.redis.defineCommand("registerIfClaimOwned", { + numberOfKeys: 1, + lua: ` +local value = redis.call('GET', KEYS[1]) +if value == ARGV[1] then + redis.call('SET', KEYS[1], ARGV[2], 'PX', ARGV[3]) + return 1 +end +return 0 + `, + }); } /** @@ -593,10 +609,24 @@ return { 0, value } async (span) => { const redisKey = this.getDebounceRedisKey(environmentId, taskIdentifier, debounceKey); + // Calculate TTL: delay until + buffer + const ttlMs = Math.max( + delayUntil.getTime() - Date.now() + 60_000, // Add 1 minute buffer + 60_000 + ); + if (claimId) { - // Verify we still own the pending claim before overwriting - const currentValue = await this.redis.get(redisKey); - if (currentValue !== `pending:${claimId}`) { + // Use atomic Lua script to verify claim and set runId in one operation. + // This prevents the TOCTOU race where another server could claim and register + // between our GET check and SET. + const result = await this.redis.registerIfClaimOwned( + redisKey, + `pending:${claimId}`, + runId, + ttlMs.toString() + ); + + if (result === 0) { // We lost the claim - another server took over or it expired this.$.logger.warn("registerDebouncedRun: lost claim, not registering", { runId, @@ -604,21 +634,15 @@ return { 0, value } taskIdentifier, debounceKey, claimId, - currentValue, }); span.setAttribute("claimLost", true); return false; } + } else { + // No claim to verify, just set directly + await this.redis.set(redisKey, runId, "PX", ttlMs); } - // Calculate TTL: delay until + buffer - const ttlMs = Math.max( - delayUntil.getTime() - Date.now() + 60_000, // Add 1 minute buffer - 60_000 - ); - - await this.redis.set(redisKey, runId, "PX", ttlMs); - this.$.logger.debug("registerDebouncedRun: stored debounce key mapping", { runId, environmentId, @@ -751,5 +775,22 @@ declare module "@internal/redis" { key: string, callback?: Callback<[number, string | null]> ): Result<[number, string | null], Context>; + + /** + * Atomically sets runId only if current value equals expected pending claim. + * Prevents TOCTOU race condition between claim verification and registration. + * @param key - The Redis key + * @param expectedClaim - Expected value "pending:{claimId}" + * @param runId - The new value (run ID) to set + * @param ttlMs - TTL in milliseconds + * @returns 1 if set succeeded, 0 if claim mismatch + */ + registerIfClaimOwned( + key: string, + expectedClaim: string, + runId: string, + ttlMs: string, + callback?: Callback + ): Result; } } diff --git a/internal-packages/run-engine/src/engine/tests/debounce.test.ts b/internal-packages/run-engine/src/engine/tests/debounce.test.ts index 3d4e56af7..92b1d4803 100644 --- a/internal-packages/run-engine/src/engine/tests/debounce.test.ts +++ b/internal-packages/run-engine/src/engine/tests/debounce.test.ts @@ -1939,5 +1939,107 @@ describe("RunEngine debounce", () => { } } ); + + containerTest( + "registerDebouncedRun: atomic claim prevents overwrite when claim is lost", + async ({ prisma, redisOptions }) => { + // This test verifies the fix for the TOCTOU race condition in registerDebouncedRun. + // The race occurs when: + // 1. Server A claims debounce key with claimId-A + // 2. Server B claims same key with claimId-B (after A's claim expires) + // 3. Server B registers runId-B successfully + // 4. Server A attempts to register runId-A with stale claimId-A + // Without the fix, step 4 would overwrite runId-B. With the fix, it fails atomically. + + const { createRedisClient } = await import("@internal/redis"); + + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + debounce: { + maxDebounceDurationMs: 60_000, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + + // Create a separate Redis client to simulate "another server" modifying keys directly + const simulatedServerRedis = createRedisClient({ + ...redisOptions, + keyPrefix: `${redisOptions.keyPrefix ?? ""}debounce:`, + }); + + try { + const taskIdentifier = "test-task"; + const debounceKey = "race-test-key"; + const environmentId = authenticatedEnvironment.id; + const delayUntil = new Date(Date.now() + 60_000); + + await setupBackgroundWorker(engine, authenticatedEnvironment, taskIdentifier); + + // Construct the Redis key (same format as DebounceSystem.getDebounceRedisKey) + const redisKey = `${environmentId}:${taskIdentifier}:${debounceKey}`; + + // Step 1: Server A claims the key with claimId-A + const claimIdA = "claim-server-A"; + await simulatedServerRedis.set(redisKey, `pending:${claimIdA}`, "PX", 60_000); + + // Step 2 & 3: Simulate Server B claiming and registering (after A's claim "expires") + // In reality, this simulates the race where B's claim overwrites A's pending claim + const runIdB = "run_server_B"; + await simulatedServerRedis.set(redisKey, runIdB, "PX", 60_000); + + // Verify Server B's registration is in place + const valueAfterB = await simulatedServerRedis.get(redisKey); + expect(valueAfterB).toBe(runIdB); + + // Step 4: Server A attempts to register with its stale claimId-A + // This should FAIL because the key no longer contains "pending:claim-server-A" + const runIdA = "run_server_A"; + const registered = await engine.debounceSystem.registerDebouncedRun({ + runId: runIdA, + environmentId, + taskIdentifier, + debounceKey, + delayUntil, + claimId: claimIdA, // Stale claim ID + }); + + // Step 5: Verify Server A's registration failed + expect(registered).toBe(false); + + // Step 6: Verify Redis still contains runId-B (not overwritten by Server A) + const finalValue = await simulatedServerRedis.get(redisKey); + expect(finalValue).toBe(runIdB); + } finally { + await simulatedServerRedis.quit(); + await engine.quit(); + } + } + ); });