From e1847ee9f1b21566c04cd66532faaade02a222da Mon Sep 17 00:00:00 2001 From: Guilherme Caulada Date: Wed, 23 Sep 2026 11:35:24 -0300 Subject: [PATCH] fix(scale-down): reselect GitHub Apps and attribute cleanup quota --- .../src/scale-runners/scale-down.test.ts | 66 +++++++++++++++++++ .../src/scale-runners/scale-down.ts | 29 +++++--- 2 files changed, 85 insertions(+), 10 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 3583247f8d..50e8d2d18d 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 @@ -6,6 +6,7 @@ import { defaultComputeProvider } from '@aws-github-runner/compute-providers/pro import { controlPlaneProviderRegistry } from '../control-plane-providers'; import * as ghAuth from '../github/auth'; +import * as rateLimit from '../github/rate-limit'; import { githubCache } from './cache'; import { newestFirstStrategy, oldestFirstStrategy, scaleDown } from './scale-down'; import type { RunnerInfo, RunnerType, ScaleDownComputeProvider } from './types'; @@ -18,6 +19,7 @@ vi.mock('../github/auth', () => ({ })); const mockOctokit = { + hook: { after: vi.fn(), error: vi.fn() }, apps: { getOrgInstallation: vi.fn(), getRepoInstallation: vi.fn(), @@ -262,6 +264,70 @@ describe('Scale down runners', () => { mockCreateClient.mockResolvedValue(mockOctokit as unknown as Octokit); }); + it('selects Apps before owner-cache lookup and attributes response quota to the selected App', async () => { + process.env.SCALE_DOWN_CONFIG = '[]'; + const runners = ['first', 'second', 'third'].map((id) => createRunnerTestData(id, 'Org', 60, true, false, false)); + mockGitHubRunners(runners); + mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce(runners).mockResolvedValue([]); + const authentication = { type: 'app' as const, token: 'token', appId: 1 }; + mockedAppAuth + .mockResolvedValueOnce({ ...authentication, appIndex: 0 }) + .mockResolvedValue({ ...authentication, appIndex: 1 }); + vi.mocked(ghAuth.getStoredInstallationId).mockResolvedValue(123); + const metric = vi.spyOn(rateLimit, 'metricGitHubAppRateLimit').mockResolvedValue(); + try { + await scaleDown(); + expect(mockedInstallationAuth).toHaveBeenCalledTimes(2); + expect(mockedInstallationAuth).toHaveBeenCalledWith(123, '', 0); + expect(mockedInstallationAuth).toHaveBeenCalledWith(123, '', 1); + expect(githubCache.clients.has(`0:Org:${runners[0].owner}`)).toBe(true); + expect(githubCache.clients.has(`1:Org:${runners[0].owner}`)).toBe(true); + const headers = { 'x-ratelimit-remaining': '17', 'x-ratelimit-limit': '5000' }; + await mockOctokit.hook.after.mock.calls[1][1]({ headers }); + expect(metric).toHaveBeenCalledWith(headers, 1); + } finally { + metric.mockRestore(); + vi.mocked(ghAuth.getStoredInstallationId).mockResolvedValue(undefined); + } + }); + + it('attributes successful and failed installation lookup and runner requests to the selected App', async () => { + const runner = createRunnerTestData('quota-hooks', 'Org', 60, true, false, false); + mockGitHubRunners([runner]); + mockListRunners.mockResolvedValueOnce([]).mockResolvedValueOnce([runner]).mockResolvedValue([]); + mockedAppAuth.mockResolvedValue({ type: 'app', token: 'token', appId: 1, appIndex: 2 }); + vi.mocked(ghAuth.getStoredInstallationId).mockResolvedValue(undefined); + const metric = vi.spyOn(rateLimit, 'metricGitHubAppRateLimit').mockResolvedValue(); + try { + mockOctokit.apps.getOrgInstallation.mockImplementationOnce(async () => { + expect(mockOctokit.hook.after).toHaveBeenCalledTimes(1); + expect(mockOctokit.hook.error).toHaveBeenCalledTimes(1); + return { data: { id: 123 } }; + }); + await scaleDown(); + expect(mockOctokit.hook.after).toHaveBeenCalledTimes(2); + const headers = { 'x-ratelimit-remaining': '0', 'x-ratelimit-limit': '5000' }; + for (const [, hook] of mockOctokit.hook.after.mock.calls) { + await hook({ headers }); + expect(metric).toHaveBeenLastCalledWith(headers, 2); + } + for (const [, hook] of mockOctokit.hook.error.mock.calls) { + const error = new RequestError('rate limited', 403, { + request: { method: 'GET', url: 'https://api.github.com/test', headers: {} }, + response: { status: 403, url: 'https://api.github.com/test', headers, data: {} }, + }); + await expect(hook(error)).rejects.toBe(error); + expect(metric).toHaveBeenLastCalledWith(headers, 2); + const networkError = new Error('network unavailable'); + metric.mockClear(); + await expect(hook(networkError)).rejects.toBe(networkError); + expect(metric).not.toHaveBeenCalled(); + } + } finally { + metric.mockRestore(); + } + }); + const endpoints = ['https://api.github.com', 'https://github.enterprise.something', 'https://companyname.ghe.com']; describe.each(endpoints)('for %s', (endpoint) => { 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 1e3e838aed..7e60b95db5 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-down.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-down.ts @@ -24,24 +24,33 @@ type OrgRunnerList = Endpoints['GET /orgs/{org}/actions/runners']['response']['d type RepoRunnerList = Endpoints['GET /repos/{owner}/{repo}/actions/runners']['response']['data']['runners']; type RunnerState = OrgRunnerList[number] | RepoRunnerList[number]; -async function getOrCreateOctokit(runner: RunnerInfo): Promise { - const key = runner.owner; - const cachedOctokit = githubCache.clients.get(key); - - if (cachedOctokit) { - logger.debug(`[createGitHubClientForRunner] Cache hit for ${key}`); - return cachedOctokit; - } +function trackAppQuota(client: Octokit, appIdx: number | undefined): void { + client.hook.after('request', async (response) => { + await metricGitHubAppRateLimit(response.headers, appIdx); + }); + client.hook.error('request', async (error) => { + if (error instanceof RequestError && error.response?.headers) { + await metricGitHubAppRateLimit(error.response.headers, appIdx); + } + throw error; + }); +} - logger.debug(`[createGitHubClientForRunner] Cache miss for ${key}`); +async function getOrCreateOctokit(runner: RunnerInfo): Promise { const { ghesApiUrl } = getGitHubEnterpriseApiUrl(); const ghAuthPre = await createGithubAppAuth(undefined, ghesApiUrl); const appIdx = ghAuthPre.appIndex; + // Re-evaluate quota before consulting the cache; an owner can use another + // installation when the previously selected app becomes exhausted. + const key = `${appIdx ?? 0}:${runner.type}:${runner.owner}`; + const cachedOctokit = githubCache.clients.get(key); + if (cachedOctokit) return cachedOctokit; // Use the pre-configured installation ID when available (avoids an API call). let installationId = await getStoredInstallationId(appIdx); if (installationId === undefined) { const githubClientPre = await createOctokitClient(ghAuthPre.token, ghesApiUrl, appIdx); + trackAppQuota(githubClientPre, appIdx); installationId = runner.type === 'Org' ? ( @@ -58,6 +67,7 @@ async function getOrCreateOctokit(runner: RunnerInfo): Promise { } const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx); const octokit = await createOctokitClient(ghAuth.token, ghesApiUrl, appIdx); + trackAppQuota(octokit, appIdx); githubCache.clients.set(key, octokit); return octokit; @@ -80,7 +90,6 @@ async function getGitHubSelfHostedRunnerState( owner: runner.owner.split('/')[0], repo: runner.owner.split('/')[1], }); - metricGitHubAppRateLimit(state.headers); return state.data; } catch (error) {