feat(github-app): select app by rate limit budget

README.md

@@ -66,22 +66,22 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 ## Requirements
 
 | Name | Version |
-| ---- | ------- |
+|------|---------|
 | <a name="requirement_terraform"></a> [terraform](#requirement\_terraform) | >= 1.3.0 |
 | <a name="requirement_aws"></a> [aws](#requirement\_aws) | >= 6.33 |
 | <a name="requirement_random"></a> [random](#requirement\_random) | ~> 3.0 |
 
 ## Providers
 
 | Name | Version |
-| ---- | ------- |
-| <a name="provider_aws"></a> [aws](#provider\_aws) | 6.60.0 |
-| <a name="provider_random"></a> [random](#provider\_random) | 3.9.0 |
+|------|---------|
+| <a name="provider_aws"></a> [aws](#provider\_aws) | >= 6.33 |
+| <a name="provider_random"></a> [random](#provider\_random) | ~> 3.0 |
 
 ## Modules
 
 | Name | Source | Version |
-| ---- | ------ | ------- |
+|------|--------|---------|
 | <a name="module_ami_housekeeper"></a> [ami\_housekeeper](#module\_ami\_housekeeper) | ./modules/ami-housekeeper | n/a |
 | <a name="module_instance_termination_watcher"></a> [instance\_termination\_watcher](#module\_instance\_termination\_watcher) | ./modules/termination-watcher | n/a |
 | <a name="module_runner_binaries"></a> [runner\_binaries](#module\_runner\_binaries) | ./modules/runner-binaries-syncer | n/a |
@@ -92,7 +92,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 ## Resources
 
 | Name | Type |
-| ---- | ---- |
+|------|------|
 | [aws_sqs_queue.queued_builds](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/sqs_queue) | resource |
 | [aws_sqs_queue.queued_builds_dlq](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/sqs_queue) | resource |
 | [aws_sqs_queue_policy.build_queue_dlq_policy](https://registry.terraform.io/providers/hashicorp/aws/latest/docs/resources/sqs_queue_policy) | resource |
@@ -103,7 +103,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 ## Inputs
 
 | Name | Description | Type | Default | Required |
-| ---- | ----------- | ---- | ------- | :------: |
+|------|-------------|------|---------|:--------:|
 | <a name="input_additional_github_apps"></a> [additional\_github\_apps](#input\_additional\_github\_apps) | Additional GitHub Apps for distributing API rate limit usage. Each must be installed on the same repos/orgs as the primary app. | <pre>list(object({<br/>    key_base64          = optional(string)<br/>    key_base64_ssm      = optional(object({ arn = string, name = string }))<br/>    id                  = optional(string)<br/>    id_ssm              = optional(object({ arn = string, name = string }))<br/>    installation_id     = optional(string)<br/>    installation_id_ssm = optional(object({ arn = string, name = string }))<br/>  }))</pre> | `[]` | no |
 | <a name="input_ami"></a> [ami](#input\_ami) | AMI configuration for the action runner instances. This object allows you to specify all AMI-related settings in one place.<br/><br/>Parameters:<br/>- `filter`: Map of lists to filter AMIs by various criteria (e.g., { name = ["ubuntu/images/hvm-ssd/ubuntu-jammy-22.04-amd64-*"], state = ["available"] })<br/>- `owners`: List of AMI owners to limit the search. Common values: ["amazon"], ["self"], or specific AWS account IDs<br/>- `id_ssm_parameter_arn`: ARN of an SSM parameter containing the AMI ID. If specified, this overrides both AMI filter and parameter name<br/>- `kms_key_arn`: Optional KMS key ARN if the AMI is encrypted with a customer managed key<br/><br/>Defaults to null, in which case the module falls back to individual AMI variables (deprecated). | <pre>object({<br/>    filter               = optional(map(list(string)), { state = ["available"] })<br/>    owners               = optional(list(string), ["amazon"])<br/>    id_ssm_parameter_arn = optional(string, null)<br/>    kms_key_arn          = optional(string, null)<br/>  })</pre> | `null` | no |
 | <a name="input_ami_housekeeper_cleanup_config"></a> [ami\_housekeeper\_cleanup\_config](#input\_ami\_housekeeper\_cleanup\_config) | Configuration for AMI cleanup.<br/><br/>    `amiFilters` - Filters to use when searching for AMIs to cleanup. Default filter for images owned by the account and that are available.<br/>    `dryRun` - If true, no AMIs will be deregistered. Default false.<br/>    `launchTemplateNames` - Launch template names to use when searching for AMIs to cleanup. Default no launch templates.<br/>    `maxItems` - The maximum number of AMIs that will be queried for cleanup. Default no maximum.<br/>    `minimumDaysOld` - Minimum number of days old an AMI must be to be considered for cleanup. Default 30.<br/>    `ssmParameterNames` - SSM parameter names to use when searching for AMIs to cleanup. This parameter should be set when using SSM to configure the AMI to use. Default no SSM parameters. | <pre>object({<br/>    amiFilters = optional(list(object({<br/>      Name   = string<br/>      Values = list(string)<br/>      })),<br/>      [{<br/>        Name : "state",<br/>        Values : ["available"],<br/>        },<br/>        {<br/>          Name : "image-type",<br/>          Values : ["machine"],<br/>      }]<br/>    )<br/>    dryRun              = optional(bool, false)<br/>    launchTemplateNames = optional(list(string))<br/>    maxItems            = optional(number)<br/>    minimumDaysOld      = optional(number, 30)<br/>    ssmParameterNames   = optional(list(string))<br/>  })</pre> | `{}` | no |
@@ -250,7 +250,7 @@ Join our discord community via [this invite link](https://discord.gg/bxgXW8jJGh)
 ## Outputs
 
 | Name | Description |
-| ---- | ----------- |
+|------|-------------|
 | <a name="output_binaries_syncer"></a> [binaries\_syncer](#output\_binaries\_syncer) | n/a |
 | <a name="output_instance_termination_handler"></a> [instance\_termination\_handler](#output\_instance\_termination\_handler) | n/a |
 | <a name="output_instance_termination_watcher"></a> [instance\_termination\_watcher](#output\_instance\_termination\_watcher) | n/a |

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.

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);
+  });
+});

lambdas/functions/control-plane/src/github/auth.ts

@@ -77,6 +77,66 @@ interface GitHubAppCredential {
 
 let appCredentialsPromise: Promise<GitHubAppCredential[]> | 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<number, AppRateLimitState>();
+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<number> {
+  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,14 +218,15 @@ export async function getAppCount(): Promise<number> {
 
 export function resetAppCredentialsCache(): void {
   appCredentialsPromise = null;
+  appRateLimitStates.clear();
 }
 
 export async function getStoredInstallationId(appIndex: number): Promise<number | undefined> {
   const credentials = await getAppCredentials();
   return credentials[appIndex]?.installationId;
 }
 
-export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise<Octokit> {
+export async function createOctokitClient(token: string, ghesApiUrl = '', appIndex?: number): Promise<Octokit> {
   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<EndpointDefaults>,
+        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<EndpointDefaults>,
+        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<AppAuthentication & { appIndex: number }> {
-  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<InstallationAccessTokenAuthentication> {
-  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<AuthInterface> {
   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`);
 

lambdas/functions/control-plane/src/github/octokit.ts

@@ -68,7 +68,7 @@ export async function getInstallationId(
   appIndex?: number,
 ): Promise<number> {
   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);
   }
 }

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' };