From 9301fabc665516fafb5421efa57dfd920b92bcad Mon Sep 17 00:00:00 2001 From: edersonbrilhante Date: Fri, 4 Sep 2026 12:06:55 +0200 Subject: [PATCH] refactor(storage): separate runner config housekeeping --- .../control-plane/src/lambda.test.ts | 20 ++--- lambdas/functions/control-plane/src/lambda.ts | 13 +-- .../src/local-ssm-housekeeper.ts | 2 +- .../src/scale-runners/ssm-housekeeper.ts | 68 ++-------------- .../aws/ssm/runner-config-housekeeper.ts | 80 +++++++++++++++++++ lambdas/libs/storage-providers/core/index.ts | 4 + lambdas/libs/storage-providers/index.ts | 2 + lambdas/libs/storage-providers/package.json | 3 +- .../runner-config-housekeeper.ts | 6 ++ 9 files changed, 119 insertions(+), 79 deletions(-) create mode 100644 lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts create mode 100644 lambdas/libs/storage-providers/runner-config-housekeeper.ts diff --git a/lambdas/functions/control-plane/src/lambda.test.ts b/lambdas/functions/control-plane/src/lambda.test.ts index 4d970f0f22..b4540a9934 100644 --- a/lambdas/functions/control-plane/src/lambda.test.ts +++ b/lambdas/functions/control-plane/src/lambda.test.ts @@ -1,4 +1,5 @@ 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'; @@ -6,7 +7,6 @@ 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(); }); }); diff --git a/lambdas/functions/control-plane/src/lambda.ts b/lambdas/functions/control-plane/src/lambda.ts index d229a0350e..4594a1289e 100644 --- a/lambdas/functions/control-plane/src/lambda.ts +++ b/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 { @@ -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 { +export async function runnerConfigHousekeeper(event: unknown, context: Context): Promise { 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 { setContext(context, 'lambda.ts'); logger.logEventIfEnabled(event); diff --git a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts b/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts index ec635b13ad..81c4cbafd5 100644 --- a/lambdas/functions/control-plane/src/local-ssm-housekeeper.ts +++ b/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({ diff --git a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts b/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts index 857b974a9d..52464b720d 100644 --- a/lambdas/functions/control-plane/src/scale-runners/ssm-housekeeper.ts +++ b/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 { - 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'; diff --git a/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts b/lambdas/libs/storage-providers/aws/ssm/runner-config-housekeeper.ts new file mode 100644 index 0000000000..a7bfdccbb0 --- /dev/null +++ b/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 { + 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 { + 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(', ')); + } +} diff --git a/lambdas/libs/storage-providers/core/index.ts b/lambdas/libs/storage-providers/core/index.ts index 25d9f23f83..a408097ed7 100644 --- a/lambdas/libs/storage-providers/core/index.ts +++ b/lambdas/libs/storage-providers/core/index.ts @@ -13,6 +13,10 @@ export interface RunnerConfigStore { create(record: RunnerConfigRecord, options?: { metadata?: RunnerConfigMetadata[] }): Promise; } +export interface RunnerConfigHousekeeper { + houseKeeper(): Promise; +} + export interface RunnerGroupCacheRecord { runnerGroupName: string; runnerGroupId: number; diff --git a/lambdas/libs/storage-providers/index.ts b/lambdas/libs/storage-providers/index.ts index ed76bccc68..0e26830879 100644 --- a/lambdas/libs/storage-providers/index.ts +++ b/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'; diff --git a/lambdas/libs/storage-providers/package.json b/lambdas/libs/storage-providers/package.json index 65e93c2c08..0ba4f8715a 100644 --- a/lambdas/libs/storage-providers/package.json +++ b/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", diff --git a/lambdas/libs/storage-providers/runner-config-housekeeper.ts b/lambdas/libs/storage-providers/runner-config-housekeeper.ts new file mode 100644 index 0000000000..040cd1b801 --- /dev/null +++ b/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(); +}