From 0b717cce4adb8ad2dce5fe76ee3ff2e72eb97252 Mon Sep 17 00:00:00 2001 From: Guilherme Caulada Date: Tue, 8 Sep 2026 14:50:35 -0300 Subject: [PATCH] feat(github-app): select app by rate limit budget (#5283) Follow-up to #5269, item 2 of the [#5038 review](https://github.com/github-aws-runners/terraform-aws-github-runner/pull/5038#pullrequestreview-4855605507). Stacked on #5282 (manifest transport); review that first. Replaces uniform random GitHub App selection with budget-aware selection, ported from our production fork: - The auth module tracks the last observed `x-ratelimit-remaining` per app, fed by `metricGitHubAppRateLimit` headers and the throttling plugin callbacks. - Selection picks the app with the most budget remaining; apps with no observed state are assumed full. - Apps that hit a secondary rate limit are skipped for 60 seconds; if every app is cooling down, the one with the most budget is used anyway. - Iteration starts at a random offset so concurrent cold-started lambdas do not converge on the same app. Uniform random kept sending ~1/N of traffic to apps that were already exhausted. Selection state is per warm container and converges within a few invocations; a cold container starts as before. - control-plane: 353 tests passed, including six new selection tests (budget preference, unobserved-assumed-full, cooldown skip, all-cooling fallback, single-app short-circuit, explicit appIndex) and a selector-feed test for the rate-limit headers. - ESLint + Prettier clean. Follow-up to #5269 / #5038. Depends on #5282. --- docs/rate-limits-and-tuning.md | 2 +- .../control-plane/src/github/auth.test.ts | 99 +++++++++++++++++++ .../control-plane/src/github/auth.ts | 95 ++++++++++++++++-- .../control-plane/src/github/octokit.ts | 8 +- .../src/github/rate-limit.test.ts | 9 +- .../control-plane/src/github/rate-limit.ts | 8 +- .../functions/control-plane/src/pool/pool.ts | 4 +- .../src/scale-runners/scale-down.ts | 4 +- .../src/scale-runners/scale-up.ts | 6 +- 9 files changed, 215 insertions(+), 20 deletions(-) diff --git a/docs/rate-limits-and-tuning.md b/docs/rate-limits-and-tuning.md index 571fd96ff1..88c74bf6d5 100644 --- a/docs/rate-limits-and-tuning.md +++ b/docs/rate-limits-and-tuning.md @@ -52,7 +52,7 @@ Without a token cache, each runner also costs a `POST /app/installations/{id}/ac ### Distributing load across multiple GitHub Apps -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. +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. > [!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 c2524503b7..3010e18abb 100644 --- a/lambdas/functions/control-plane/src/github/auth.test.ts +++ b/lambdas/functions/control-plane/src/github/auth.test.ts @@ -12,6 +12,8 @@ import { getStoredInstallationId, onRateLimit, onSecondaryRateLimit, + reportAppRateLimit, + reportAppSecondaryRateLimit, resetAppCredentialsCache, } from './auth'; import { describe, it, expect, beforeEach, vi } from 'vitest'; @@ -435,3 +437,100 @@ describe('Test getStoredInstallationId', () => { expect(result1).toBe(67890); }); }); + +describe('Test rate-limit aware app selection', () => { + 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(() => { + const mockedAuth = vi.fn(); + mockedAuth.mockResolvedValue({ token: 'token' }); + const mockWithHook = Object.assign(mockedAuth, { hook: vi.fn() }); + vi.mocked(createAppAuth).mockReturnValue(mockWithHook); + + 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], + ]), + ); + + // Pin the random start offset to 0 so selection is deterministic. + vi.spyOn(Math, 'random').mockReturnValue(0); + }); + + it('selects the app with the most rate limit budget remaining', async () => { + reportAppRateLimit(0, 100); + reportAppRateLimit(1, 5000); + + const result = await createGithubAppAuth(undefined); + expect(result.appIndex).toBe(1); + }); + + it('selects from the supplied credentials store without reading SSM', async () => { + const credentialsStore = { + get: vi.fn().mockResolvedValue([ + { appId: 10, privateKey: 'first-key' }, + { appId: 20, privateKey: 'second-key' }, + ]), + }; + reportAppRateLimit(0, 100); + reportAppRateLimit(1, 5000); + + const result = await createGithubAppAuth(undefined, '', undefined, credentialsStore); + + expect(result.appIndex).toBe(1); + expect(createAppAuth).toHaveBeenCalledWith(expect.objectContaining({ appId: 20, createJwt: expect.any(Function) })); + expect(mockedGetParameter).not.toHaveBeenCalled(); + expect(mockedGetParameters).not.toHaveBeenCalled(); + }); + + it('assumes full budget for apps without observed state', async () => { + reportAppRateLimit(0, 100); + // App 1 has no observed state and is assumed full. + + const result = await createGithubAppAuth(undefined); + expect(result.appIndex).toBe(1); + }); + + it('skips an app cooling down after a secondary rate limit', async () => { + reportAppRateLimit(0, 100); + reportAppRateLimit(1, 5000); + reportAppSecondaryRateLimit(1); + + const result = await createGithubAppAuth(undefined); + expect(result.appIndex).toBe(0); + }); + + it('falls back to the most budget when every app is cooling down', async () => { + reportAppRateLimit(0, 100); + reportAppRateLimit(1, 5000); + reportAppSecondaryRateLimit(0); + reportAppSecondaryRateLimit(1); + + const result = await createGithubAppAuth(undefined); + expect(result.appIndex).toBe(1); + }); + + it('short-circuits to the primary app in single-app deployments', async () => { + delete process.env.PARAMETER_GITHUB_APPS_MANIFEST_NAME; + reportAppRateLimit(0, 0); + + const result = await createGithubAppAuth(undefined); + expect(result.appIndex).toBe(0); + }); + + it('respects an explicitly provided appIndex', async () => { + reportAppRateLimit(0, 5000); + reportAppRateLimit(1, 100); + + const result = await createGithubAppAuth(undefined, '', 1); + expect(result.appIndex).toBe(1); + }); +}); diff --git a/lambdas/functions/control-plane/src/github/auth.ts b/lambdas/functions/control-plane/src/github/auth.ts index 3e177ab253..a0452280c9 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -59,6 +59,66 @@ export function onSecondaryRateLimit( let appCredentialsPromise: Promise | null = null; +interface AppRateLimitState { + remaining: number; + cooldownUntil: number; +} + +// Last known primary rate limit remaining and secondary rate limit cooldown +// per app index. Fed by response headers and throttling callbacks; persists +// across invocations in a warm lambda so selection converges quickly. +const appRateLimitStates = new Map(); +const SECONDARY_RATE_LIMIT_COOLDOWN_MS = 60_000; + +export function reportAppRateLimit(appIndex: number, remaining: number): void { + const state = appRateLimitStates.get(appIndex) ?? { remaining, cooldownUntil: 0 }; + state.remaining = remaining; + appRateLimitStates.set(appIndex, state); +} + +export function reportAppSecondaryRateLimit(appIndex: number): void { + const state = appRateLimitStates.get(appIndex) ?? { remaining: 0, cooldownUntil: 0 }; + state.cooldownUntil = Date.now() + SECONDARY_RATE_LIMIT_COOLDOWN_MS; + appRateLimitStates.set(appIndex, state); + 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 { + const credentials = await getAppCredentials(credentialsStore); + if (credentials.length === 1) return 0; + const now = Date.now(); + const offset = Math.floor(Math.random() * credentials.length); + let best = -1; + let bestRemaining = -1; + for (let n = 0; n < credentials.length; n++) { + const i = (offset + n) % credentials.length; + const state = appRateLimitStates.get(i); + if (state && state.cooldownUntil > now) continue; + const remaining = state?.remaining ?? Number.MAX_SAFE_INTEGER; + if (remaining > bestRemaining) { + bestRemaining = remaining; + best = i; + } + } + if (best === -1) { + // Every app is cooling down; pick the one with the most remaining anyway. + for (let i = 0; i < credentials.length; i++) { + const remaining = appRateLimitStates.get(i)?.remaining ?? Number.MAX_SAFE_INTEGER; + if (remaining > bestRemaining) { + bestRemaining = remaining; + best = i; + } + } + } + // 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; +} + async function loadAppCredentials(): Promise { const credentials = await createCommonStorage().githubAppCredentials.get(); logger.info(`Loaded ${credentials.length} GitHub App credential(s)`); @@ -79,6 +139,7 @@ export async function getAppCount(credentialsStore?: GitHubAppCredentialsStore): export function resetAppCredentialsCache(): void { appCredentialsPromise = null; + appRateLimitStates.clear(); } export async function getStoredInstallationId( @@ -97,7 +158,7 @@ export async function getAppId(appIndex = 0, credentialsStore?: GitHubAppCredent return credential.appId.toString(); } -export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise { +export async function createOctokitClient(token: string, ghesApiUrl = '', appIndex?: number): Promise { const CustomOctokit = Octokit.plugin(retry, throttling); const octokitOptions: OctokitOptions = { auth: token }; if (ghesApiUrl) { @@ -119,7 +180,31 @@ export async function createOctokitClient(token: string, ghesApiUrl = ''): Promi }); }, }, - throttle: { onRateLimit, onSecondaryRateLimit }, + throttle: { + onRateLimit: ( + retryAfter: number, + options: Required, + octokit: CoreOctokit, + retryCount: number, + ) => { + if (appIndex !== undefined) { + // Primary budget exhausted for this app; steer new flows elsewhere. + reportAppRateLimit(appIndex, 0); + } + return onRateLimit(retryAfter, options, octokit, retryCount); + }, + onSecondaryRateLimit: ( + retryAfter: number, + options: Required, + octokit: CoreOctokit, + retryCount: number, + ) => { + if (appIndex !== undefined) { + reportAppSecondaryRateLimit(appIndex); + } + return onSecondaryRateLimit(retryAfter, options, octokit, retryCount); + }, + }, }); } @@ -129,8 +214,7 @@ export async function createGithubAppAuth( appIndex?: number, credentialsStore?: GitHubAppCredentialsStore, ): Promise { - const credentials = await getAppCredentials(credentialsStore); - const idx = appIndex ?? Math.floor(Math.random() * credentials.length); + const idx = appIndex ?? (await selectAppIndex(credentialsStore)); const auth = await createAuth(installationId, ghesApiUrl, idx, credentialsStore); return { ...(await auth({ type: 'app' })), appIndex: idx }; } @@ -141,8 +225,7 @@ export async function createGithubInstallationAuth( appIndex?: number, credentialsStore?: GitHubAppCredentialsStore, ): Promise { - const credentials = await getAppCredentials(credentialsStore); - const idx = appIndex ?? Math.floor(Math.random() * credentials.length); + const idx = appIndex ?? (await selectAppIndex(credentialsStore)); const auth = await createAuth(installationId, ghesApiUrl, idx, credentialsStore); return auth({ type: 'installation', installationId }); } diff --git a/lambdas/functions/control-plane/src/github/octokit.ts b/lambdas/functions/control-plane/src/github/octokit.ts index 010516e436..46b292686c 100644 --- a/lambdas/functions/control-plane/src/github/octokit.ts +++ b/lambdas/functions/control-plane/src/github/octokit.ts @@ -68,7 +68,7 @@ export async function getInstallationId( appIndex?: number, ): Promise { const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl, appIndex); - const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl); + const githubClient = await createOctokitClient(ghAuth.token, ghesApiUrl, appIndex); return resolveInstallationId(githubClient, enableOrgLevel, payload, appIndex); } @@ -88,13 +88,13 @@ export async function getOctokit( // Select one app for this entire auth flow const ghAuth = await createGithubAppAuth(undefined, ghesApiUrl); const appIdx = ghAuth.appIndex; - const githubAppClient = await createOctokitClient(ghAuth.token, ghesApiUrl); + const githubAppClient = await createOctokitClient(ghAuth.token, ghesApiUrl, appIdx); const installationId = await resolveInstallationId(githubAppClient, enableOrgLevel, payload, appIdx); try { const installationAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx); - return await createOctokitClient(installationAuth.token, ghesApiUrl); + return await createOctokitClient(installationAuth.token, ghesApiUrl, 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 +117,6 @@ export async function getOctokit( }); const installationAuth = await createGithubInstallationAuth(resolvedInstallationId, ghesApiUrl, appIdx); - return await createOctokitClient(installationAuth.token, ghesApiUrl); + return await createOctokitClient(installationAuth.token, ghesApiUrl, appIdx); } } diff --git a/lambdas/functions/control-plane/src/github/rate-limit.test.ts b/lambdas/functions/control-plane/src/github/rate-limit.test.ts index 40c83ff6a3..fe54759bfd 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.test.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.test.ts @@ -3,11 +3,12 @@ import { createSingleMetric } from '@aws-github-runner/aws-powertools-util'; import { MetricUnit } from '@aws-lambda-powertools/metrics'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { getAppId } from './auth'; +import { getAppId, reportAppRateLimit } from './auth'; import { metricGitHubAppRateLimit } from './rate-limit'; vi.mock('./auth', () => ({ getAppId: vi.fn(), + reportAppRateLimit: vi.fn(), })); vi.mock('@aws-github-runner/aws-powertools-util', () => ({ logger: { debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn() }, @@ -61,4 +62,10 @@ describe('metricGitHubAppRateLimit', () => { expect(mockedGetAppId).toHaveBeenNthCalledWith(1, 0); expect(mockedGetAppId).toHaveBeenNthCalledWith(2, 1); }); + it('feeds the app selector even when metrics are disabled', async () => { + process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'false'; + await metricGitHubAppRateLimit({ 'x-ratelimit-remaining': '4200', 'x-ratelimit-limit': '5000' }, 1); + expect(reportAppRateLimit).toHaveBeenCalledWith(1, 4200); + expect(createSingleMetric).not.toHaveBeenCalled(); + }); }); diff --git a/lambdas/functions/control-plane/src/github/rate-limit.ts b/lambdas/functions/control-plane/src/github/rate-limit.ts index 710d7cf80a..4cf2ad923f 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.ts @@ -2,7 +2,7 @@ import { ResponseHeaders } from '@octokit/types'; import { createSingleMetric, logger } from '@aws-github-runner/aws-powertools-util'; import { MetricUnit } from '@aws-lambda-powertools/metrics'; import yn from 'yn'; -import { getAppId } from './auth'; +import { getAppId, reportAppRateLimit } from './auth'; export async function metricGitHubAppRateLimit(headers: ResponseHeaders, appIndex?: number): Promise { try { @@ -11,6 +11,12 @@ export async function metricGitHubAppRateLimit(headers: ResponseHeaders, appInde logger.debug(`Rate limit remaining: ${remaining}, limit: ${limit}`); + // Feed the app selector so new auth flows prefer the app with the most + // budget left. Headers without an appIndex belong to the primary app. + if (!isNaN(remaining)) { + reportAppRateLimit(appIndex ?? 0, remaining); + } + const updateMetric = yn(process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT); if (updateMetric) { const appId = await getAppId(appIndex); diff --git a/lambdas/functions/control-plane/src/pool/pool.ts b/lambdas/functions/control-plane/src/pool/pool.ts index c02787426f..9f91c6cbcd 100644 --- a/lambdas/functions/control-plane/src/pool/pool.ts +++ b/lambdas/functions/control-plane/src/pool/pool.ts @@ -51,7 +51,7 @@ export async function adjust(event: PoolEvent): Promise { const installationId = await getInstallationId(ghAppAuth.token, ghesApiUrl, runnerOwner, appIdx, storage); const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx, storage.githubAppCredentials); - const githubInstallationClient = await createOctokitClient(ghAuth.token, ghesApiUrl); + const githubInstallationClient = await createOctokitClient(ghAuth.token, ghesApiUrl, appIdx); // Get statuses of runners registered in GitHub const runnerStatusses = await getGitHubRegisteredRunnnerStatusses( @@ -120,7 +120,7 @@ async function getInstallationId( const storedId = await getStoredInstallationId(appIndex, storage?.githubAppCredentials); if (storedId !== undefined) return storedId; - const githubClient = await createOctokitClient(appToken, ghesApiUrl); + const githubClient = await createOctokitClient(appToken, ghesApiUrl, appIndex); return ( await githubClient.apps.getOrgInstallation({ 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 3e387bce06..329ce694d9 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-down.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-down.ts @@ -41,7 +41,7 @@ async function getOrCreateOctokit(runner: RunnerInfo): Promise { // 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); + const githubClientPre = await createOctokitClient(ghAuthPre.token, ghesApiUrl, appIdx); installationId = runner.type === 'Org' ? ( @@ -57,7 +57,7 @@ async function getOrCreateOctokit(runner: RunnerInfo): Promise { ).data.id; } const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx); - const octokit = await createOctokitClient(ghAuth.token, ghesApiUrl); + const octokit = await createOctokitClient(ghAuth.token, ghesApiUrl, appIdx); githubCache.clients.set(key, octokit); return octokit; 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 6608bacb02..d4e3889f19 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -58,7 +58,7 @@ async function createGithubInstallationClient( appIndex, storage?.githubAppCredentials, ); - return await createOctokitClient(ghAuth.token, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl, appIndex); } 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 @@ -86,7 +86,7 @@ async function createGithubInstallationClient( appIndex, storage?.githubAppCredentials, ); - return await createOctokitClient(ghAuth.token, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl, appIndex); } } @@ -117,7 +117,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise