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