fix: enable Octokit throttle retry and SSM adaptive retry under burst

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

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<EndpointDefaults>,
+  // 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<EndpointDefaults>,
+  // 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<Octokit> {
   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<EndpointDefaults>) => {
-        logger.warn(
-          `GitHub rate limit: Request quota exhausted for request ${options.method} ${options.url}. Requested `,
-        );
-      },
-      onSecondaryRateLimit: (retryAfter: number, options: Required<EndpointDefaults>) => {
-        logger.warn(`GitHub rate limit: SecondaryRateLimit detected for request ${options.method} ${options.url}`);
-      },
+      onRateLimit,
+      onSecondaryRateLimit,
     },
   });
 }

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

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<string> {
-  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<Map<stri
     return new Map();
   }
 
-  const ssmClient = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION }));
   const result = new Map<string, string>();
 
   // 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<void> {
-  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');