From f4a65991d9cf44af13f9eda6df05f76bb3b45e65 Mon Sep 17 00:00:00 2001 From: Nelson Wolf Date: Mon, 6 Jul 2026 12:17:44 -0700 Subject: [PATCH] fix(lambda): add post-deregister busy check to prevent terminating active runners The scale-down lambda had a TOCTOU race condition where a job could be assigned to a runner between checking its busy state and terminating the EC2 instance. This caused in-flight jobs to be killed mid-execution. The fix adds a post-deregistration busy re-check: 1. Check busy (fast-path to skip busy runners) 2. Deregister from GitHub (prevents new job assignment) 3. Re-check busy (now stable since no new jobs can be assigned) If the runner became busy between step 1 and step 2, the in-flight job completes using its job-scoped OAuth token and the instance is left for orphan cleanup. Rebased on top of the per-runner de-registration error handling added to main since this fix was originally proposed. Fixes #5085 Co-authored-by: Jack Venberg --- .../src/scale-runners/scale-down.test.ts | 55 +++++++++++++++++ .../src/scale-runners/scale-down.ts | 61 +++++++++++++------ 2 files changed, 98 insertions(+), 18 deletions(-) diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-down.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-down.test.ts index 730cf29bf7..286e662ac3 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-down.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-down.test.ts @@ -360,6 +360,61 @@ describe('Scale down runners', () => { checkNonTerminated(runners); }); + it(`Should not terminate a runner that became busy between deregister and post-deregister check.`, async () => { + // setup: runner appears idle on the pre-deregister check, deregister succeeds, + // but the post-deregister re-check finds it busy (a job was assigned in between) + const runners = [ + createRunnerTestData('race-condition-1', type, MINIMUM_TIME_RUNNING_IN_MINUTES + 1, true, false, false), + ]; + + mockGitHubRunners(runners); + mockAwsRunners(runners); + + // First call (pre-deregister) returns not-busy, second call (post-deregister) returns busy + const busyCheckMock = + mockOctokit.actions[type === 'Repo' ? 'getSelfHostedRunnerForRepo' : 'getSelfHostedRunnerForOrg']; + busyCheckMock + .mockImplementationOnce(() => ({ data: { busy: false } })) + .mockImplementationOnce(() => ({ data: { busy: true } })); + + // act + await expect(scaleDown()).resolves.not.toThrow(); + + // assert: runner should NOT be terminated + checkTerminated(runners); + checkNonTerminated(runners); + }); + + it(`Should terminate a runner when the post-deregister busy check returns 404.`, async () => { + // setup: after deregistration, GitHub API returns 404 (runner fully removed from GitHub) + const runners = [ + createRunnerTestData('deregistered-404', type, MINIMUM_TIME_RUNNING_IN_MINUTES + 1, true, false, true), + ]; + + mockGitHubRunners(runners); + mockAwsRunners(runners); + + const error404 = new RequestError('Runner not found', 404, { + request: { + method: 'GET', + url: 'https://api.github.com/test', + headers: {}, + }, + }); + + // First call (pre-deregister) returns not-busy, second call (post-deregister) throws 404 + const busyCheckMock = + mockOctokit.actions[type === 'Repo' ? 'getSelfHostedRunnerForRepo' : 'getSelfHostedRunnerForOrg']; + busyCheckMock.mockImplementationOnce(() => ({ data: { busy: false } })).mockRejectedValueOnce(error404); + + // act + await expect(scaleDown()).resolves.not.toThrow(); + + // assert: runner should be terminated (404 post-deregister is treated as not busy) + checkTerminated(runners); + checkNonTerminated(runners); + }); + it(`Should terminate orphan (Non JIT)`, async () => { // setup const orphanRunner = createRunnerTestData('orphan-1', type, MINIMUM_BOOT_TIME + 1, false, false, false); diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-down.ts b/lambdas/functions/control-plane/src/scale-runners/scale-down.ts index 5fecc14c99..88ae8dccd9 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-down.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-down.ts @@ -172,6 +172,7 @@ async function removeRunner(ec2runner: RunnerInfo, ghRunnerIds: number[]): Promi return; } + // Step 1: Check busy state as a fast-path to skip runners that are obviously busy. const states = await Promise.all( ghRunnerIds.map(async (ghRunnerId) => { // Get busy state instead of using the output of listGitHubRunners(...) to minimize to race condition. @@ -179,28 +180,52 @@ async function removeRunner(ec2runner: RunnerInfo, ghRunnerIds: number[]): Promi }), ); - if (states.every((busy) => busy === false)) { - const results = await Promise.all( - ghRunnerIds.map((ghRunnerId) => deleteGitHubRunner(githubInstallationClient, ec2runner, ghRunnerId)), + if (!states.every((busy) => busy === false)) { + logger.info(`Runner '${ec2runner.instanceId}' cannot be de-registered, because it is still busy.`); + return; + } + + // Step 2: De-register the runner from GitHub. This prevents GitHub from assigning new jobs + // to this runner, closing the race window where a job could be assigned between the busy + // check above and the termination below. + const results = await Promise.all( + ghRunnerIds.map((ghRunnerId) => deleteGitHubRunner(githubInstallationClient, ec2runner, ghRunnerId)), + ); + + const allSucceeded = results.every((r) => r.success); + const failedRunners = results.filter((r) => !r.success); + + if (!allSucceeded) { + // Only terminate EC2 if we successfully de-registered from GitHub + // Otherwise, leave the instance running so the next scale-down cycle can retry + logger.error( + `Failed to de-register ${failedRunners.length} GitHub runner(s) for instance '${ec2runner.instanceId}'. ` + + `Instance will NOT be terminated to allow retry on next scale-down cycle. ` + + `Failed runner IDs: ${failedRunners.map((r) => r.ghRunnerId).join(', ')}`, ); + return; + } - const allSucceeded = results.every((r) => r.success); - const failedRunners = results.filter((r) => !r.success); + // Step 3: Re-check busy state after de-registration. A job may have been assigned between + // step 1 and step 2. After de-registration no new jobs can be assigned, so this check is + // now stable. If the runner is busy, the in-flight job will complete using its job-scoped + // OAuth token (the runner worker uses credentials from the job message, not the runner + // registration). We leave the instance running and it will be cleaned up as an orphan. + const postDeregisterStates = await Promise.all( + ghRunnerIds.map(async (ghRunnerId) => { + return await getGitHubRunnerBusyState(githubInstallationClient, ec2runner, ghRunnerId); + }), + ); - if (allSucceeded) { - await terminateRunner(ec2runner.instanceId); - logger.info(`AWS runner instance '${ec2runner.instanceId}' is terminated and GitHub runner is de-registered.`); - } else { - // Only terminate EC2 if we successfully de-registered from GitHub - // Otherwise, leave the instance running so the next scale-down cycle can retry - logger.error( - `Failed to de-register ${failedRunners.length} GitHub runner(s) for instance '${ec2runner.instanceId}'. ` + - `Instance will NOT be terminated to allow retry on next scale-down cycle. ` + - `Failed runner IDs: ${failedRunners.map((r) => r.ghRunnerId).join(', ')}`, - ); - } + if (postDeregisterStates.every((busy) => busy === false)) { + await terminateRunner(ec2runner.instanceId); + logger.info(`AWS runner instance '${ec2runner.instanceId}' is terminated and GitHub runner is de-registered.`); } else { - logger.info(`Runner '${ec2runner.instanceId}' cannot be de-registered, because it is still busy.`); + logger.warn( + `Runner '${ec2runner.instanceId}' became busy between idle check and de-registration. ` + + `Skipping termination to allow the in-flight job to complete. ` + + `The instance will be cleaned up as an orphan on a subsequent cycle.`, + ); } } catch (e) { logger.error(