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