fix: gracefully handle JIT config failures and terminate unconfigured instance (#4990)

This pull request enhances the robustness and reliability of the GitHub
Actions runner scaling logic by improving error handling and retry
mechanisms for GitHub API calls. It introduces the
`@octokit/plugin-retry` plugin to automatically retry failed API
requests, adds detailed logging for retry attempts, and ensures that
failures in creating JIT configs for individual runner instances do not
halt the entire scaling process. Additionally, new tests are added to
verify handling of various API failure scenarios.

**GitHub API client improvements:**

* Added `@octokit/plugin-retry` to dependencies (`package.json`) and
integrated it into the Octokit client initialization to enable automatic
retries for failed GitHub API requests.
[[1]](diffhunk://#diff-37d09418dae74ded5678eabfa3b60993ee491e2fd9e49e11142f639b078ac9b2R41)
[[2]](diffhunk://#diff-cf7cdd79fe0ed0e3a2e8928c0c7667a096c47c47abdb2354ddadee67e80a226dR21)
[[3]](diffhunk://#diff-cf7cdd79fe0ed0e3a2e8928c0c7667a096c47c47abdb2354ddadee67e80a226dL29-R30)
* Configured the retry plugin to log detailed warnings on each retry
attempt, including the HTTP method, URL, error message, and status code.

**Error handling and resilience in JIT config creation:**

* Updated `createJitConfig` in `scale-up.ts` to catch and log errors for
individual runner instances when creating JIT configs, allowing the
process to continue for remaining instances and logging a summary of
failed attempts at the end.
[[1]](diffhunk://#diff-fbc68af2a40bf14ad13a80b13958c0b52d1d0fde5f0009416a693fb4b691ceaeR537-R542)
[[2]](diffhunk://#diff-fbc68af2a40bf14ad13a80b13958c0b52d1d0fde5f0009416a693fb4b691ceaeR582-R596)
* Instances that failed to generate a configuration, will now be
terminated to avoid generating waste.

**Testing improvements:**

* Added comprehensive tests to `scale-up.test.ts` to verify correct
behavior when GitHub API calls fail for some instances, including
retryable errors (e.g., 5xx), non-retryable errors (e.g., 4xx), and
partial failures, ensuring only successful JIT configs are stored.

Author: Brend-Smits

lambdas/functions/control-plane/package.json

@@ -38,6 +38,7 @@
     "@middy/core": "^6.4.5",
     "@octokit/auth-app": "8.2.0",
     "@octokit/core": "7.0.6",
+    "@octokit/plugin-retry": "8.0.3",
     "@octokit/plugin-throttling": "11.0.3",
     "@octokit/rest": "22.0.1",
     "cron-parser": "^5.4.0"

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

@@ -19,6 +19,7 @@ type StrategyOptions = {
 import { createSign, randomUUID } from 'node:crypto';
 import { request } from '@octokit/request';
 import { Octokit } from '@octokit/rest';
+import { retry } from '@octokit/plugin-retry';
 import { throttling } from '@octokit/plugin-throttling';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import { getParameter } from '@aws-github-runner/aws-ssm-util';
@@ -27,7 +28,7 @@ import { EndpointDefaults } from '@octokit/types';
 const logger = createChildLogger('gh-auth');
 
 export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise<Octokit> {
-  const CustomOctokit = Octokit.plugin(throttling);
+  const CustomOctokit = Octokit.plugin(retry, throttling);
   const ocktokitOptions: OctokitOptions = {
     auth: token,
   };
@@ -39,6 +40,17 @@ export async function createOctokitClient(token: string, ghesApiUrl = ''): Promi
   return new CustomOctokit({
     ...ocktokitOptions,
     userAgent: process.env.USER_AGENT || 'github-aws-runners',
+    retry: {
+      onRetry: (retryCount: number, error: Error, request: { method: string; url: string }) => {
+        logger.warn('GitHub API request retry attempt', {
+          retryCount,
+          method: request.method,
+          url: request.url,
+          error: error.message,
+          status: (error as Error & { status?: number }).status,
+        });
+      },
+    },
     throttle: {
       onRateLimit: (retryAfter: number, options: Required<EndpointDefaults>) => {
         logger.warn(

lambdas/functions/control-plane/src/scale-runners/scale-up.test.ts

@@ -343,6 +343,172 @@ describe('scaleUp with GHES', () => {
         ],
       });
     });
+
+    it('should create JIT config for all remaining instances even when GitHub API fails for one instance', async () => {
+      process.env.RUNNERS_MAXIMUM_COUNT = '5';
+      mockCreateRunner.mockImplementation(async () => {
+        return ['i-instance-1', 'i-instance-2', 'i-instance-3'];
+      });
+      mockListRunners.mockImplementation(async () => {
+        return [];
+      });
+
+      mockOctokit.actions.generateRunnerJitconfigForOrg.mockImplementation(({ name }) => {
+        if (name === 'unit-test-i-instance-2') {
+          // Simulate a 503 Service Unavailable error from GitHub
+          const error = new Error('Service Unavailable') as Error & {
+            status: number;
+            response: { status: number; data: { message: string } };
+          };
+          error.status = 503;
+          error.response = {
+            status: 503,
+            data: { message: 'Service temporarily unavailable' },
+          };
+          throw error;
+        }
+        return {
+          data: {
+            runner: { id: 9876543210 },
+            encoded_jit_config: `TEST_JIT_CONFIG_${name}`,
+          },
+          headers: {},
+        };
+      });
+
+      await scaleUpModule.scaleUp(TEST_DATA);
+
+      expect(mockOctokit.actions.generateRunnerJitconfigForOrg).toHaveBeenCalledWith({
+        org: TEST_DATA_SINGLE.repositoryOwner,
+        name: 'unit-test-i-instance-1',
+        runner_group_id: 1,
+        labels: ['label1', 'label2'],
+      });
+
+      expect(mockOctokit.actions.generateRunnerJitconfigForOrg).toHaveBeenCalledWith({
+        org: TEST_DATA_SINGLE.repositoryOwner,
+        name: 'unit-test-i-instance-2',
+        runner_group_id: 1,
+        labels: ['label1', 'label2'],
+      });
+
+      expect(mockOctokit.actions.generateRunnerJitconfigForOrg).toHaveBeenCalledWith({
+        org: TEST_DATA_SINGLE.repositoryOwner,
+        name: 'unit-test-i-instance-3',
+        runner_group_id: 1,
+        labels: ['label1', 'label2'],
+      });
+
+      expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-1',
+        Value: 'TEST_JIT_CONFIG_unit-test-i-instance-1',
+        Type: 'SecureString',
+        Tags: [{ Key: 'InstanceId', Value: 'i-instance-1' }],
+      });
+
+      expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-3',
+        Value: 'TEST_JIT_CONFIG_unit-test-i-instance-3',
+        Type: 'SecureString',
+        Tags: [{ Key: 'InstanceId', Value: 'i-instance-3' }],
+      });
+
+      expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-2',
+      });
+    });
+
+    it('should handle retryable errors with error handling logic', async () => {
+      process.env.RUNNERS_MAXIMUM_COUNT = '5';
+      mockCreateRunner.mockImplementation(async () => {
+        return ['i-instance-1', 'i-instance-2'];
+      });
+      mockListRunners.mockImplementation(async () => {
+        return [];
+      });
+
+      mockOctokit.actions.generateRunnerJitconfigForOrg.mockImplementation(({ name }) => {
+        if (name === 'unit-test-i-instance-1') {
+          const error = new Error('Internal Server Error') as Error & {
+            status: number;
+            response: { status: number; data: { message: string } };
+          };
+          error.status = 500;
+          error.response = {
+            status: 500,
+            data: { message: 'Internal server error' },
+          };
+          throw error;
+        }
+        return {
+          data: {
+            runner: { id: 9876543210 },
+            encoded_jit_config: `TEST_JIT_CONFIG_${name}`,
+          },
+          headers: {},
+        };
+      });
+
+      await scaleUpModule.scaleUp(TEST_DATA);
+
+      expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-2',
+        Value: 'TEST_JIT_CONFIG_unit-test-i-instance-2',
+        Type: 'SecureString',
+        Tags: [{ Key: 'InstanceId', Value: 'i-instance-2' }],
+      });
+
+      expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-1',
+      });
+    });
+
+    it('should handle non-retryable 4xx errors gracefully', async () => {
+      process.env.RUNNERS_MAXIMUM_COUNT = '5';
+      mockCreateRunner.mockImplementation(async () => {
+        return ['i-instance-1', 'i-instance-2'];
+      });
+      mockListRunners.mockImplementation(async () => {
+        return [];
+      });
+
+      mockOctokit.actions.generateRunnerJitconfigForOrg.mockImplementation(({ name }) => {
+        if (name === 'unit-test-i-instance-1') {
+          // 404 is not retryable - will fail immediately
+          const error = new Error('Not Found') as Error & {
+            status: number;
+            response: { status: number; data: { message: string } };
+          };
+          error.status = 404;
+          error.response = {
+            status: 404,
+            data: { message: 'Resource not found' },
+          };
+          throw error;
+        }
+        return {
+          data: {
+            runner: { id: 9876543210 },
+            encoded_jit_config: `TEST_JIT_CONFIG_${name}`,
+          },
+          headers: {},
+        };
+      });
+
+      await scaleUpModule.scaleUp(TEST_DATA);
+
+      expect(mockSSMClient).toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-2',
+        Value: 'TEST_JIT_CONFIG_unit-test-i-instance-2',
+        Type: 'SecureString',
+        Tags: [{ Key: 'InstanceId', Value: 'i-instance-2' }],
+      });
+
+      expect(mockSSMClient).not.toHaveReceivedCommandWith(PutParameterCommand, {
+        Name: '/github-action-runners/default/runners/config/i-instance-1',
+      });
+    });
+
     it.each(RUNNER_TYPES)(
       'calls create start runner config of 40' + ' instances (ssm rate limit condition) to test time delay ',
       async (type: RunnerType) => {