feat: introduce github caching

Author: dblinkhorn

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

@@ -1,10 +1,9 @@
-import { createAppAuth } from '@octokit/auth-app';
-import { StrategyOptions } from '@octokit/auth-app/dist-types/types';
+import { createAppAuth, type StrategyOptions } from '@octokit/auth-app';
 import { request } from '@octokit/request';
 import { RequestInterface, RequestParameters } from '@octokit/types';
 import { getParameter } from '@aws-github-runner/aws-ssm-util';
 import * as nock from 'nock';
-
+import { reset } from './cache';
 import { createGithubAppAuth, createOctokitClient } from './auth';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 
@@ -29,6 +28,7 @@ const PARAMETER_GITHUB_APP_KEY_BASE64_NAME = `/actions-runner/${ENVIRONMENT}/git
 const mockedGet = vi.mocked(getParameter);
 
 beforeEach(() => {
+  reset(); // clear all caches before each test
   vi.resetModules();
   vi.clearAllMocks();
   process.env = { ...cleanEnv };

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

@@ -22,74 +22,106 @@ import { throttling } from '@octokit/plugin-throttling';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import { getParameter } from '@aws-github-runner/aws-ssm-util';
 import { EndpointDefaults } from '@octokit/types';
+import {
+  getClient,
+  getAuthObject,
+  getAuthConfig,
+  createTokenCacheKey,
+  createAuthCacheKey,
+  createAuthConfigCacheKey,
+} from './cache';
+import type { GithubAppConfig } from './types';
 
 const logger = createChildLogger('gh-auth');
 
 export async function createOctokitClient(token: string, ghesApiUrl = ''): Promise<Octokit> {
-  const CustomOctokit = Octokit.plugin(throttling);
-  const ocktokitOptions: OctokitOptions = {
-    auth: token,
-  };
-  if (ghesApiUrl) {
-    ocktokitOptions.baseUrl = ghesApiUrl;
-    ocktokitOptions.previews = ['antiope'];
-  }
+  const cacheKey = createTokenCacheKey(token, ghesApiUrl);
 
-  return new CustomOctokit({
-    ...ocktokitOptions,
-    userAgent: process.env.USER_AGENT || 'github-aws-runners',
-    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}`);
+  return getClient(cacheKey, async () => {
+    const CustomOctokit = Octokit.plugin(throttling);
+    const ocktokitOptions: OctokitOptions = {
+      auth: token,
+    };
+    if (ghesApiUrl) {
+      ocktokitOptions.baseUrl = ghesApiUrl;
+      ocktokitOptions.previews = ['antiope'];
+    }
+
+    return new CustomOctokit({
+      ...ocktokitOptions,
+      userAgent: process.env.USER_AGENT || 'github-aws-runners',
+      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}`);
+        },
       },
-    },
+    });
   });
 }
 
 export async function createGithubAppAuth(
   installationId: number | undefined,
   ghesApiUrl = '',
 ): Promise<AppAuthentication> {
-  const auth = await createAuth(installationId, ghesApiUrl);
-  const appAuthOptions: AppAuthOptions = { type: 'app' };
-  return auth(appAuthOptions);
+  const cacheKey = createAuthCacheKey('app', installationId, ghesApiUrl);
+
+  return getAuthObject<AppAuthentication>(cacheKey, async () => {
+    const auth = await createAuth(installationId, ghesApiUrl);
+    const appAuthOptions: AppAuthOptions = { type: 'app' };
+    return auth(appAuthOptions);
+  });
 }
 
 export async function createGithubInstallationAuth(
   installationId: number | undefined,
   ghesApiUrl = '',
 ): Promise<InstallationAccessTokenAuthentication> {
-  const auth = await createAuth(installationId, ghesApiUrl);
-  const installationAuthOptions: InstallationAuthOptions = { type: 'installation', installationId };
-  return auth(installationAuthOptions);
+  const cacheKey = createAuthCacheKey('installation', installationId, ghesApiUrl);
+
+  return getAuthObject<InstallationAccessTokenAuthentication>(cacheKey, async () => {
+    const auth = await createAuth(installationId, ghesApiUrl);
+    const installationAuthOptions: InstallationAuthOptions = { type: 'installation', installationId };
+    return auth(installationAuthOptions);
+  });
 }
 
 async function createAuth(installationId: number | undefined, ghesApiUrl: string): Promise<AuthInterface> {
-  const appId = parseInt(await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME));
-  let authOptions: StrategyOptions = {
-    appId,
-    privateKey: Buffer.from(
+  const configCacheKey = createAuthConfigCacheKey(ghesApiUrl);
+
+  const config = await getAuthConfig(configCacheKey, async (): Promise<GithubAppConfig> => {
+    const appId = parseInt(await getParameter(process.env.PARAMETER_GITHUB_APP_ID_NAME));
+    const privateKey = Buffer.from(
       await getParameter(process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME),
       'base64',
       // replace literal \n characters with new lines to allow the key to be stored as a
       // single line variable. This logic should match how the GitHub Terraform provider
       // processes private keys to retain compatibility between the projects
     )
       .toString()
-      .replace('/[\\n]/g', String.fromCharCode(10)),
+      .replace('/[\\n]/g', String.fromCharCode(10));
+
+    return {
+      appId,
+      privateKey,
+    };
+  });
+
+  let authOptions: StrategyOptions = {
+    appId: config.appId,
+    privateKey: config.privateKey,
   };
   if (installationId) authOptions = { ...authOptions, installationId };
 
   logger.debug(`GHES API URL: ${ghesApiUrl}`);
   if (ghesApiUrl) {
     authOptions.request = request.defaults({
       baseUrl: ghesApiUrl,
-    });
+    }) as RequestInterface;
   }
   return createAppAuth(authOptions);
 }

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

@@ -0,0 +1,55 @@
+import { Octokit } from '@octokit/rest';
+import type { GithubAppConfig } from './types';
+
+const clients = new Map<string, Octokit>();
+const authObjects = new Map<string, any>();
+const authConfigs = new Map<string, GithubAppConfig>();
+
+export function createTokenCacheKey(token: string, ghesApiUrl: string = '') {
+  return `octokit-${token}-${ghesApiUrl}`;
+}
+
+export function createAuthCacheKey(type: 'app' | 'installation', installationId?: number, ghesApiUrl: string = '') {
+  const id = installationId || 'none';
+  return `${type}-auth-${id}-${ghesApiUrl}`;
+}
+
+export function createAuthConfigCacheKey(ghesApiUrl: string = '') {
+  return `auth-config-${ghesApiUrl}`;
+}
+
+export async function getClient(key: string, creator: () => Promise<Octokit>): Promise<Octokit> {
+  if (clients.has(key)) {
+    return clients.get(key)!;
+  }
+
+  const client = await creator();
+  clients.set(key, client);
+  return client;
+}
+
+export async function getAuthObject<T>(key: string, creator: () => Promise<T>) {
+  if (authObjects.has(key)) {
+    return authObjects.get(key) as T;
+  }
+
+  const authObj = await creator();
+  authObjects.set(key, authObj);
+  return authObj;
+}
+
+export async function getAuthConfig(key: string, creator: () => Promise<GithubAppConfig>) {
+  if (authConfigs.has(key)) {
+    return authConfigs.get(key)!;
+  }
+
+  const config = await creator();
+  authConfigs.set(key, config);
+  return config;
+}
+
+export function reset() {
+  clients.clear();
+  authObjects.clear();
+  authConfigs.clear();
+}

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

@@ -0,0 +1,2 @@
+export * from './cache';
+export * from './types';

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

@@ -0,0 +1,5 @@
+export interface GithubAppConfig {
+  appId: number;
+  privateKey: string;
+  installationId?: number;
+}

modules/multi-runner/README.md

@@ -171,7 +171,7 @@ module "multi-runner" {
 | <a name="input_runners_ssm_housekeeper"></a> [runners\_ssm\_housekeeper](#input\_runners\_ssm\_housekeeper) | Configuration for the SSM housekeeper lambda. This lambda deletes token / JIT config from SSM.<br/><br/>  `schedule_expression`: is used to configure the schedule for the lambda.<br/>  `enabled`: enable or disable the lambda trigger via the EventBridge.<br/>  `lambda_memory_size`: lambda memory size limit.<br/>  `lambda_timeout`: timeout for the lambda in seconds.<br/>  `config`: configuration for the lambda function. Token path will be read by default from the module. | <pre>object({<br/>    schedule_expression = optional(string, "rate(1 day)")<br/>    enabled             = optional(bool, true)<br/>    lambda_memory_size  = optional(number, 512)<br/>    lambda_timeout      = optional(number, 60)<br/>    config = object({<br/>      tokenPath      = optional(string)<br/>      minimumDaysOld = optional(number, 1)<br/>      dryRun         = optional(bool, false)<br/>    })<br/>  })</pre> | <pre>{<br/>  "config": {}<br/>}</pre> | no |
 | <a name="input_scale_down_lambda_memory_size"></a> [scale\_down\_lambda\_memory\_size](#input\_scale\_down\_lambda\_memory\_size) | Memory size limit in MB for scale down. | `number` | `512` | no |
 | <a name="input_scale_up_lambda_memory_size"></a> [scale\_up\_lambda\_memory\_size](#input\_scale\_up\_lambda\_memory\_size) | Memory size limit in MB for scale\_up lambda. | `number` | `512` | no |
-| <a name="input_ssm_paths"></a> [ssm\_paths](#input\_ssm\_paths) | The root path used in SSM to store configuration and secreets. | <pre>object({<br/>    root    = optional(string, "github-action-runners")<br/>    app     = optional(string, "app")<br/>    runners = optional(string, "runners")<br/>    webhook = optional(string, "webhook")<br/>  })</pre> | `{}` | no |
+| <a name="input_ssm_paths"></a> [ssm\_paths](#input\_ssm\_paths) | The root path used in SSM to store configuration and secrets. | <pre>object({<br/>    root    = optional(string, "github-action-runners")<br/>    app     = optional(string, "app")<br/>    runners = optional(string, "runners")<br/>    webhook = optional(string, "webhook")<br/>  })</pre> | `{}` | no |
 | <a name="input_state_event_rule_binaries_syncer"></a> [state\_event\_rule\_binaries\_syncer](#input\_state\_event\_rule\_binaries\_syncer) | Option to disable EventBridge Lambda trigger for the binary syncer, useful to stop automatic updates of binary distribution | `string` | `"ENABLED"` | no |
 | <a name="input_subnet_ids"></a> [subnet\_ids](#input\_subnet\_ids) | List of subnets in which the action runners will be launched, the subnets needs to be subnets in the `vpc_id`. | `list(string)` | n/a | yes |
 | <a name="input_syncer_lambda_s3_key"></a> [syncer\_lambda\_s3\_key](#input\_syncer\_lambda\_s3\_key) | S3 key for syncer lambda function. Required if using S3 bucket to specify lambdas. | `string` | `null` | no |