From 2ec44dc35e4aa0bcfa3a91ada95ea9300d93cd6e Mon Sep 17 00:00:00 2001 From: Vegard Hansen Date: Wed, 27 May 2026 06:44:19 +0200 Subject: [PATCH 1/3] fix: enable Octokit throttle retry and SSM adaptive retry under burst Two resilience improvements for burst load: 1. Octokit plugin-throttling callbacks now return true (with caps) so rate-limited GitHub API requests are retried instead of immediately failing. Primary: up to 2 retries. Secondary: 1 retry. 2. SSM client switched to adaptive retry mode with maxAttempts=10, giving ~30s of retry budget for PutParameter under throttling. Standard mode (3 attempts, ~3s) is insufficient when multiple concurrent Lambdas write JIT configs during burst. Fixes #5135 --- .../control-plane/src/github/auth.ts | 8 +++++-- lambdas/libs/aws-ssm-util/src/index.ts | 21 ++++++++++++++++--- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/lambdas/functions/control-plane/src/github/auth.ts b/lambdas/functions/control-plane/src/github/auth.ts index 9a572c48a8..fea7f3fb07 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -54,11 +54,15 @@ 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 `, + `GitHub rate limit: Request quota exhausted for request ${options.method} ${options.url}, retrying after ${retryAfter}s`, ); + return (options as unknown as { request: { retryCount: number } }).request.retryCount < 2; }, onSecondaryRateLimit: (retryAfter: number, options: Required) => { - logger.warn(`GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}`); + logger.warn( + `GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}, retrying after ${retryAfter}s`, + ); + return (options as unknown as { request: { retryCount: number } }).request.retryCount < 1; }, }, }); diff --git a/lambdas/libs/aws-ssm-util/src/index.ts b/lambdas/libs/aws-ssm-util/src/index.ts index 9173cbb210..c4cfb67a00 100644 --- a/lambdas/libs/aws-ssm-util/src/index.ts +++ b/lambdas/libs/aws-ssm-util/src/index.ts @@ -2,8 +2,23 @@ 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. +const SSM_CLIENT_CONFIG = { + region: process.env.AWS_REGION, + maxAttempts: 10, + retryMode: 'adaptive' as const, +}; + export async function getParameter(parameter_name: string): Promise { - const ssmClient = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION })); + const ssmClient = getTracedAWSV3Client(new SSMClient(SSM_CLIENT_CONFIG)); const client = new SSMProvider({ awsSdkV3Client: ssmClient }); //getTracedAWSV3Client(); const result = await client.get(parameter_name, { decrypt: true, @@ -48,7 +63,7 @@ export async function getParameters(parameter_names: string[]): Promise(); // AWS SSM GetParameters API has a limit of 10 parameters per call @@ -80,7 +95,7 @@ export async function putParameter( secure: boolean, options: { tags?: Tag[] } = {}, ): Promise { - const client = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION })); + const client = getTracedAWSV3Client(new SSMClient(SSM_CLIENT_CONFIG)); // Determine tier based on parameter_value size const valueSizeBytes = Buffer.byteLength(parameter_value, 'utf8'); From 1ca6aef08d004557b4b4edc513f4fd2acbedb406 Mon Sep 17 00:00:00 2001 From: Vegard Hansen Date: Tue, 21 Jul 2026 11:14:14 +0200 Subject: [PATCH 2/3] refactor: address review feedback on throttle and SSM retry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Use the retryCount argument the throttling plugin already passes as its 4th parameter instead of casting options through unknown. Extract the two limit handlers so the retry caps are testable at all. - Memoise the SSMClient. Adaptive retry keeps its rate-sensing token bucket on the client instance, so constructing a fresh client per call discarded it and reduced retryMode: 'adaptive' to plain retries — the change was close to a no-op as written. - Cover both: retry-cap boundaries for the limit handlers, and the client's retry configuration and reuse. --- .../control-plane/src/github/auth.test.ts | 26 +++++++++- .../control-plane/src/github/auth.ts | 50 ++++++++++++++----- lambdas/libs/aws-ssm-util/src/index.test.ts | 30 ++++++++++- lambdas/libs/aws-ssm-util/src/index.ts | 36 +++++++++---- 4 files changed, 118 insertions(+), 24 deletions(-) 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 fea7f3fb07..09aeb40554 100644 --- a/lambdas/functions/control-plane/src/github/auth.ts +++ b/lambdas/functions/control-plane/src/github/auth.ts @@ -27,6 +27,42 @@ 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, + _octokit: Octokit, + 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, + _octokit: Octokit, + 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,18 +88,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}, retrying after ${retryAfter}s`, - ); - return (options as unknown as { request: { retryCount: number } }).request.retryCount < 2; - }, - onSecondaryRateLimit: (retryAfter: number, options: Required) => { - logger.warn( - `GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}, retrying after ${retryAfter}s`, - ); - return (options as unknown as { request: { retryCount: number } }).request.retryCount < 1; - }, + 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 c4cfb67a00..71b33cbf41 100644 --- a/lambdas/libs/aws-ssm-util/src/index.ts +++ b/lambdas/libs/aws-ssm-util/src/index.ts @@ -11,15 +11,32 @@ import { SSMProvider } from '@aws-lambda-powertools/parameters/ssm'; // 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. -const SSM_CLIENT_CONFIG = { - region: process.env.AWS_REGION, - maxAttempts: 10, - retryMode: 'adaptive' as const, -}; +// +// 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(SSM_CLIENT_CONFIG)); - 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 @@ -63,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, @@ -95,7 +111,7 @@ export async function putParameter( secure: boolean, options: { tags?: Tag[] } = {}, ): Promise { - const client = getTracedAWSV3Client(new SSMClient(SSM_CLIENT_CONFIG)); + const client = ssmClient(); // Determine tier based on parameter_value size const valueSizeBytes = Buffer.byteLength(parameter_value, 'utf8'); From f3f65e998b80f7014927521ed0e469218cabdca5 Mon Sep 17 00:00:00 2001 From: Vegard Hansen Date: Wed, 22 Jul 2026 15:06:41 +0200 Subject: [PATCH 3/3] fix(auth): type throttle handler octokit param as core Octokit The extracted onRateLimit/onSecondaryRateLimit handlers typed their unused octokit parameter as @octokit/rest's Octokit, which is wider than the @octokit/core Octokit the plugin's LimitHandler expects. Parameter contravariance made the handlers unassignable, failing the dist build (tsc) while vitest's transpile-only path missed it. --- lambdas/functions/control-plane/src/github/auth.ts | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/lambdas/functions/control-plane/src/github/auth.ts b/lambdas/functions/control-plane/src/github/auth.ts index 09aeb40554..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 @@ -40,7 +40,10 @@ const MAX_SECONDARY_RATE_LIMIT_RETRIES = 1; export function onRateLimit( retryAfter: number, options: Required, - _octokit: Octokit, + // 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( @@ -53,7 +56,10 @@ export function onRateLimit( export function onSecondaryRateLimit( retryAfter: number, options: Required, - _octokit: Octokit, + // 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(