feat: add retry logic to Octokit client and add ScaleUp tests

Author: Brend-Smits

lambdas/functions/control-plane/package.json

@@ -38,6 +38,7 @@
     "@middy/core": "^6.4.5",
     "@octokit/auth-app": "8.1.2",
     "@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

@@ -18,6 +18,7 @@ type StrategyOptions = {
 };
 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';
@@ -26,7 +27,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,
   };
@@ -38,6 +39,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

@@ -346,12 +346,13 @@ describe('scaleUp with GHES', () => {
         return [];
       });
 
-      let callCount = 0;
       mockOctokit.actions.generateRunnerJitconfigForOrg.mockImplementation(({ name }) => {
-        callCount++;
         if (name === 'unit-test-i-instance-2') {
           // Simulate a 503 Service Unavailable error from GitHub
-          const error = new Error('Service Unavailable') as any;
+          const error = new Error('Service Unavailable') as Error & {
+            status: number;
+            response: { status: number; data: { message: string } };
+          };
           error.status = 503;
           error.response = {
             status: 503,
@@ -361,7 +362,7 @@ describe('scaleUp with GHES', () => {
         }
         return {
           data: {
-            runner: { id: 9876543210 + callCount },
+            runner: { id: 9876543210 },
             encoded_jit_config: `TEST_JIT_CONFIG_${name}`,
           },
           headers: {},
@@ -410,6 +411,97 @@ describe('scaleUp with GHES', () => {
       });
     });
 
+    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) => {

lambdas/yarn.lock

@@ -155,6 +155,7 @@ __metadata:
     "@middy/core": "npm:^6.4.5"
     "@octokit/auth-app": "npm:8.1.2"
     "@octokit/core": "npm:7.0.6"
+    "@octokit/plugin-retry": "npm:8.0.3"
     "@octokit/plugin-throttling": "npm:11.0.3"
     "@octokit/rest": "npm:22.0.1"
     "@octokit/types": "npm:^16.0.0"
@@ -3685,6 +3686,19 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@octokit/plugin-retry@npm:8.0.3":
+  version: 8.0.3
+  resolution: "@octokit/plugin-retry@npm:8.0.3"
+  dependencies:
+    "@octokit/request-error": "npm:^7.0.2"
+    "@octokit/types": "npm:^16.0.0"
+    bottleneck: "npm:^2.15.3"
+  peerDependencies:
+    "@octokit/core": ">=7"
+  checksum: 10c0/24d35d85f750f9e3e52f63b8ddd8fc8aa7bdd946c77b9ea4d6894d026c5c2c69109e8de3880a9970c906f624eb777c7d0c0a2072e6d41dadc7b36cce104b978c
+  languageName: node
+  linkType: hard
+
 "@octokit/plugin-throttling@npm:11.0.3":
   version: 11.0.3
   resolution: "@octokit/plugin-throttling@npm:11.0.3"