From fb3555f065fa680657ae9df6d56de92b9057ef64 Mon Sep 17 00:00:00 2001 From: Guilherme Caulada Date: Wed, 19 Aug 2026 15:00:26 -0300 Subject: [PATCH] feat(github-app): select app by rate limit budget Replace uniform random GitHub App selection with budget-aware selection. The auth module tracks the last observed x-ratelimit-remaining per app (fed by response headers and the throttling plugin callbacks) and selects the app with the most budget left. Apps that hit a secondary rate limit are skipped for 60 seconds. Iteration starts at a random offset so concurrent cold-started lambdas do not converge on the same app. Uniform random selection kept sending ~1/N of traffic to apps that were already exhausted or throttled. Selection state is per warm container and converges within a few invocations. --- docs/rate-limits-and-tuning.md | 2 +- .../control-plane/src/github/auth.test.ts | 81 ++++++++++++++++ .../control-plane/src/github/auth.ts | 97 +++++++++++++++++-- .../control-plane/src/github/octokit.ts | 8 +- .../src/github/rate-limit.test.ts | 12 ++- .../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, 199 insertions(+), 23 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..04864be8f2 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,82 @@ 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('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 b40b120bdd..6bf5220604 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -77,6 +77,66 @@ interface GitHubAppCredential { 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(): Promise { + const credentials = await getAppCredentials(); + 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; +} + // One entry per additional app in the manifest parameter. The manifest keeps // the lambda environment size constant regardless of the number of apps: the // environment carries only the manifest's parameter name, and the manifest @@ -158,6 +218,7 @@ export async function getAppCount(): Promise { export function resetAppCredentialsCache(): void { appCredentialsPromise = null; + appRateLimitStates.clear(); } export async function getStoredInstallationId(appIndex: number): Promise { @@ -165,7 +226,7 @@ export async function getStoredInstallationId(appIndex: number): Promise { +export async function createOctokitClient(token: string, ghesApiUrl = '', appIndex?: number): Promise { const CustomOctokit = Octokit.plugin(retry, throttling); const ocktokitOptions: OctokitOptions = { auth: token, @@ -190,8 +251,29 @@ export async function createOctokitClient(token: string, ghesApiUrl = ''): Promi }, }, throttle: { - onRateLimit, - onSecondaryRateLimit, + 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); + }, }, }); } @@ -201,8 +283,7 @@ export async function createGithubAppAuth( ghesApiUrl = '', appIndex?: number, ): Promise { - const credentials = await getAppCredentials(); - const idx = appIndex ?? Math.floor(Math.random() * credentials.length); + const idx = appIndex ?? (await selectAppIndex()); const auth = await createAuth(installationId, ghesApiUrl, idx); const result = await auth({ type: 'app' }); return { ...result, appIndex: idx }; @@ -213,8 +294,7 @@ export async function createGithubInstallationAuth( ghesApiUrl = '', appIndex?: number, ): Promise { - const credentials = await getAppCredentials(); - const idx = appIndex ?? Math.floor(Math.random() * credentials.length); + const idx = appIndex ?? (await selectAppIndex()); const auth = await createAuth(installationId, ghesApiUrl, idx); return auth({ type: 'installation', installationId }); } @@ -233,8 +313,7 @@ async function createAuth( appIndex?: number, ): Promise { const credentials = await getAppCredentials(); - const selected = - appIndex !== undefined ? credentials[appIndex] : credentials[Math.floor(Math.random() * credentials.length)]; + const selected = credentials[appIndex ?? (await selectAppIndex())]; logger.debug(`Selected GitHub App ${selected.appId} for authentication`); 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 457a367c89..c3398d57e7 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 { metricGitHubAppRateLimit } from './rate-limit'; import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { getLoadedAppId } from './auth'; +import { getLoadedAppId, reportAppRateLimit } from './auth'; vi.mock('./auth', async () => ({ // App ids per index, as loaded by the auth module from SSM. getLoadedAppId: vi.fn(async (appIndex: number) => [1234, 5678][appIndex]), + reportAppRateLimit: vi.fn(), })); vi.mock('@aws-github-runner/aws-powertools-util', async () => { @@ -91,6 +92,15 @@ describe('metricGitHubAppRateLimit', () => { }); }); + it('feeds the app selector with the remaining budget even when metrics are disabled', async () => { + process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'false'; + const headers: ResponseHeaders = { 'x-ratelimit-remaining': '4200', 'x-ratelimit-limit': '5000' }; + + await metricGitHubAppRateLimit(headers, 1); + + expect(reportAppRateLimit).toHaveBeenCalledWith(1, 4200); + }); + it('should label metric with an empty AppId when the appIndex is unknown', async () => { process.env.ENABLE_METRIC_GITHUB_APP_RATE_LIMIT = 'true'; const headers: ResponseHeaders = { 'x-ratelimit-remaining': '75', 'x-ratelimit-limit': '5000' }; diff --git a/lambdas/functions/control-plane/src/github/rate-limit.ts b/lambdas/functions/control-plane/src/github/rate-limit.ts index 2cd194fbeb..2860905059 100644 --- a/lambdas/functions/control-plane/src/github/rate-limit.ts +++ b/lambdas/functions/control-plane/src/github/rate-limit.ts @@ -3,7 +3,7 @@ import { createSingleMetric, logger } from '@aws-github-runner/aws-powertools-ut import { MetricUnit } from '@aws-lambda-powertools/metrics'; import yn from 'yn'; -import { getLoadedAppId } from './auth'; +import { getLoadedAppId, reportAppRateLimit } from './auth'; // App ids come from the credentials already loaded by the auth module, so no // additional SSM reads are needed here. Index 0 is the primary app. @@ -19,6 +19,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 da5d2ee9b1..7d64aefc85 100644 --- a/lambdas/functions/control-plane/src/pool/pool.ts +++ b/lambdas/functions/control-plane/src/pool/pool.ts @@ -55,7 +55,7 @@ export async function adjust(event: PoolEvent): Promise { const installationId = await getInstallationId(ghAppAuth.token, ghesApiUrl, runnerOwner, appIdx); const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIdx); - 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(appToken: string, ghesApiUrl: string, org: stri const storedId = await getStoredInstallationId(appIndex); 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 44d522a1f0..06ce69dc7c 100644 --- a/lambdas/functions/control-plane/src/scale-runners/scale-up.ts +++ b/lambdas/functions/control-plane/src/scale-runners/scale-up.ts @@ -44,7 +44,7 @@ async function createGithubInstallationClient( try { const ghAuth = await createGithubInstallationAuth(installationId, ghesApiUrl, appIndex); - 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 @@ -67,7 +67,7 @@ async function createGithubInstallationClient( }); const ghAuth = await createGithubInstallationAuth(resolvedInstallationId, ghesApiUrl, appIndex); - return await createOctokitClient(ghAuth.token, ghesApiUrl); + return await createOctokitClient(ghAuth.token, ghesApiUrl, appIndex); } } @@ -103,7 +103,7 @@ export async function scaleUp(payloads: ActionRequestMessageSQS[]): Promise