refactor(storage): separate runner config housekeeping

lambdas/functions/control-plane/src/lambda.test.ts

@@ -1,12 +1,12 @@
 import { captureLambdaHandler, logger } from '@aws-github-runner/aws-powertools-util';
+import { createRunnerConfigHousekeeper } from '@aws-github-runner/storage-providers';
 import { Context, SQSEvent, SQSRecord } from 'aws-lambda';
 
 import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, ssmHousekeeper, jobRetryCheck } from './lambda';
 import { adjust } from './pool/pool';
 import { scaleDown } from './scale-runners/scale-down';
 import { scaleUp } from './scale-runners/scale-up';
 import type { ActionRequestMessage } from './scale-runners/types';
-import { cleanSSMTokens } from './scale-runners/ssm-housekeeper';
 import { checkAndRetryJob } from './scale-runners/job-retry';
 import { describe, it, expect, vi, MockedFunction, beforeEach } from 'vitest';
 
@@ -64,10 +64,14 @@ const context: Context = {
 vi.mock('./pool/pool');
 vi.mock('./scale-runners/scale-down');
 vi.mock('./scale-runners/scale-up');
-vi.mock('./scale-runners/ssm-housekeeper');
 vi.mock('./scale-runners/job-retry');
 vi.mock('@aws-github-runner/aws-powertools-util');
 vi.mock('@aws-github-runner/aws-ssm-util');
+vi.mock('@aws-github-runner/storage-providers', () => ({
+  createRunnerConfigHousekeeper: vi.fn(),
+}));
+
+const mockedCreateRunnerConfigHousekeeper = vi.mocked(createRunnerConfigHousekeeper);
 
 describe('Test scale up lambda wrapper.', () => {
   it('Do not handle empty record sets.', async () => {
@@ -298,19 +302,15 @@ describe('Test middleware', () => {
 
 describe('Test ssm housekeeper lambda wrapper.', () => {
   it('Invoke without errors.', async () => {
-    vi.mocked(cleanSSMTokens).mockResolvedValue();
-
-    process.env.SSM_CLEANUP_CONFIG = JSON.stringify({
-      dryRun: false,
-      minimumDaysOld: 1,
-      tokenPath: '/path/to/tokens/',
-    });
+    const houseKeeper = vi.fn().mockResolvedValue();
+    mockedCreateRunnerConfigHousekeeper.mockReturnValue({ houseKeeper });
 
     await expect(ssmHousekeeper({}, context)).resolves.not.toThrow();
+    expect(houseKeeper).toHaveBeenCalledOnce();
   });
 
   it('Errors not throws.', async () => {
-    vi.mocked(cleanSSMTokens).mockRejectedValue(new Error());
+    mockedCreateRunnerConfigHousekeeper.mockReturnValue({ houseKeeper: vi.fn().mockRejectedValue(new Error()) });
     await expect(ssmHousekeeper({}, context)).resolves.not.toThrow();
   });
 });

lambdas/functions/control-plane/src/lambda.ts

@@ -1,13 +1,13 @@
 import middy from '@middy/core';
 import { logger, setContext } from '@aws-github-runner/aws-powertools-util';
 import { captureLambdaHandler, tracer } from '@aws-github-runner/aws-powertools-util';
+import { createRunnerConfigHousekeeper } from '@aws-github-runner/storage-providers';
 import { Context, type SQSBatchItemFailure, type SQSBatchResponse, SQSEvent } from 'aws-lambda';
 
 import { PoolEvent, adjust } from './pool/pool';
 import { scaleDown } from './scale-runners/scale-down';
 import { scaleUp } from './scale-runners/scale-up';
 import type { ActionRequestMessage, ActionRequestMessageSQS } from './scale-runners/types';
-import { SSMCleanupOptions, cleanSSMTokens } from './scale-runners/ssm-housekeeper';
 import { checkAndRetryJob } from './scale-runners/job-retry';
 
 export async function scaleUpHandler(event: SQSEvent, context: Context): Promise<SQSBatchResponse> {
@@ -114,22 +114,25 @@ export const addMiddleware = () => {
   middy(scaleUpHandler).use(handler);
   middy(scaleDownHandler).use(handler);
   middy(adjustPool).use(handler);
-  middy(ssmHousekeeper).use(handler);
+  middy(runnerConfigHousekeeper).use(handler);
 };
 addMiddleware();
 
-export async function ssmHousekeeper(event: unknown, context: Context): Promise<void> {
+export async function runnerConfigHousekeeper(event: unknown, context: Context): Promise<void> {
   setContext(context, 'lambda.ts');
   logger.logEventIfEnabled(event);
-  const config = JSON.parse(process.env.SSM_CLEANUP_CONFIG) as SSMCleanupOptions;
+  const housekeeper = createRunnerConfigHousekeeper();
 
   try {
-    await cleanSSMTokens(config);
+    await housekeeper.houseKeeper();
   } catch (e) {
     logger.error(`${(e as Error).message}`, { error: e as Error });
   }
 }
 
+/** @deprecated Use runnerConfigHousekeeper. Kept for existing Terraform handler configuration. */
+export const ssmHousekeeper = runnerConfigHousekeeper;
+
 export async function jobRetryCheck(event: SQSEvent, context: Context): Promise<void> {
   setContext(context, 'lambda.ts');
   logger.logEventIfEnabled(event);

lambdas/functions/control-plane/src/local-ssm-housekeeper.ts

@@ -1,4 +1,4 @@
-import { cleanSSMTokens } from './scale-runners/ssm-housekeeper';
+import { cleanSSMTokens } from '@aws-github-runner/storage-providers/aws/ssm/runner-config-housekeeper';
 
 export function run(): void {
   cleanSSMTokens({

lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts

@@ -1,62 +1,6 @@
-import { DeleteParameterCommand, GetParametersByPathCommand, SSMClient } from '@aws-sdk/client-ssm';
-import { logger } from '@aws-github-runner/aws-powertools-util';
-import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util';
-
-export interface SSMCleanupOptions {
-  dryRun: boolean;
-  minimumDaysOld: number;
-  tokenPath: string;
-}
-
-function validateOptions(options: SSMCleanupOptions): void {
-  const errorMessages: string[] = [];
-  if (!options.minimumDaysOld || options.minimumDaysOld < 1) {
-    errorMessages.push(`minimumDaysOld must be greater then 0, value is set to "${options.minimumDaysOld}"`);
-  }
-  if (!options.tokenPath) {
-    errorMessages.push('tokenPath must be defined');
-  }
-  if (errorMessages.length > 0) {
-    throw new Error(errorMessages.join(', '));
-  }
-}
-
-export async function cleanSSMTokens(options: SSMCleanupOptions): Promise<void> {
-  logger.info(`Cleaning tokens / JIT config older then ${options.minimumDaysOld} days, dryRun: ${options.dryRun}`);
-  logger.debug('Cleaning with options', { options });
-  validateOptions(options);
-
-  const client = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION }));
-  const parameters = await client.send(new GetParametersByPathCommand({ Path: options.tokenPath }));
-  while (parameters.NextToken) {
-    const nextParameters = await client.send(
-      new GetParametersByPathCommand({ Path: options.tokenPath, NextToken: parameters.NextToken }),
-    );
-    parameters.Parameters?.push(...(nextParameters.Parameters ?? []));
-    parameters.NextToken = nextParameters.NextToken;
-  }
-  logger.info(`Found #${parameters.Parameters?.length} parameters in path ${options.tokenPath}`);
-  logger.debug('Found parameters', { parameters });
-
-  // minimumDate = today - minimumDaysOld
-  const minimumDate = new Date();
-  minimumDate.setDate(minimumDate.getDate() - options.minimumDaysOld);
-
-  for (const parameter of parameters.Parameters ?? []) {
-    if (parameter.LastModifiedDate && new Date(parameter.LastModifiedDate) < minimumDate) {
-      logger.info(`Deleting parameter ${parameter.Name} with last modified date ${parameter.LastModifiedDate}`);
-      try {
-        if (!options.dryRun) {
-          // sleep 50ms to avoid rait limit
-          await new Promise((resolve) => setTimeout(resolve, 50));
-          await client.send(new DeleteParameterCommand({ Name: parameter.Name }));
-        }
-      } catch (e) {
-        logger.warn(`Failed to delete parameter ${parameter.Name} with error ${(e as Error).message}`);
-        logger.debug('Failed to delete parameter', { e });
-      }
-    } else {
-      logger.debug(`Skipping parameter ${parameter.Name} with last modified date ${parameter.LastModifiedDate}`);
-    }
-  }
-}
+/** @deprecated Import the SSM runner-config housekeeper from storage-providers. */
+export {
+  cleanSSMTokens,
+  createAwsSsmRunnerConfigHousekeeper,
+  type SSMCleanupOptions,
+} from '@aws-github-runner/storage-providers/aws/ssm/runner-config-housekeeper';

lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts

@@ -0,0 +1,80 @@
+import { DeleteParameterCommand, GetParametersByPathCommand, SSMClient } from '@aws-sdk/client-ssm';
+import { getTracedAWSV3Client, logger } from '@aws-github-runner/aws-powertools-util';
+
+import type { RunnerConfigHousekeeper } from '../../core';
+
+export interface SSMCleanupOptions {
+  dryRun: boolean;
+  minimumDaysOld: number;
+  tokenPath: string;
+}
+
+export function createAwsSsmRunnerConfigHousekeeper(options?: SSMCleanupOptions): RunnerConfigHousekeeper {
+  return new AwsSsmRunnerConfigHousekeeper(options ?? loadCleanupOptions());
+}
+
+export async function cleanSSMTokens(options: SSMCleanupOptions): Promise<void> {
+  logger.info(`Cleaning tokens / JIT config older then ${options.minimumDaysOld} days, dryRun: ${options.dryRun}`);
+  logger.debug('Cleaning with options', { options });
+  validateOptions(options);
+
+  const client = getTracedAWSV3Client(new SSMClient({ region: process.env.AWS_REGION }));
+  const parameters = await client.send(new GetParametersByPathCommand({ Path: options.tokenPath }));
+  while (parameters.NextToken) {
+    const nextParameters = await client.send(
+      new GetParametersByPathCommand({ Path: options.tokenPath, NextToken: parameters.NextToken }),
+    );
+    parameters.Parameters?.push(...(nextParameters.Parameters ?? []));
+    parameters.NextToken = nextParameters.NextToken;
+  }
+  logger.info(`Found #${parameters.Parameters?.length} parameters in path ${options.tokenPath}`);
+
+  const minimumDate = new Date();
+  minimumDate.setDate(minimumDate.getDate() - options.minimumDaysOld);
+
+  for (const parameter of parameters.Parameters ?? []) {
+    if (parameter.LastModifiedDate && new Date(parameter.LastModifiedDate) < minimumDate) {
+      logger.info(`Deleting parameter ${parameter.Name} with last modified date ${parameter.LastModifiedDate}`);
+      try {
+        if (!options.dryRun) {
+          await new Promise((resolve) => setTimeout(resolve, 50));
+          await client.send(new DeleteParameterCommand({ Name: parameter.Name }));
+        }
+      } catch (error) {
+        logger.warn(`Failed to delete parameter ${parameter.Name} with error ${(error as Error).message}`);
+        logger.debug('Failed to delete parameter', { error });
+      }
+    } else {
+      logger.debug(`Skipping parameter ${parameter.Name} with last modified date ${parameter.LastModifiedDate}`);
+    }
+  }
+}
+
+class AwsSsmRunnerConfigHousekeeper implements RunnerConfigHousekeeper {
+  constructor(private readonly options: SSMCleanupOptions) {}
+
+  houseKeeper(): Promise<void> {
+    return cleanSSMTokens(this.options);
+  }
+}
+
+function loadCleanupOptions(): SSMCleanupOptions {
+  const value = process.env.SSM_CLEANUP_CONFIG;
+  if (!value || value.trim() === '') {
+    throw new Error('Environment variable SSM_CLEANUP_CONFIG is not set');
+  }
+  return JSON.parse(value) as SSMCleanupOptions;
+}
+
+function validateOptions(options: SSMCleanupOptions): void {
+  const errorMessages: string[] = [];
+  if (!options.minimumDaysOld || options.minimumDaysOld < 1) {
+    errorMessages.push(`minimumDaysOld must be greater then 0, value is set to "${options.minimumDaysOld}"`);
+  }
+  if (!options.tokenPath) {
+    errorMessages.push('tokenPath must be defined');
+  }
+  if (errorMessages.length > 0) {
+    throw new Error(errorMessages.join(', '));
+  }
+}

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

@@ -13,6 +13,10 @@ export interface RunnerConfigStore {
   create(record: RunnerConfigRecord, options?: { metadata?: RunnerConfigMetadata[] }): Promise<void>;
 }
 
+export interface RunnerConfigHousekeeper {
+  houseKeeper(): Promise<void>;
+}
+
 export interface RunnerGroupCacheRecord {
   runnerGroupName: string;
   runnerGroupId: number;

lambdas/libs/storage-providers/index.ts

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

lambdas/libs/storage-providers/package.json

@@ -3,7 +3,8 @@
   "version": "1.0.0",
   "main": "index.ts",
   "exports": {
-    ".": "./index.ts"
+    ".": "./index.ts",
+    "./aws/ssm/runner-config-housekeeper": "./aws/ssm/runner-config-housekeeper.ts"
   },
   "type": "module",
   "license": "MIT",

lambdas/libs/storage-providers/runner-config-housekeeper.ts

@@ -0,0 +1,6 @@
+import { createAwsSsmRunnerConfigHousekeeper } from './aws/ssm/runner-config-housekeeper';
+import type { RunnerConfigHousekeeper } from './core';
+
+export function createRunnerConfigHousekeeper(): RunnerConfigHousekeeper {
+  return createAwsSsmRunnerConfigHousekeeper();
+}