From 4727dfd7d486cc78cab560547e66e3a32177f7aa Mon Sep 17 00:00:00 2001 From: Anatoli Tsikhamirau Date: Sat, 12 Sep 2026 13:17:50 +0200 Subject: [PATCH] fix(github-app): fail over to another app on rate limit, attribute retry-path calls correctly Two related defects in GitHub App selection: 1. App selection already picks the app with the most budget up front, but once an app is selected it stays pinned for the whole request: if it hits a rate limit mid-request, the code sleeps out the retry-after window even when another configured app still has budget. Concurrent invocations that independently pick the same "best" app compound this, since the in-memory rate-limit state only updates after a request completes or fails. 2. getOctokit() selected a GitHub App internally but only returned the Octokit client, not which App was chosen. The retry Lambda called isJobQueued() without an app index, so metricGitHubAppRateLimit() always attributed the call to app 0 regardless of which app actually handled it, corrupting the per-app rate-limit metrics that app selection itself relies on for every request going through the retry path. The throttle callbacks now check whether another configured app has headroom when a rate limit hits; if so, they stop retrying immediately instead of sleeping, and getOctokitWithFailover() retries the call against a freshly selected, different app, returning both the client and the app index it used. With 3+ configured apps, the failover loop tracks and excludes every app already tried in the current sequence (not just the most recent one), bounded by a small attempt cap as a safety net, so a second exhausted app doesn't mean giving up before trying a third. job-retry.ts uses this for its job-status check. scale-up.ts's own job-status check gets the same treatment: each message group tracks its own current app/client and fails over independently the first time its check actually hits a rate limit, reusing the failed-over client for the rest of that group's messages and for the registration-token/JIT-config calls that follow. Those registration-token/JIT-config calls themselves don't get retried on a rate limit yet - they go through a plugin interface (ScaleUpComputeProvider) that would need its own design work to make failover-aware without a breaking change; noted as a follow-up. --- docs/rate-limits-and-tuning.md | 2 + .../control-plane/src/github/auth.test.ts | 141 ++++++++++++++++++ .../control-plane/src/github/auth.ts | 70 +++++++-- .../control-plane/src/github/octokit.test.ts | 103 ++++++++++++- .../control-plane/src/github/octokit.ts | 56 ++++++- .../src/scale-runners/job-retry.test.ts | 53 ++++++- .../src/scale-runners/job-retry.ts | 8 +- .../src/scale-runners/scale-up.test.ts | 73 +++++++++ .../src/scale-runners/scale-up.ts | 70 ++++++++- 9 files changed, 546 insertions(+), 30 deletions(-) diff --git a/docs/rate-limits-and-tuning.md b/docs/rate-limits-and-tuning.md index 88c74bf6d5..76dd814c7d 100644 --- a/docs/rate-limits-and-tuning.md +++ b/docs/rate-limits-and-tuning.md @@ -54,6 +54,8 @@ Without a token cache, each runner also costs a `POST /app/installations/{id}/ac Rate limits are per App installation and cannot be raised. To scale beyond one App's budget, configure extra Apps with `additional_github_apps`. The control-plane lambdas select one App per invocation, making the effective limit N × the per-App limit. Selection prefers the App with the most rate-limit budget remaining, based on the `x-ratelimit-remaining` headers observed by the running Lambda container; Apps that hit a secondary rate limit are skipped for 60 seconds. +If the selected App still gets rate-limited mid-invocation (its budget was exhausted since the last time headers were observed, or a burst raced past that check), the job-status check (`isJobQueued`, in both the scale-up and retry lambdas) fails over immediately to another configured App with headroom instead of sleeping out the `retry-after` window. Failover keeps trying further Apps (bounded) if a second one also turns out exhausted, so it isn't limited to two-App deployments. The registration-token and JIT-config API calls (the ones that create the actual runner) don't yet get this reactive failover — they still rely on the proactive selection above and, if rate-limited, sleep out the window like before. + > [!IMPORTANT] > Every additional App must be installed on the same organizations or repositories as the primary App. The module cannot verify this. A missing installation surfaces at runtime as installation lookup 404s on the fraction of invocations that select the misconfigured App, which is hard to trace back to the installation. diff --git a/lambdas/functions/control-plane/src/github/auth.test.ts b/lambdas/functions/control-plane/src/github/auth.test.ts index 3010e18abb..d1034b780d 100644 --- a/lambdas/functions/control-plane/src/github/auth.test.ts +++ b/lambdas/functions/control-plane/src/github/auth.test.ts @@ -10,6 +10,8 @@ import { createGithubAppAuth, createOctokitClient, getStoredInstallationId, + hasAlternativeAppWithHeadroom, + isGitHubRateLimitError, onRateLimit, onSecondaryRateLimit, reportAppRateLimit, @@ -533,4 +535,143 @@ describe('Test rate-limit aware app selection', () => { const result = await createGithubAppAuth(undefined, '', 1); expect(result.appIndex).toBe(1); }); + + it('excludes the given app index from selection, e.g. one that just got rate-limited', async () => { + reportAppRateLimit(0, 5000); + reportAppRateLimit(1, 100); + + const result = await createGithubAppAuth(undefined, '', undefined, undefined, 0); + expect(result.appIndex).toBe(1); + }); + + it('falls back to the excluded app when it is the only one configured', async () => { + delete process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME; + + const result = await createGithubAppAuth(undefined, '', undefined, undefined, 0); + expect(result.appIndex).toBe(0); + }); +}); + +describe('Test app selection with 3+ apps (multi-app failover)', () => { + const decryptedValue = 'decryptedValue'; + const b64 = Buffer.from(decryptedValue, 'binary').toString('base64'); + const app2IdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_id`; + const app2KeyParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_key_base64`; + const app3IdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_1_id`; + const app3KeyParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_1_key_base64`; + + beforeEach(() => { + const mockedAuth = vi.fn().mockResolvedValue({ token: 'token' }); + vi.mocked(createAppAuth).mockReturnValue(Object.assign(mockedAuth, { hook: vi.fn() })); + + process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME = `/actions-runner/${ENVIRONMENT}/additional_github_apps_manifest`; + mockedGetParameter.mockResolvedValue( + JSON.stringify([ + { idParamName: app2IdParam, keyParamName: app2KeyParam }, + { idParamName: app3IdParam, keyParamName: app3KeyParam }, + ]), + ); + mockedGetParameters.mockResolvedValue( + new Map([ + [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], + [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], + [app2IdParam, '2'], + [app2KeyParam, b64], + [app3IdParam, '3'], + [app3KeyParam, b64], + ]), + ); + vi.spyOn(Math, 'random').mockReturnValue(0); + }); + + it('excludes every already-tried app, not just the most recent one, when given an array', async () => { + reportAppRateLimit(0, 5000); + reportAppRateLimit(1, 4000); + reportAppRateLimit(2, 3000); + + const result = await createGithubAppAuth(undefined, '', undefined, undefined, [0, 1]); + expect(result.appIndex).toBe(2); + }); + + it('hasAlternativeAppWithHeadroom finds the third app once the first two are excluded', async () => { + await createGithubAppAuth(undefined); // populate the credentials cache + reportAppRateLimit(0, 0); + reportAppRateLimit(1, 0); + reportAppRateLimit(2, 100); + + expect(hasAlternativeAppWithHeadroom(0)).toBe(true); // apps 1/2 still uninspected in this call + expect(hasAlternativeAppWithHeadroom([0, 1])).toBe(true); // app 2 still has budget + expect(hasAlternativeAppWithHeadroom([0, 1, 2])).toBe(false); // nothing left to try + }); +}); + +describe('Test hasAlternativeAppWithHeadroom', () => { + const decryptedValue = 'decryptedValue'; + const b64 = Buffer.from(decryptedValue, 'binary').toString('base64'); + const app2IdParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_id`; + const app2KeyParam = `/actions-runner/${ENVIRONMENT}/additional_github_app_0_key_base64`; + + beforeEach(async () => { + const mockedAuth = vi.fn().mockResolvedValue({ token: 'token' }); + vi.mocked(createAppAuth).mockReturnValue(Object.assign(mockedAuth, { hook: vi.fn() })); + + process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME = `/actions-runner/${ENVIRONMENT}/additional_github_apps_manifest`; + mockedGetParameter.mockResolvedValue(JSON.stringify([{ idParamName: app2IdParam, keyParamName: app2KeyParam }])); + mockedGetParameters.mockResolvedValue( + new Map([ + [PARAMETER_GITHUB_APP_ID_NAME, GITHUB_APP_ID], + [PARAMETER_GITHUB_APP_KEY_BASE64_NAME, b64], + [app2IdParam, '2'], + [app2KeyParam, b64], + ]), + ); + // Populates the credentials cache that the sync check reads. + await createGithubAppAuth(undefined); + }); + + it('returns false before any credentials have been loaded', () => { + resetAppCredentialsCache(); + expect(hasAlternativeAppWithHeadroom(0)).toBe(false); + }); + + it('returns true when another app has remaining budget', () => { + reportAppRateLimit(1, 100); + expect(hasAlternativeAppWithHeadroom(0)).toBe(true); + }); + + it('returns false when the only other app is exhausted', () => { + reportAppRateLimit(1, 0); + expect(hasAlternativeAppWithHeadroom(0)).toBe(false); + }); + + it('returns false when the only other app is cooling down from a secondary rate limit', () => { + reportAppRateLimit(1, 100); + reportAppSecondaryRateLimit(1); + expect(hasAlternativeAppWithHeadroom(0)).toBe(false); + }); + + it('returns false in a single-app deployment', async () => { + resetAppCredentialsCache(); + delete process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME; + await createGithubAppAuth(undefined); + + expect(hasAlternativeAppWithHeadroom(0)).toBe(false); + }); +}); + +describe('Test isGitHubRateLimitError', () => { + it.each([ + [ + 'a 403 with x-ratelimit-remaining: 0', + { status: 403, response: { headers: { 'x-ratelimit-remaining': '0' } } }, + true, + ], + ['a 429 with a rate limit message', { status: 429, message: 'You have exceeded a secondary rate limit' }, true], + ['a 403 that is a plain permission error', { status: 403, response: { headers: {} } }, false], + ['a 404', { status: 404 }, false], + ['a non-error value', 'not an error', false], + ['null', null, false], + ])('%s -> %s', (_description, error, expected) => { + expect(isGitHubRateLimitError(error)).toBe(expected); + }); }); diff --git a/lambdas/functions/control-plane/src/github/auth.ts b/lambdas/functions/control-plane/src/github/auth.ts index a0452280c9..85a1eb63a6 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -83,12 +83,21 @@ export function reportAppSecondaryRateLimit(appIndex: number): void { logger.warn(`GitHub App index ${appIndex} put in secondary rate limit cooldown`); } -// Select the app with the most primary rate limit budget remaining, skipping -// apps cooling down after a secondary rate limit. Apps with no observed state -// are assumed full. Iteration starts at a random offset so concurrent -// cold-started lambdas do not all converge on the same app. -async function selectAppIndex(credentialsStore?: GitHubAppCredentialsStore): Promise { +// One or more apps to skip during selection, e.g. apps already tried in a failover sequence. +export type AppIndexExclusion = number | number[]; + +function toExcludeSet(exclude?: AppIndexExclusion): Set { + if (exclude === undefined) return new Set(); + return new Set(Array.isArray(exclude) ? exclude : [exclude]); +} + +// Picks the app with the most rate-limit budget left, skipping excluded/cooling-down apps; random offset avoids concurrent cold starts converging on the same app. +async function selectAppIndex( + credentialsStore?: GitHubAppCredentialsStore, + excludeAppIndexes?: AppIndexExclusion, +): Promise { const credentials = await getAppCredentials(credentialsStore); + const exclude = toExcludeSet(excludeAppIndexes); if (credentials.length === 1) return 0; const now = Date.now(); const offset = Math.floor(Math.random() * credentials.length); @@ -96,6 +105,7 @@ async function selectAppIndex(credentialsStore?: GitHubAppCredentialsStore): Pro let bestRemaining = -1; for (let n = 0; n < credentials.length; n++) { const i = (offset + n) % credentials.length; + if (exclude.has(i)) continue; const state = appRateLimitStates.get(i); if (state && state.cooldownUntil > now) continue; const remaining = state?.remaining ?? Number.MAX_SAFE_INTEGER; @@ -105,8 +115,9 @@ async function selectAppIndex(credentialsStore?: GitHubAppCredentialsStore): Pro } } if (best === -1) { - // Every app is cooling down; pick the one with the most remaining anyway. + // Every non-excluded app is cooling down; pick the one with the most remaining anyway. for (let i = 0; i < credentials.length; i++) { + if (exclude.has(i)) continue; const remaining = appRateLimitStates.get(i)?.remaining ?? Number.MAX_SAFE_INTEGER; if (remaining > bestRemaining) { bestRemaining = remaining; @@ -114,17 +125,35 @@ async function selectAppIndex(credentialsStore?: GitHubAppCredentialsStore): Pro } } } + if (best === -1) best = exclude.size > 0 ? [...exclude][0] : 0; // only excluded apps exist; nothing else to pick // Info so the app selection distribution is observable at default log level. logger.info(`Selected GitHub App index ${best} with ${bestRemaining} rate limit remaining`); return best; } +let cachedAppCredentials: GitHubAppCredential[] | null = null; + async function loadAppCredentials(): Promise { const credentials = await createCommonStorage().githubAppCredentials.get(); logger.info(`Loaded ${credentials.length} GitHub App credential(s)`); + cachedAppCredentials = credentials; return credentials; } +// Sync (for the throttle plugin's callbacks); relies on createGithubAppAuth() having already cached credentials earlier in the same auth flow. +export function hasAlternativeAppWithHeadroom(excludeAppIndexes: AppIndexExclusion): boolean { + if (!cachedAppCredentials || cachedAppCredentials.length <= 1) return false; + const exclude = toExcludeSet(excludeAppIndexes); + const now = Date.now(); + for (let i = 0; i < cachedAppCredentials.length; i++) { + if (exclude.has(i)) continue; + const state = appRateLimitStates.get(i); + if (state && state.cooldownUntil > now) continue; + if ((state?.remaining ?? Number.MAX_SAFE_INTEGER) > 0) return true; + } + return false; +} + function getAppCredentials(credentialsStore?: GitHubAppCredentialsStore): Promise { if (credentialsStore) { return credentialsStore.get(); @@ -139,6 +168,7 @@ export async function getAppCount(credentialsStore?: GitHubAppCredentialsStore): export function resetAppCredentialsCache(): void { appCredentialsPromise = null; + cachedAppCredentials = null; appRateLimitStates.clear(); } @@ -188,8 +218,14 @@ export async function createOctokitClient(token: string, ghesApiUrl = '', appInd retryCount: number, ) => { if (appIndex !== undefined) { - // Primary budget exhausted for this app; steer new flows elsewhere. - reportAppRateLimit(appIndex, 0); + reportAppRateLimit(appIndex, 0); // primary budget exhausted for this app; steer new flows elsewhere + if (hasAlternativeAppWithHeadroom(appIndex)) { + logger.warn( + `GitHub App index ${appIndex} rate-limited with an alternate app available; ` + + `failing over instead of waiting ${retryAfter}s`, + ); + return false; + } } return onRateLimit(retryAfter, options, octokit, retryCount); }, @@ -201,6 +237,13 @@ export async function createOctokitClient(token: string, ghesApiUrl = '', appInd ) => { if (appIndex !== undefined) { reportAppSecondaryRateLimit(appIndex); + if (hasAlternativeAppWithHeadroom(appIndex)) { + logger.warn( + `GitHub App index ${appIndex} secondary rate-limited with an alternate app available; ` + + `failing over instead of waiting ${retryAfter}s`, + ); + return false; + } } return onSecondaryRateLimit(retryAfter, options, octokit, retryCount); }, @@ -213,12 +256,21 @@ export async function createGithubAppAuth( ghesApiUrl = '', appIndex?: number, credentialsStore?: GitHubAppCredentialsStore, + excludeAppIndexes?: AppIndexExclusion, ): Promise { - const idx = appIndex ?? (await selectAppIndex(credentialsStore)); + const idx = appIndex ?? (await selectAppIndex(credentialsStore, excludeAppIndexes)); const auth = await createAuth(installationId, ghesApiUrl, idx, credentialsStore); return { ...(await auth({ type: 'app' })), appIndex: idx }; } +export function isGitHubRateLimitError(error: unknown): boolean { + if (typeof error !== 'object' || error === null) return false; + const err = error as { status?: number; response?: { headers?: Record }; message?: string }; + if (err.status !== 403 && err.status !== 429) return false; + if (err.response?.headers?.['x-ratelimit-remaining'] === '0') return true; + return typeof err.message === 'string' && /rate limit/i.test(err.message); +} + export async function createGithubInstallationAuth( installationId: number | undefined, ghesApiUrl = '', diff --git a/lambdas/functions/control-plane/src/github/octokit.test.ts b/lambdas/functions/control-plane/src/github/octokit.test.ts index 351ce84159..6032887158 100644 --- a/lambdas/functions/control-plane/src/github/octokit.test.ts +++ b/lambdas/functions/control-plane/src/github/octokit.test.ts @@ -1,8 +1,14 @@ import { Octokit } from '@octokit/rest'; import type { ActionRequestMessage } from '../scale-runners/types'; -import { getOctokit } from './octokit'; +import { getOctokit, getOctokitWithFailover } from './octokit'; import { describe, it, expect, beforeEach, vi, Mock } from 'vitest'; -import { createGithubAppAuth, createGithubInstallationAuth, getStoredInstallationId } from '../github/auth'; +import { + createGithubAppAuth, + createGithubInstallationAuth, + getStoredInstallationId, + hasAlternativeAppWithHeadroom, + isGitHubRateLimitError, +} from '../github/auth'; const mockOctokit = { apps: { @@ -19,6 +25,8 @@ vi.mock('../github/auth', async () => ({ createGithubAppAuth: vi.fn().mockResolvedValue({ token: 'token', appIndex: 0 }), getAppCount: vi.fn().mockResolvedValue(1), getStoredInstallationId: vi.fn().mockResolvedValue(undefined), + hasAlternativeAppWithHeadroom: vi.fn().mockReturnValue(false), + isGitHubRateLimitError: vi.fn().mockReturnValue(false), })); vi.mock('@octokit/rest', async () => ({ @@ -190,3 +198,94 @@ describe('Test getOctokit stale installation fallback', () => { expect(createGithubInstallationAuth).toHaveBeenCalledTimes(1); }); }); + +describe('Test getOctokitWithFailover', () => { + const payload = { + eventType: 'workflow_job', + id: 0, + installationId: 5, + repositoryOwner: 'owner', + repositoryName: 'repo', + } as ActionRequestMessage; + + beforeEach(() => { + vi.clearAllMocks(); + (getStoredInstallationId as Mock).mockResolvedValue(undefined); + }); + + it('passes every already-tried app index through to app selection when retrying', async () => { + (createGithubAppAuth as Mock).mockResolvedValue({ token: 'token', appIndex: 0 }); + (isGitHubRateLimitError as Mock).mockReturnValue(true); + (hasAlternativeAppWithHeadroom as Mock).mockReturnValue(true); + + const rateLimitError = Object.assign(new Error('rate limit exceeded'), { status: 403 }); + const work = vi.fn().mockRejectedValueOnce(rateLimitError).mockResolvedValueOnce('done'); + + await expect(getOctokitWithFailover('', true, payload, work)).resolves.toBe('done'); + + expect(createGithubAppAuth).toHaveBeenNthCalledWith(1, undefined, '', undefined, undefined, []); + expect(createGithubAppAuth).toHaveBeenNthCalledWith(2, undefined, '', undefined, undefined, [0]); + expect(work).toHaveBeenCalledTimes(2); + }); + + it('keeps failing over past a second exhausted app with 3+ configured apps', async () => { + (createGithubAppAuth as Mock) + .mockResolvedValueOnce({ token: 'token', appIndex: 0 }) + .mockResolvedValueOnce({ token: 'token', appIndex: 1 }) + .mockResolvedValueOnce({ token: 'token', appIndex: 2 }); + (isGitHubRateLimitError as Mock).mockReturnValue(true); + (hasAlternativeAppWithHeadroom as Mock).mockReturnValue(true); + + const rateLimitError = Object.assign(new Error('rate limit exceeded'), { status: 403 }); + const work = vi + .fn() + .mockRejectedValueOnce(rateLimitError) // app 0 exhausted + .mockRejectedValueOnce(rateLimitError) // app 1 also exhausted + .mockResolvedValueOnce('done'); // app 2 has headroom + + await expect(getOctokitWithFailover('', true, payload, work)).resolves.toBe('done'); + + expect(createGithubAppAuth).toHaveBeenNthCalledWith(1, undefined, '', undefined, undefined, []); + expect(createGithubAppAuth).toHaveBeenNthCalledWith(2, undefined, '', undefined, undefined, [0]); + expect(createGithubAppAuth).toHaveBeenNthCalledWith(3, undefined, '', undefined, undefined, [0, 1]); + expect(work).toHaveBeenCalledTimes(3); + }); + + it('gives up after a bounded number of failover attempts instead of looping forever', async () => { + let nextAppIndex = 0; + (createGithubAppAuth as Mock).mockImplementation(async () => ({ token: 'token', appIndex: nextAppIndex++ })); + (isGitHubRateLimitError as Mock).mockReturnValue(true); + (hasAlternativeAppWithHeadroom as Mock).mockReturnValue(true); // pretend there's always another app + + const rateLimitError = Object.assign(new Error('rate limit exceeded'), { status: 403 }); + const work = vi.fn().mockRejectedValue(rateLimitError); + + await expect(getOctokitWithFailover('', true, payload, work)).rejects.toThrow('rate limit exceeded'); + expect(work.mock.calls.length).toBeLessThan(20); + }); + + it('does not retry when the error is not a rate limit error', async () => { + (createGithubAppAuth as Mock).mockResolvedValue({ token: 'token', appIndex: 0 }); + (isGitHubRateLimitError as Mock).mockReturnValue(false); + + const otherError = new Error('boom'); + const work = vi.fn().mockRejectedValueOnce(otherError); + + await expect(getOctokitWithFailover('', true, payload, work)).rejects.toThrow('boom'); + expect(work).toHaveBeenCalledTimes(1); + expect(createGithubAppAuth).toHaveBeenCalledTimes(1); + }); + + it('does not retry when no alternate app has headroom', async () => { + (createGithubAppAuth as Mock).mockResolvedValue({ token: 'token', appIndex: 0 }); + (isGitHubRateLimitError as Mock).mockReturnValue(true); + (hasAlternativeAppWithHeadroom as Mock).mockReturnValue(false); + + const rateLimitError = Object.assign(new Error('rate limit exceeded'), { status: 403 }); + const work = vi.fn().mockRejectedValueOnce(rateLimitError); + + await expect(getOctokitWithFailover('', true, payload, work)).rejects.toThrow('rate limit exceeded'); + expect(work).toHaveBeenCalledTimes(1); + expect(createGithubAppAuth).toHaveBeenCalledTimes(1); + }); +}); diff --git a/lambdas/functions/control-plane/src/github/octokit.ts b/lambdas/functions/control-plane/src/github/octokit.ts index 46b292686c..7e1f9087b8 100644 --- a/lambdas/functions/control-plane/src/github/octokit.ts +++ b/lambdas/functions/control-plane/src/github/octokit.ts @@ -6,8 +6,13 @@ import { createGithubInstallationAuth, createOctokitClient, getStoredInstallationId, + hasAlternativeAppWithHeadroom, + isGitHubRateLimitError, + type AppIndexExclusion, } from './auth'; +const MAX_FAILOVER_ATTEMPTS = 5; + const logger = createChildLogger('octokit'); function getErrorStatus(error: unknown): number | undefined { @@ -78,15 +83,18 @@ export async function getInstallationId( * phase out the usages of methods in gh-auth.ts outside of this module. Main purpose to make * mocking of the octokit client easier. * - * @returns ockokit client + * @param excludeAppIndexes skip these apps during selection (used to fail over away from apps that + * were already tried and found rate-limited) + * @returns octokit client and the index of the GitHub App used to authenticate it */ export async function getOctokit( ghesApiUrl: string, enableOrgLevel: boolean, payload: ActionRequestMessage, -): Promise { + excludeAppIndexes?: AppIndexExclusion, +): Promise<{ client: Octokit; appIndex: number }> { // Select one app for this entire auth flow - const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl); + const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl, undefined, undefined, excludeAppIndexes); const appIdx = ghAuth.appIndex; const githubAppClient = await createOctokitClient(ghAuth.token, ghesApiUrl, appIdx); @@ -94,7 +102,8 @@ export async function getOctokit( try { const installationAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx); - return await createOctokitClient(installationAuth.token, ghesApiUrl, appIdx); + const client = await createOctokitClient(installationAuth.token, ghesApiUrl, appIdx); + return { client, appIndex: appIdx }; } catch (error) { // The installation id can be stale when it was reused from the webhook payload or from the // pre-configured per-app value while the app was uninstalled and reinstalled. Re-resolve the @@ -117,6 +126,43 @@ export async function getOctokit( }); const installationAuth = await createGithubInstallationAuth(resolvedInstallationId, ghesApiUrl, appIdx); - return await createOctokitClient(installationAuth.token, ghesApiUrl, appIdx); + const client = await createOctokitClient(installationAuth.token, ghesApiUrl, appIdx); + return { client, appIndex: appIdx }; + } +} + +/** + * Runs `work` against the app selected for this auth flow. If it fails with a rate-limit error + * and another configured app still has budget, fails over to that app and retries, instead of + * waiting out the exhausted app's retry-after window. Every app tried so far is excluded from + * later selections, so with 3+ configured apps a failover that lands on an app that turns out to + * also be exhausted keeps failing over rather than giving up after a single retry. + */ +export async function getOctokitWithFailover( + ghesApiUrl: string, + enableOrgLevel: boolean, + payload: ActionRequestMessage, + work: (client: Octokit, appIndex: number) => Promise, +): Promise { + const triedAppIndexes: number[] = []; + for (let attempt = 0; ; attempt++) { + // Pass a snapshot: getOctokit call args must not change out from under a caller inspecting them. + const { client, appIndex } = await getOctokit(ghesApiUrl, enableOrgLevel, payload, [...triedAppIndexes]); + triedAppIndexes.push(appIndex); + try { + return await work(client, appIndex); + } catch (error) { + if ( + attempt >= MAX_FAILOVER_ATTEMPTS || + !isGitHubRateLimitError(error) || + !hasAlternativeAppWithHeadroom(triedAppIndexes) + ) { + throw error; + } + logger.warn('Rate limit hit on selected app, failing over to an alternate app', { + appIndex, + triedAppIndexes: [...triedAppIndexes], + }); + } } } diff --git a/lambdas/functions/control-plane/src/scale-runners/job-retry.test.ts b/lambdas/functions/control-plane/src/scale-runners/job-retry.test.ts index c4ff1e5d76..9470412fac 100644 --- a/lambdas/functions/control-plane/src/scale-runners/job-retry.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/job-retry.test.ts @@ -1,13 +1,18 @@ import { publishMessage } from '../aws/sqs'; import { publishRetryMessage, checkAndRetryJob } from './job-retry'; import type { ActionRequestMessage, ActionRequestMessageRetry } from './types'; -import { getOctokit } from '../github/octokit'; +import { getOctokitWithFailover } from '../github/octokit'; +import { metricGitHubAppRateLimit } from '../github/rate-limit'; import { jobRetryCheck } from '../lambda'; import { Octokit } from '@octokit/rest'; import { createSingleMetric } from '@aws-github-runner/aws-powertools-util'; import { describe, it, expect, beforeEach, vi } from 'vitest'; import type { SQSRecord } from 'aws-lambda'; +vi.mock('../github/rate-limit', async () => ({ + metricGitHubAppRateLimit: vi.fn(), +})); + vi.mock('../aws/sqs', async () => ({ publishMessage: vi.fn(), })); @@ -49,11 +54,11 @@ vi.mock('@octokit/rest', async () => ({ }), })); vi.mock('../github/octokit', async () => ({ - getOctokit: vi.fn(), + getOctokitWithFailover: vi.fn(), })); -const mockCreateOctokitClient = vi.mocked(getOctokit); -mockCreateOctokitClient.mockResolvedValue(new Octokit()); +const mockGetOctokitWithFailover = vi.mocked(getOctokitWithFailover); +mockGetOctokitWithFailover.mockImplementation((_ghesApiUrl, _enableOrgLevel, _payload, work) => work(new Octokit(), 0)); describe('Test job retry publish message', () => { const data = [ @@ -245,6 +250,40 @@ describe(`Test job retry check`, () => { expect(publishMessage).not.toHaveBeenCalled(); }); + it(`should attribute the job-status check to the app getOctokitWithFailover actually selected.`, async () => { + // setup: failover selected app index 1 for this attempt (e.g. app 0 was rate-limited) + mockGetOctokitWithFailover.mockImplementationOnce((_ghesApiUrl, _enableOrgLevel, _payload, work) => + work(new Octokit(), 1), + ); + mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({ + data: { status: 'queued' }, + headers: { 'x-ratelimit-remaining': '10', 'x-ratelimit-limit': '60' }, + })); + + const message: ActionRequestMessageRetry = { + eventType: 'workflow_job', + id: 0, + installationId: 0, + repositoryName: 'test', + repositoryOwner: 'github-aws-runners', + repoOwnerType: 'Organization', + retryCounter: 0, + }; + process.env.ENABLE_ORGANIZATION_RUNNERS = 'true'; + process.env.RUNNER_NAME_PREFIX = 'test'; + process.env.JOB_QUEUE_SCALE_UP_URL = + 'https://sqs.eu-west-1.amazonaws.com/123456789/webhook_events_workflow_job_queue'; + + // act + await checkAndRetryJob(message); + + // assert + expect(metricGitHubAppRateLimit).toHaveBeenCalledWith( + { 'x-ratelimit-remaining': '10', 'x-ratelimit-limit': '60' }, + 1, + ); + }); + it(`should not publish a message for retry if job is no longer queued.`, async () => { // setup mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({ @@ -340,10 +379,10 @@ describe('Test job retry handler (batch processing)', () => { }); it('should continue processing other records when one fails', async () => { - mockCreateOctokitClient - .mockResolvedValueOnce(new Octokit()) // First record succeeds + mockGetOctokitWithFailover + .mockImplementationOnce((_ghesApiUrl, _enableOrgLevel, _payload, work) => work(new Octokit(), 0)) // First record succeeds .mockRejectedValueOnce(new Error('API error')) // Second record fails - .mockResolvedValueOnce(new Octokit()); // Third record succeeds + .mockImplementationOnce((_ghesApiUrl, _enableOrgLevel, _payload, work) => work(new Octokit(), 0)); // Third record succeeds mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({ data: { diff --git a/lambdas/functions/control-plane/src/scale-runners/job-retry.ts b/lambdas/functions/control-plane/src/scale-runners/job-retry.ts index 8f7d6e2289..99fbd5bade 100644 --- a/lambdas/functions/control-plane/src/scale-runners/job-retry.ts +++ b/lambdas/functions/control-plane/src/scale-runners/job-retry.ts @@ -2,7 +2,7 @@ import { addPersistentContextToChildLogger, createSingleMetric, logger } from '@ import { publishMessage } from '../aws/sqs'; import { getGitHubEnterpriseApiUrl, isJobQueued } from './github-runner'; import type { ActionRequestMessage, ActionRequestMessageRetry } from './types'; -import { getOctokit } from '../github/octokit'; +import { getOctokitWithFailover } from '../github/octokit'; import { MetricUnit } from '@aws-lambda-powertools/metrics'; import yn from 'yn'; @@ -61,10 +61,12 @@ export async function checkAndRetryJob(payload: ActionRequestMessageRetry): Prom logger.info(`Received event`); const { ghesApiUrl } = getGitHubEnterpriseApiUrl(); - const ghClient = await getOctokit(ghesApiUrl, enableOrgLevel, payload); // check job is still queued - if (await isJobQueued(ghClient, payload)) { + const jobQueued = await getOctokitWithFailover(ghesApiUrl, enableOrgLevel, payload, (client, appIndex) => + isJobQueued(client, payload, appIndex), + ); + if (jobQueued) { await publishMessage(JSON.stringify(payload), jobQueueUrl); createMetric(enableMetrics, environment, payload); logger.info(`Job is still queued, message published to build queue and will be handled by scale-up.`, { payload }); diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts index 664bac60fb..973adf3171 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts @@ -85,6 +85,8 @@ vi.mock('../github/auth', async () => ({ createOctokitClient: vi.fn(), getAppCount: vi.fn().mockResolvedValue(1), getStoredInstallationId: vi.fn().mockResolvedValue(undefined), + hasAlternativeAppWithHeadroom: vi.fn().mockReturnValue(false), + isGitHubRateLimitError: vi.fn().mockReturnValue(false), })); vi.mock('@aws-github-runner/aws-ssm-util', async () => { @@ -2341,6 +2343,77 @@ describe('Multi-app round-robin', () => { }); }); +describe('App failover for isJobQueued', () => { + const mockedHasAlternativeAppWithHeadroom = vi.mocked(ghAuth.hasAlternativeAppWithHeadroom); + const mockedIsGitHubRateLimitError = vi.mocked(ghAuth.isGitHubRateLimitError); + const rateLimitError = Object.assign(new Error('rate limit exceeded'), { status: 403 }); + + const groupMessages = (count: number): ActionRequestMessageSQS[] => + Array.from({ length: count }, (_, i) => ({ ...TEST_DATA_SINGLE, messageId: `message-${i}` })); + + it('fails over to a re-selected app and treats the job as queued on retry success', async () => { + mockedIsGitHubRateLimitError.mockReturnValue(true); + mockedHasAlternativeAppWithHeadroom.mockReturnValue(true); + mockOctokit.actions.getJobForWorkflowRun + .mockRejectedValueOnce(rateLimitError) + .mockResolvedValueOnce({ data: { status: 'queued' } }); + + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockOctokit.actions.getJobForWorkflowRun).toHaveBeenCalledTimes(2); + expect(mockedAppAuth).toHaveBeenCalledTimes(2); // invocation-level select + one failover + expect(mockedAppAuth).toHaveBeenNthCalledWith(2, undefined, expect.any(String), undefined, expect.anything(), [0]); + expect(createRunner).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); + }); + + it('reuses the failed-over client for the rest of the group without re-selecting again', async () => { + mockedIsGitHubRateLimitError.mockReturnValue(true); + mockedHasAlternativeAppWithHeadroom.mockReturnValue(true); + mockOctokit.actions.getJobForWorkflowRun + .mockRejectedValueOnce(rateLimitError) // first message triggers failover + .mockResolvedValue({ data: { status: 'queued' } }); // retry and second message both succeed + + await scaleUpModule.scaleUp(groupMessages(2)); + + expect(mockOctokit.actions.getJobForWorkflowRun).toHaveBeenCalledTimes(3); + expect(mockedAppAuth).toHaveBeenCalledTimes(2); // invocation-level select + one failover, not two + expect(createRunner).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 2 })); + }); + + it('does not fail over when the error is not rate-limit shaped (falls open as before)', async () => { + mockedIsGitHubRateLimitError.mockReturnValue(false); + mockedHasAlternativeAppWithHeadroom.mockReturnValue(true); + mockOctokit.actions.getJobForWorkflowRun.mockRejectedValue(new Error('GitHub API 502')); + + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockedAppAuth).toHaveBeenCalledTimes(1); // no failover attempted + expect(createRunner).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); + }); + + it('does not fail over when no alternate app has headroom (falls open as before)', async () => { + mockedIsGitHubRateLimitError.mockReturnValue(true); + mockedHasAlternativeAppWithHeadroom.mockReturnValue(false); + mockOctokit.actions.getJobForWorkflowRun.mockRejectedValue(rateLimitError); + + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockedAppAuth).toHaveBeenCalledTimes(1); // no failover attempted + expect(createRunner).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); + }); + + it('gives up after a bounded number of failover attempts instead of looping forever', async () => { + mockedIsGitHubRateLimitError.mockReturnValue(true); + mockedHasAlternativeAppWithHeadroom.mockReturnValue(true); // pretend there's always another app + mockOctokit.actions.getJobForWorkflowRun.mockRejectedValue(rateLimitError); + + await scaleUpModule.scaleUp(TEST_DATA); + + expect(mockOctokit.actions.getJobForWorkflowRun.mock.calls.length).toBeLessThan(20); + expect(createRunner).toHaveBeenCalledWith(expect.objectContaining({ numberOfRunners: 1 })); // fail-open + }); +}); + function defaultOctokitMockImpl() { mockOctokit.actions.getJobForWorkflowRun.mockImplementation(() => ({ data: { diff --git a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts index d4e3889f19..c5c441e2c5 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -5,7 +5,13 @@ import { createStorageProviders, type StorageProviders } from '@aws-github-runne import { Octokit } from '@octokit/rest'; import yn from 'yn'; -import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from '../github/auth'; +import { + createGithubAppAuth, + createGithubInstallationAuth, + createOctokitClient, + hasAlternativeAppWithHeadroom, + isGitHubRateLimitError, +} from '../github/auth'; import { controlPlaneProviderRegistry } from '../control-plane-providers'; import { getGitHubEnterpriseApiUrl, @@ -25,6 +31,12 @@ import type { } from './types'; const logger = createChildLogger('scale-up'); +const MAX_FAILOVER_ATTEMPTS = 5; + +interface GithubClientHolder { + client: Octokit; + appIndex: number; +} function getErrorStatus(error: unknown): number | undefined { if (typeof error !== 'object' || error === null) { @@ -90,6 +102,55 @@ async function createGithubInstallationClient( } } +// On a rate limit, retries against a different app instead of sleeping; mutates clientHolder so later messages in the group reuse it. +async function checkJobQueuedWithFailover( + clientHolder: GithubClientHolder, + message: ActionRequestMessageSQS, + enableOrgLevel: boolean, + ghesApiUrl: string, + storage: StorageProviders, +): Promise { + const triedAppIndexes: number[] = []; + for (let attempt = 0; ; attempt++) { + try { + return await isJobQueued(clientHolder.client, message, clientHolder.appIndex); + } catch (error) { + if ( + error instanceof UnsupportedEventError || + attempt >= MAX_FAILOVER_ATTEMPTS || + !isGitHubRateLimitError(error) || + !hasAlternativeAppWithHeadroom([...triedAppIndexes, clientHolder.appIndex]) + ) { + throw error; + } + + triedAppIndexes.push(clientHolder.appIndex); + logger.warn('Rate limit hit checking job status, failing over to an alternate app', { + appIndex: clientHolder.appIndex, + triedAppIndexes: [...triedAppIndexes], + }); + + const failoverAuth = await createGithubAppAuth( + undefined, + ghesApiUrl, + undefined, + storage.githubAppCredentials, + triedAppIndexes, + ); + const failoverAppClient = await createOctokitClient(failoverAuth.token, ghesApiUrl, failoverAuth.appIndex); + clientHolder.client = await createGithubInstallationClient( + failoverAppClient, + enableOrgLevel, + message, + ghesApiUrl, + failoverAuth.appIndex, + storage, + ); + clientHolder.appIndex = failoverAuth.appIndex; + } + } +} + export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise { const storage = createStorageProviders(); logger.info('Received scale up requests', { @@ -212,6 +273,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise