diff --git a/lambdas/functions/control-plane/src/github/auth.test.ts b/lambdas/functions/control-plane/src/github/auth.test.ts index 274496ea20..f18c9a9d86 100644 --- a/lambdas/functions/control-plane/src/github/auth.test.ts +++ b/lambdas/functions/control-plane/src/github/auth.test.ts @@ -6,7 +6,7 @@ import { getParameters } from '@aws-github-runner/aws-ssm-util'; import { generateKeyPairSync } from 'node:crypto'; import * as nock from 'nock'; -import { createGithubAppAuth, createOctokitClient } from './auth'; +import { createGithubAppAuth, createOctokitClient, onRateLimit, onSecondaryRateLimit } from './auth'; import { describe, it, expect, beforeEach, vi } from 'vitest'; type MockProxy = T & { @@ -297,3 +297,27 @@ describe('Test createGithubAppAuth', () => { expect(result.token).toBe(token); }); }); + +describe('Test throttling retry caps', () => { + // The plugin passes retryCount as the 4th argument and uses the boolean return + // value to decide whether to retry (see @octokit/plugin-throttling wrap-request). + const options = { method: 'GET', url: '/repos/o/r' } as never; + const octokit = {} as never; + + it.each([ + [0, true], + [1, true], + [2, false], + [3, false], + ])('onRateLimit retries at retryCount=%i -> %s', (retryCount, expected) => { + expect(onRateLimit(60, options, octokit, retryCount)).toBe(expected); + }); + + it.each([ + [0, true], + [1, false], + [2, false], + ])('onSecondaryRateLimit retries at retryCount=%i -> %s', (retryCount, expected) => { + expect(onSecondaryRateLimit(60, options, octokit, retryCount)).toBe(expected); + }); +}); diff --git a/lambdas/functions/control-plane/src/github/auth.ts b/lambdas/functions/control-plane/src/github/auth.ts index 9a572c48a8..9d4e25be95 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -1,5 +1,5 @@ import { createAppAuth, type AppAuthentication, type InstallationAccessTokenAuthentication } from '@octokit/auth-app'; -import type { OctokitOptions } from '@octokit/core'; +import type { OctokitOptions, Octokit as CoreOctokit } from '@octokit/core'; import type { RequestInterface } from '@octokit/types'; // Define types that are not directly exported @@ -27,6 +27,48 @@ import { EndpointDefaults } from '@octokit/types'; const logger = createChildLogger('gh-auth'); +// Retry caps for the throttling plugin. Returning `true` from a limit handler tells +// the plugin to retry after the interval GitHub asked for; returning `false` gives up. +// Primary rate limits reset on a fixed schedule, so a couple of retries is worthwhile. +// Secondary rate limits are abuse-detection signals — retry once, then back off and +// let the message return to the queue rather than pushing harder. +const MAX_RATE_LIMIT_RETRIES = 2; +const MAX_SECONDARY_RATE_LIMIT_RETRIES = 1; + +// Exported for tests: the plugin only surfaces these via the client constructor, +// so there is no other seam to assert the retry cap against. +export function onRateLimit( + retryAfter: number, + options: Required, + // The throttling plugin types this as @octokit/core's Octokit, not the wider + // @octokit/rest one imported above; matching it keeps the handler assignable to + // the plugin's LimitHandler. Unused here regardless. + _octokit: CoreOctokit, + retryCount: number, +): boolean { + logger.warn( + `GitHub rate limit: Request quota exhausted for request ${options.method} ${options.url}, ` + + `retrying after ${retryAfter}s`, + ); + return retryCount < MAX_RATE_LIMIT_RETRIES; +} + +export function onSecondaryRateLimit( + retryAfter: number, + options: Required, + // The throttling plugin types this as @octokit/core's Octokit, not the wider + // @octokit/rest one imported above; matching it keeps the handler assignable to + // the plugin's LimitHandler. Unused here regardless. + _octokit: CoreOctokit, + retryCount: number, +): boolean { + logger.warn( + `GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}, ` + + `retrying after ${retryAfter}s`, + ); + return retryCount < MAX_SECONDARY_RATE_LIMIT_RETRIES; +} + export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise { const CustomOctokit = Octokit.plugin(retry, throttling); const ocktokitOptions: OctokitOptions = { @@ -52,14 +94,8 @@ export async function createOctokitClient(token: string, ghesApiUrl = ''): Promi }, }, throttle: { - onRateLimit: (retryAfter: number, options: Required) => { - logger.warn( - `GitHub rate limit: Request quota exhausted for request ${options.method} ${options.url}. Requested `, - ); - }, - onSecondaryRateLimit: (retryAfter: number, options: Required) => { - logger.warn(`GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}`); - }, + onRateLimit, + onSecondaryRateLimit, }, }); } diff --git a/lambdas/libs/aws-ssm-util/src/index.test.ts b/lambdas/libs/aws-ssm-util/src/index.test.ts index 51a9d65027..8a1d8d3864 100644 --- a/lambdas/libs/aws-ssm-util/src/index.test.ts +++ b/lambdas/libs/aws-ssm-util/src/index.test.ts @@ -10,7 +10,7 @@ import 'aws-sdk-client-mock-jest/vitest'; import { mockClient } from 'aws-sdk-client-mock'; import nock from 'nock'; -import { getParameter, getParameters, putParameter, SSM_ADVANCED_TIER_THRESHOLD } from '.'; +import { getParameter, getParameters, putParameter, resetSSMClient, ssmClient, SSM_ADVANCED_TIER_THRESHOLD } from '.'; import { describe, it, expect, beforeEach, vi } from 'vitest'; const mockSSMClient = mockClient(SSMClient); @@ -20,6 +20,7 @@ beforeEach(() => { vi.resetModules(); vi.clearAllMocks(); process.env = { ...cleanEnv }; + resetSSMClient(); nock.disableNetConnect(); }); @@ -254,3 +255,30 @@ describe('Test getParameters (batch)', () => { expect(result).toEqual(new Map()); }); }); + +describe('SSM client configuration', () => { + it('configures adaptive retry with a raised attempt cap', async () => { + const config = ssmClient().config; + + expect(await config.maxAttempts()).toBe(10); + expect(config.retryMode).toBe('adaptive'); + }); + + it('reuses one client so adaptive rate-sensing state survives across calls', async () => { + // Adaptive retry keeps its token bucket on the client instance; a fresh client + // per call would silently downgrade `adaptive` to plain retries. + mockSSMClient.on(GetParametersCommand).resolves({ Parameters: [] }); + + await getParameters(['/app/one']); + await getParameters(['/app/two']); + + expect(ssmClient()).toBe(ssmClient()); + }); + + it('picks up the region from the environment on first use', async () => { + process.env.AWS_REGION = 'eu-north-1'; + resetSSMClient(); + + expect(await ssmClient().config.region()).toBe('eu-north-1'); + }); +}); diff --git a/lambdas/libs/aws-ssm-util/src/index.ts b/lambdas/libs/aws-ssm-util/src/index.ts index 9173cbb210..71b33cbf41 100644 --- a/lambdas/libs/aws-ssm-util/src/index.ts +++ b/lambdas/libs/aws-ssm-util/src/index.ts @@ -2,9 +2,41 @@ import { GetParametersCommand, PutParameterCommand, SSMClient, Tag } from '@aws- import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util'; import { SSMProvider } from '@aws-lambda-powertools/parameters/ssm'; +// SSM PutParameter has a per-account, per-region rate limit (~40 TPS standard +// throughput). Under burst load with multiple concurrent Lambdas each writing +// JIT configs, the default retry (standard, 3 attempts, ~3s budget) is +// insufficient and throws ThrottlingException. +// +// `adaptive` retry mode adds client-side rate-sensing via a token bucket: +// when the SDK sees ThrottlingException it slows further calls to match the +// observed budget. Combined with maxAttempts=10 this gives ~30s of retry +// per call without hammering the API. +// +// The client is memoised deliberately. Adaptive retry keeps its rate-sensing +// token bucket on the client instance, so constructing a fresh SSMClient per +// call would discard what it learned and reduce `adaptive` to plain retries. +// Reusing one client per Lambda container lets the backoff carry across calls +// (and saves the per-call construction cost). +let memoisedClient: SSMClient | undefined; + +export function ssmClient(): SSMClient { + memoisedClient ??= getTracedAWSV3Client( + new SSMClient({ + region: process.env.AWS_REGION, + maxAttempts: 10, + retryMode: 'adaptive', + }), + ); + return memoisedClient; +} + +// Exposed for tests, which need a fresh client per case to assert construction. +export function resetSSMClient(): void { + memoisedClient = undefined; +} + export async function getParameter(parameter_name: string): Promise { - const ssmClient = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION })); - const client = new SSMProvider({ awsSdkV3Client: ssmClient }); //getTracedAWSV3Client(); + const client = new SSMProvider({ awsSdkV3Client: ssmClient() }); const result = await client.get(parameter_name, { decrypt: true, maxAge: 30, // 30 seconds override default 5 seconds @@ -48,14 +80,13 @@ export async function getParameters(parameter_names: string[]): Promise(); // AWS SSM GetParameters API has a limit of 10 parameters per call const chunkSize = 10; for (let i = 0; i < parameter_names.length; i += chunkSize) { const chunk = parameter_names.slice(i, i + chunkSize); - const response = await ssmClient.send( + const response = await ssmClient().send( new GetParametersCommand({ Names: chunk, WithDecryption: true, @@ -80,7 +111,7 @@ export async function putParameter( secure: boolean, options: { tags?: Tag[] } = {}, ): Promise { - const client = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION })); + const client = ssmClient(); // Determine tier based on parameter_value size const valueSizeBytes = Buffer.byteLength(parameter_value, 'utf8');