feat: add lambda function to cleanup org runners

lambdas/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "jest.jestCommandLine": "yarn run test --",
+}

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

@@ -1,14 +1,23 @@
 import { captureLambdaHandler, logger } from '@aws-github-runner/aws-powertools-util';
 import { Context, SQSEvent, SQSRecord } from 'aws-lambda';
 
-import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, ssmHousekeeper, jobRetryCheck } from './lambda';
+import {
+  addMiddleware,
+  adjustPool,
+  scaleDownHandler,
+  scaleUpHandler,
+  ssmHousekeeper,
+  jobRetryCheck,
+  cleanupOrgRunnersHandler,
+} from './lambda';
 import { adjust } from './pool/pool';
 import ScaleError from './scale-runners/ScaleError';
 import { scaleDown } from './scale-runners/scale-down';
 import { ActionRequestMessage, scaleUp } from './scale-runners/scale-up';
 import { cleanSSMTokens } from './scale-runners/ssm-housekeeper';
 import { checkAndRetryJob } from './scale-runners/job-retry';
 import { describe, it, expect, vi, MockedFunction } from 'vitest';
+import { cleanupOrgRunners } from './scale-runners/cleanup-org-runners';
 
 const body: ActionRequestMessage = {
   eventType: 'workflow_job',
@@ -66,6 +75,7 @@ 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('./scale-runners/cleanup-org-runners');
 vi.mock('@aws-github-runner/aws-powertools-util');
 vi.mock('@aws-github-runner/aws-ssm-util');
 
@@ -226,3 +236,26 @@ describe('Test job retry check wrapper', () => {
     expect(logSpyWarn).toHaveBeenCalledWith(`Error processing job retry: ${error.message}`, { error });
   });
 });
+
+describe('Test cleanupOrgRunnersHandler lambda wrapper', () => {
+  it('Cleanup without error should resolve.', async () => {
+    const mock = mocked(cleanupOrgRunners);
+    mock.mockImplementation(() => {
+      return new Promise((resolve) => {
+        resolve();
+      });
+    });
+    await expect(cleanupOrgRunnersHandler({}, context)).resolves.not.toThrow();
+  });
+
+  it('Cleanup with error should resolve and log error.', async () => {
+    const logSpyError = jest.spyOn(logger, 'error');
+
+    const mock = mocked(cleanupOrgRunners);
+    const error = new Error('Error cleaning up org runners.');
+    mock.mockRejectedValue(error);
+
+    await expect(cleanupOrgRunnersHandler({}, context)).resolves.not.toThrow();
+    expect(logSpyError).toHaveBeenCalledWith(expect.stringContaining(error.message), expect.anything());
+  });
+});

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

@@ -9,6 +9,7 @@ import { scaleDown } from './scale-runners/scale-down';
 import { scaleUp } from './scale-runners/scale-up';
 import { SSMCleanupOptions, cleanSSMTokens } from './scale-runners/ssm-housekeeper';
 import { checkAndRetryJob } from './scale-runners/job-retry';
+import { cleanupOrgRunners } from './scale-runners/cleanup-org-runners';
 
 export async function scaleUpHandler(event: SQSEvent, context: Context): Promise<void> {
   setContext(context, 'lambda.ts');
@@ -64,6 +65,8 @@ export const addMiddleware = () => {
   middy(scaleDownHandler).use(handler);
   middy(adjustPool).use(handler);
   middy(ssmHousekeeper).use(handler);
+  middy(jobRetryCheck).use(handler);
+  middy(cleanupOrgRunnersHandler).use(handler);
 };
 addMiddleware();
 
@@ -91,3 +94,14 @@ export async function jobRetryCheck(event: SQSEvent, context: Context): Promise<
   }
   return Promise.resolve();
 }
+
+export async function cleanupOrgRunnersHandler(event: unknown, context: Context): Promise<void> {
+  setContext(context, 'lambda.ts');
+  logger.logEventIfEnabled(event);
+
+  try {
+    await cleanupOrgRunners();
+  } catch (e) {
+    logger.error(`${(e as Error).message}`, { error: e as Error });
+  }
+}