refactor(storage): extract runner group cache store

lambdas/functions/control-plane/src/scale-runners/github-runner.ts

@@ -1,9 +1,10 @@
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
-import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util';
 import {
   createRunnerConfigStore,
+  createRunnerGroupCacheStore,
   type RunnerConfigMetadata,
   type RunnerConfigStore,
+  type RunnerGroupCacheStore,
 } from '@aws-github-runner/storage-providers';
 import { Octokit } from '@octokit/rest';
 
@@ -19,6 +20,7 @@ export interface GitHubRunnerMetadata {
 }
 
 export interface StartRunnerConfigOptions {
+  runnerGroupCacheStore?: RunnerGroupCacheStore;
   getRunnerConfigMetadata?: (runnerId: string) => RunnerConfigMetadata[];
   onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise<void>;
 }
@@ -187,43 +189,22 @@ export async function isJobQueued(
 export async function getRunnerGroupId(
   githubRunnerConfig: CreateGitHubRunnerConfig,
   ghClient: Octokit,
+  runnerGroupCacheStore?: RunnerGroupCacheStore,
 ): Promise<number> {
   // if the runnerType is Repo, then runnerGroupId is default to 1
   let runnerGroupId: number | undefined = 1;
   if (githubRunnerConfig.runnerType === 'Org' && githubRunnerConfig.runnerGroup !== undefined) {
-    let runnerGroup: string | undefined;
-    // check if runner group id is already stored in SSM Parameter Store and
-    // use it if it exists to avoid API call to GitHub
-    try {
-      runnerGroup = await getParameter(
-        `${githubRunnerConfig.ssmConfigPath}/runner-group/${githubRunnerConfig.runnerGroup}`,
-      );
-    } catch (err) {
-      logger.debug('Handling error:', err as Error);
-      logger.warn(
-        `SSM Parameter "${githubRunnerConfig.ssmConfigPath}/runner-group/${githubRunnerConfig.runnerGroup}"
-         for Runner group ${githubRunnerConfig.runnerGroup} does not exist`,
-      );
-    }
+    const cacheStore = runnerGroupCacheStore ?? createRunnerGroupCacheStore();
+    const runnerGroup = await cacheStore.get(githubRunnerConfig.runnerGroup);
     if (runnerGroup === undefined) {
       // get runner group id from GitHub
       runnerGroupId = await getRunnerGroupByName(ghClient, githubRunnerConfig);
-      // store runner group id in SSM
-      try {
-        await putParameter(
-          `${githubRunnerConfig.ssmConfigPath}/runner-group/${githubRunnerConfig.runnerGroup}`,
-          runnerGroupId.toString(),
-          false,
-          {
-            tags: githubRunnerConfig.ssmParameterStoreTags,
-          },
-        );
-      } catch (err) {
-        logger.debug('Error storing runner group id in SSM Parameter Store', err as Error);
-        throw err;
-      }
+      await cacheStore.create({
+        runnerGroupName: githubRunnerConfig.runnerGroup,
+        runnerGroupId,
+      });
     } else {
-      runnerGroupId = parseInt(runnerGroup);
+      runnerGroupId = runnerGroup;
     }
   }
   return runnerGroupId;
@@ -318,7 +299,7 @@ async function createJitConfig(
   runnerConfigStore: RunnerConfigStore,
   options: StartRunnerConfigOptions,
 ): Promise<string[]> {
-  const runnerGroupId = await getRunnerGroupId(githubRunnerConfig, ghClient);
+  const runnerGroupId = await getRunnerGroupId(githubRunnerConfig, ghClient, options.runnerGroupCacheStore);
   const { isDelay, delay, delayMilliseconds } = addDelay(runnerIds, runnerConfigStore);
   const runnerLabels = githubRunnerConfig.runnerLabels.split(',');
   const failedRunnerIds: string[] = [];

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

@@ -141,13 +141,15 @@ let expectedRunnerParams = { ...EXPECTED_RUNNER_PARAMS };
 function setDefaults() {
   process.env = { ...cleanEnv };
   process.env.PARAMETER_GITHUB_APP_ID_NAME = 'github-app-id';
+  process.env.PARAMETER_GITHUB_APP_KEY_BASE64_NAME = 'github-app-key';
   process.env.GITHUB_APP_KEY_BASE64 = 'TEST_CERTIFICATE_DATA';
   process.env.GITHUB_APP_ID = '1337';
   process.env.GITHUB_APP_CLIENT_ID = 'TEST_CLIENT_ID';
   process.env.GITHUB_APP_CLIENT_SECRET = 'TEST_CLIENT_SECRET';
   process.env.RUNNERS_MAXIMUM_COUNT = '3';
   process.env.ENVIRONMENT = EXPECTED_RUNNER_PARAMS.environment;
   process.env.SSM_TOKEN_PATH = '/github-action-runners/default/runners/config';
+  process.env.SSM_CONFIG_PATH = '/github-action-runners/default/runners/config';
 }
 
 async function createTestProviderRunners(input: CreateScaleUpRunnersInput<unknown>): Promise<CreateRunnerResult> {
@@ -330,7 +332,9 @@ describe('scaleUp with GHES', () => {
     it('returns a retryable failure if runner group lookup fails for ephemeral runners', async () => {
       process.env.RUNNER_GROUP_NAME = 'test-runner-group';
       mockSSMgetParameter.mockImplementation(async () => {
-        throw new Error('ParameterNotFound');
+        const error = new Error('ParameterNotFound');
+        error.name = 'ParameterNotFound';
+        throw error;
       });
 
       await expect(scaleUpModule.scaleUp(TEST_DATA)).resolves.toEqual(['foobar']);
@@ -348,7 +352,9 @@ describe('scaleUp with GHES', () => {
 
     it('create SSM parameter for runner group id if it does not exist', async () => {
       mockSSMgetParameter.mockImplementation(async () => {
-        throw new Error('ParameterNotFound');
+        const error = new Error('ParameterNotFound');
+        error.name = 'ParameterNotFound';
+        throw error;
       });
       await scaleUpModule.scaleUp(TEST_DATA);
       expect(mockOctokit.paginate).toHaveBeenCalledTimes(1);
@@ -1585,7 +1591,9 @@ describe('scaleUp with Github Data Residency', () => {
 
     it('create SSM parameter for runner group id if it does not exist', async () => {
       mockSSMgetParameter.mockImplementation(async () => {
-        throw new Error('ParameterNotFound');
+        const error = new Error('ParameterNotFound');
+        error.name = 'ParameterNotFound';
+        throw error;
       });
       await scaleUpModule.scaleUp(TEST_DATA);
       expect(mockOctokit.paginate).toHaveBeenCalledTimes(1);
@@ -2339,12 +2347,15 @@ function defaultOctokitMockImpl() {
 
 function defaultSSMGetParameterMockImpl() {
   mockSSMgetParameter.mockImplementation(async (name: string) => {
-    if (name === `${process.env.SSM_CONFIG_PATH}/runner-group/${process.env.RUNNER_GROUP_NAME}`) {
+    const runnerGroupName = process.env.RUNNER_GROUP_NAME || 'Default';
+    if (name === `${process.env.SSM_CONFIG_PATH}/runner-group/${runnerGroupName}`) {
       return '1';
     } else if (name === `${process.env.PARAMETER_GITHUB_APP_ID_NAME}`) {
       return `${process.env.GITHUB_APP_ID}`;
     } else {
-      throw new Error(`ParameterNotFound: ${name}`);
+      const error = new Error(`ParameterNotFound: ${name}`);
+      error.name = 'ParameterNotFound';
+      throw error;
     }
   });
 }

lambdas/libs/compute-providers/core/index.ts

@@ -31,6 +31,7 @@ export interface GitHubRunnerMetadata {
 }
 
 export interface StartRunnerConfigOptions {
+  runnerGroupCacheStore?: import('@aws-github-runner/storage-providers').RunnerGroupCacheStore;
   getRunnerConfigMetadata?: (runnerId: string) => { key: string; value: string }[];
   onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise<void>;
 }

lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.test.ts

@@ -0,0 +1,63 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util';
+
+import { createAwsSsmRunnerGroupCacheStore } from './runner-group-cache-store';
+
+vi.mock('@aws-github-runner/aws-ssm-util', () => ({
+  getParameter: vi.fn(),
+  putParameter: vi.fn(),
+}));
+
+const getParameterMock = vi.mocked(getParameter);
+const putParameterMock = vi.mocked(putParameter);
+
+describe('aws_ssm runner group cache store', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    process.env.SSM_CONFIG_PATH = '/runner/config';
+    delete process.env.SSM_PARAMETER_STORE_TAGS;
+  });
+
+  it('returns a cached numeric ID and preserves configured tags on create', async () => {
+    process.env.SSM_PARAMETER_STORE_TAGS = JSON.stringify([{ Key: 'Environment', Value: 'test' }]);
+    getParameterMock.mockResolvedValue('42');
+    const store = createAwsSsmRunnerGroupCacheStore();
+
+    await expect(store.get('Default')).resolves.toBe(42);
+    await store.create({ runnerGroupName: 'Default', runnerGroupId: 42 });
+
+    expect(getParameterMock).toHaveBeenCalledWith('/runner/config/runner-group/Default');
+    expect(putParameterMock).toHaveBeenCalledWith('/runner/config/runner-group/Default', '42', false, {
+      tags: [{ Key: 'Environment', Value: 'test' }],
+    });
+  });
+
+  it('returns undefined only for ParameterNotFound', async () => {
+    getParameterMock.mockRejectedValue(Object.assign(new Error('missing'), { name: 'ParameterNotFound' }));
+
+    await expect(createAwsSsmRunnerGroupCacheStore().get('Default')).resolves.toBeUndefined();
+  });
+
+  it('returns undefined when ParameterNotFound is wrapped by the SSM provider', async () => {
+    const cause = Object.assign(new Error('missing'), { name: 'ParameterNotFound' });
+    getParameterMock.mockRejectedValue(
+      Object.assign(new Error('failed to get parameter'), { name: 'GetParameterError', cause }),
+    );
+
+    await expect(createAwsSsmRunnerGroupCacheStore().get('Default')).resolves.toBeUndefined();
+  });
+
+  it('propagates access and service errors', async () => {
+    const error = Object.assign(new Error('denied'), { name: 'AccessDeniedException' });
+    getParameterMock.mockRejectedValue(error);
+
+    await expect(createAwsSsmRunnerGroupCacheStore().get('Default')).rejects.toBe(error);
+  });
+
+  it('rejects a non-numeric cached ID', async () => {
+    getParameterMock.mockResolvedValue('not-a-number');
+
+    await expect(createAwsSsmRunnerGroupCacheStore().get('Default')).rejects.toThrow(/cached runner group ID/i);
+  });
+});

lambdas/libs/storage-providers/aws/ssm/runner-group-cache-store.ts

@@ -0,0 +1,67 @@
+import { getParameter, putParameter } from '@aws-github-runner/aws-ssm-util';
+
+import type { RunnerGroupCacheRecord, RunnerGroupCacheStore } from '../../core';
+import type {} from './environment';
+import { loadSsmParameterStoreTagsFromEnvironment } from './parameter-store-tags';
+
+interface AwsSsmRunnerGroupCacheStoreConfig {
+  configPath: string;
+  parameterStoreTags: { Key: string; Value: string }[];
+}
+
+export function createAwsSsmRunnerGroupCacheStore(): RunnerGroupCacheStore {
+  const configPath = process.env.SSM_CONFIG_PATH;
+  if (!configPath || configPath.trim() === '') {
+    throw new Error('Environment variable SSM_CONFIG_PATH is not set');
+  }
+
+  return new AwsSsmRunnerGroupCacheStore({
+    configPath,
+    parameterStoreTags: loadSsmParameterStoreTagsFromEnvironment(),
+  });
+}
+
+class AwsSsmRunnerGroupCacheStore implements RunnerGroupCacheStore {
+  constructor(private readonly config: AwsSsmRunnerGroupCacheStoreConfig) {}
+
+  async get(runnerGroupName: string): Promise<number | undefined> {
+    try {
+      const value = await getParameter(this.parameterName(runnerGroupName));
+      const runnerGroupId = Number.parseInt(value, 10);
+      if (Number.isNaN(runnerGroupId)) {
+        throw new Error(`Cached runner group ID for ${runnerGroupName} is invalid`);
+      }
+      return runnerGroupId;
+    } catch (error) {
+      if (isParameterNotFoundError(error)) {
+        return undefined;
+      }
+      throw error;
+    }
+  }
+
+  async create(record: RunnerGroupCacheRecord): Promise<void> {
+    await putParameter(this.parameterName(record.runnerGroupName), record.runnerGroupId.toString(), false, {
+      tags: this.config.parameterStoreTags,
+    });
+  }
+
+  private parameterName(runnerGroupName: string): string {
+    return `${this.config.configPath}/runner-group/${runnerGroupName}`;
+  }
+}
+
+function isParameterNotFoundError(error: unknown): boolean {
+  const seen = new Set<object>();
+  let current: unknown = error;
+
+  while (current !== null && typeof current === 'object' && !seen.has(current)) {
+    seen.add(current);
+    if ('name' in current && current.name === 'ParameterNotFound') {
+      return true;
+    }
+    current = 'cause' in current ? current.cause : undefined;
+  }
+
+  return false;
+}

lambdas/libs/storage-providers/core/index.ts

@@ -12,3 +12,13 @@ export interface RunnerConfigStore {
   readonly maxWritesPerSecond?: number;
   create(record: RunnerConfigRecord, options?: { metadata?: RunnerConfigMetadata[] }): Promise<void>;
 }
+
+export interface RunnerGroupCacheRecord {
+  runnerGroupName: string;
+  runnerGroupId: number;
+}
+
+export interface RunnerGroupCacheStore {
+  get(runnerGroupName: string): Promise<number | undefined>;
+  create(record: RunnerGroupCacheRecord): Promise<void>;
+}

lambdas/libs/storage-providers/index.ts

@@ -1,2 +1,9 @@
-export type { RunnerConfigMetadata, RunnerConfigRecord, RunnerConfigStore } from './core';
+export type {
+  RunnerConfigMetadata,
+  RunnerConfigRecord,
+  RunnerConfigStore,
+  RunnerGroupCacheRecord,
+  RunnerGroupCacheStore,
+} from './core';
 export { createRunnerConfigStore } from './runner-config';
+export { createRunnerGroupCacheStore } from './runner-group-cache';

lambdas/libs/storage-providers/runner-group-cache.ts

@@ -0,0 +1,6 @@
+import { createAwsSsmRunnerGroupCacheStore } from './aws/ssm/runner-group-cache-store';
+import type { RunnerGroupCacheStore } from './core';
+
+export function createRunnerGroupCacheStore(): RunnerGroupCacheStore {
+  return createAwsSsmRunnerGroupCacheStore();
+}