refactor(compute-providers): reuse EC2 runner clients

lambdas/libs/compute-providers/aws/ec2/control-plane.ts

@@ -1,20 +1,26 @@
 import type { CreateStartRunnerConfig, ComputeProviderPlugin } from '../../core';
+import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util';
+import { EC2Client } from '@aws-sdk/client-ec2';
 
 import type { ControlPlaneProviderCapabilities, ControlPlaneProviderModule } from '../../contracts';
 import type {} from './src/environment';
-import { createEc2PoolProvider } from './src/control-plane/pool';
-import { createEc2ScaleDownProvider } from './src/control-plane/scale-down';
-import { createEc2ScaleUpProvider } from './src/control-plane/scale-up';
+import { createEc2PoolCapability } from './src/control-plane/pool';
+import { createEc2ScaleDownCapability } from './src/control-plane/scale-down';
+import { createEc2ScaleUpCapability } from './src/control-plane/scale-up';
+import { createEc2RunnerClient } from './src/runners';
 
 export function createEc2ControlPlanePlugin(
   createStartRunnerConfig: CreateStartRunnerConfig,
 ): ComputeProviderPlugin<ControlPlaneProviderCapabilities, 'ec2'> {
+  const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION }));
+  const ec2Operations = createEc2RunnerClient(ec2Client).forRequest({ signal: undefined });
+
   return {
     type: 'ec2',
     capabilities: {
-      pool: () => createEc2PoolProvider(createStartRunnerConfig),
-      scaleUp: () => createEc2ScaleUpProvider(createStartRunnerConfig),
-      scaleDown: createEc2ScaleDownProvider,
+      pool: () => createEc2PoolCapability(ec2Operations, createStartRunnerConfig),
+      scaleUp: () => createEc2ScaleUpCapability(ec2Operations, createStartRunnerConfig),
+      scaleDown: () => createEc2ScaleDownCapability(ec2Operations),
     },
   };
 }

lambdas/libs/compute-providers/aws/ec2/src/control-plane/dynamic-labels.ts

@@ -11,8 +11,6 @@ import {
   BurstablePerformance,
   CpuManufacturer,
   CpuPerformanceFactorRequest,
-  DescribeLaunchTemplateVersionsCommand,
-  EC2Client,
   FleetBlockDeviceMappingRequest,
   FleetEbsBlockDeviceRequest,
   InstanceGeneration,
@@ -30,7 +28,6 @@ import {
   VCpuCountRangeRequest,
   VolumeType,
 } from '@aws-sdk/client-ec2';
-import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util';
 
 import { Ec2OverrideConfig } from '../runners.d';
 
@@ -365,24 +362,3 @@ export function shouldLoadLaunchTemplateBlockDeviceName(labels: string[]): boole
 
   return hasBlockDeviceOverride && !hasBlockDeviceName;
 }
-
-export async function getDefaultBlockDeviceNameFromLaunchTemplate(launchTemplateName: string): Promise<string> {
-  const ec2Client = getTracedAWSV3Client(new EC2Client({ region: process.env.AWS_REGION }));
-  const launchTemplateVersions = await ec2Client.send(
-    new DescribeLaunchTemplateVersionsCommand({
-      LaunchTemplateName: launchTemplateName,
-      Versions: ['$Default'],
-    }),
-  );
-  const blockDeviceMappings =
-    launchTemplateVersions.LaunchTemplateVersions?.[0]?.LaunchTemplateData?.BlockDeviceMappings;
-  const blockDeviceName =
-    blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName && blockDeviceMapping.Ebs)
-      ?.DeviceName ?? blockDeviceMappings?.find((blockDeviceMapping) => blockDeviceMapping.DeviceName)?.DeviceName;
-
-  if (!blockDeviceName) {
-    throw new Error(`Failed to determine block device name from launch template '${launchTemplateName}'.`);
-  }
-
-  return blockDeviceName;
-}

lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.test.ts

@@ -1,26 +1,34 @@
 import type { Octokit } from '@octokit/rest';
 import type { CreateGitHubRunnerConfig, CreateStartRunnerConfig, RunnerInfo } from '../../../../core';
-import { bootTimeExceeded, listEC2Runners } from '../runners';
-import { calculateEc2PoolSize, createEc2PoolProvider } from './pool';
-import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-config';
+import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners';
+import { createEc2PoolCapability } from './pool';
+import { createRunners, type Ec2ProviderConfig, loadEc2ProviderConfig } from './runner-creation';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 vi.mock('../runners', () => ({
   bootTimeExceeded: vi.fn(),
-  listEC2Runners: vi.fn(),
 }));
 
-vi.mock('./runner-config', () => ({
+vi.mock('./runner-creation', () => ({
   createRunners: vi.fn(),
   loadEc2ProviderConfig: vi.fn(),
 }));
 
 const mockBootTimeExceeded = vi.mocked(bootTimeExceeded);
-const mockListRunners = vi.mocked(listEC2Runners);
 const mockCreateRunners = vi.mocked(createRunners);
 const mockLoadProviderConfig = vi.mocked(loadEc2ProviderConfig);
 
-describe('calculateEc2PoolSize', () => {
+const ec2Operations = {
+  list: vi.fn<Ec2RunnerResourceOperations['list']>(),
+  create: vi.fn<Ec2RunnerResourceOperations['create']>(),
+  terminate: vi.fn<Ec2RunnerResourceOperations['terminate']>(),
+  tag: vi.fn<Ec2RunnerResourceOperations['tag']>(),
+  untag: vi.fn<Ec2RunnerResourceOperations['untag']>(),
+} satisfies Ec2RunnerResourceOperations;
+const createStartRunnerConfig = vi.fn<CreateStartRunnerConfig>();
+const capability = createEc2PoolCapability(ec2Operations, createStartRunnerConfig);
+
+describe('createEc2PoolCapability.countAvailableRunners', () => {
   beforeEach(() => {
     vi.clearAllMocks();
   });
@@ -29,7 +37,7 @@ describe('calculateEc2PoolSize', () => {
     const runners: RunnerInfo[] = [{ id: 'i-idle', owner: 'owner', type: 'Org' }];
     const runnerStatus = new Map([['i-idle', { busy: false, status: 'online' }]]);
 
-    expect(calculateEc2PoolSize(runners, runnerStatus)).toBe(1);
+    expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(1);
     expect(mockBootTimeExceeded).not.toHaveBeenCalled();
   });
 
@@ -43,35 +51,59 @@ describe('calculateEc2PoolSize', () => {
       ['i-offline', { busy: false, status: 'offline' }],
     ]);
 
-    expect(calculateEc2PoolSize(runners, runnerStatus)).toBe(0);
+    expect(capability.countAvailableRunners(runners, runnerStatus)).toBe(0);
     expect(mockBootTimeExceeded).not.toHaveBeenCalled();
   });
 
   it('counts registered busy runners when busy runners are included', () => {
     const runners: RunnerInfo[] = [{ id: 'i-busy', owner: 'owner', type: 'Org' }];
     const runnerStatus = new Map([['i-busy', { busy: true, status: 'online' }]]);
 
-    expect(calculateEc2PoolSize(runners, runnerStatus, true)).toBe(1);
+    expect(capability.countAvailableRunners(runners, runnerStatus, true)).toBe(1);
     expect(mockBootTimeExceeded).not.toHaveBeenCalled();
   });
 
   it('counts unregistered runners that are still booting', () => {
     const runners: RunnerInfo[] = [{ id: 'i-booting', owner: 'owner', type: 'Org' }];
     mockBootTimeExceeded.mockReturnValue(false);
 
-    expect(calculateEc2PoolSize(runners, new Map())).toBe(1);
+    expect(capability.countAvailableRunners(runners, new Map())).toBe(1);
   });
 
   it('does not count unregistered runners whose boot time expired', () => {
     const runners: RunnerInfo[] = [{ id: 'i-expired', owner: 'owner', type: 'Org' }];
     mockBootTimeExceeded.mockReturnValue(true);
 
-    expect(calculateEc2PoolSize(runners, new Map())).toBe(0);
+    expect(capability.countAvailableRunners(runners, new Map())).toBe(0);
   });
 });
 
-describe('createEc2PoolProvider', () => {
-  const createStartRunnerConfig = vi.fn<CreateStartRunnerConfig>();
+describe('createEc2PoolCapability.listRunners', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('lists only running instances managed for the requested pool', async () => {
+    const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }];
+    ec2Operations.list.mockResolvedValue(runners);
+
+    await expect(
+      capability.listRunners({
+        environment: 'test-environment',
+        runnerOwner: 'owner',
+        runnerType: 'Org',
+      }),
+    ).resolves.toBe(runners);
+    expect(ec2Operations.list).toHaveBeenCalledWith({
+      environment: 'test-environment',
+      runnerOwner: 'owner',
+      runnerType: 'Org',
+      statuses: ['running'],
+    });
+  });
+});
+
+describe('createEc2PoolCapability.createRunners', () => {
   const githubInstallationClient = {} as Octokit;
   const githubRunnerConfig: CreateGitHubRunnerConfig = {
     ephemeral: true,
@@ -105,42 +137,22 @@ describe('createEc2PoolProvider', () => {
     mockLoadProviderConfig.mockReturnValue(providerConfig);
   });
 
-  it('lists only running instances managed for the requested pool', async () => {
-    const runners: RunnerInfo[] = [{ id: 'i-running', owner: 'owner', type: 'Org' }];
-    mockListRunners.mockResolvedValue(runners);
-    const provider = createEc2PoolProvider(createStartRunnerConfig);
-
-    await expect(
-      provider.listRunners({
-        environment: 'test-environment',
-        runnerOwner: 'owner',
-        runnerType: 'Org',
-      }),
-    ).resolves.toBe(runners);
-    expect(mockListRunners).toHaveBeenCalledWith({
-      environment: 'test-environment',
-      runnerOwner: 'owner',
-      runnerType: 'Org',
-      statuses: ['running'],
-    });
-  });
-
   it('creates pool runners with the pool source and returns their instance IDs', async () => {
     mockCreateRunners.mockResolvedValue({
       instances: ['i-created'],
       retryableErrorCount: 0,
       nonRetryableErrorCount: 0,
     });
-    const provider = createEc2PoolProvider(createStartRunnerConfig);
 
     await expect(
-      provider.createRunners({
+      capability.createRunners({
         githubRunnerConfig,
         numberOfRunners: 1,
         githubInstallationClient,
       }),
     ).resolves.toEqual(['i-created']);
     expect(mockCreateRunners).toHaveBeenCalledWith(
+      ec2Operations,
       githubRunnerConfig,
       providerConfig,
       1,

lambdas/libs/compute-providers/aws/ec2/src/control-plane/pool.ts

@@ -1,67 +1,11 @@
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
-import type {
-  CreateStartRunnerConfig,
-  CreatePoolRunnersInput,
-  ListPoolRunnersInput,
-  PoolComputeProvider,
-  RunnerInfo,
-  RunnerStatus,
-} from '../../../../core';
-import { createRunners, loadEc2ProviderConfig } from './runner-config';
-import { bootTimeExceeded, listEC2Runners } from '../runners';
+import type { CreateStartRunnerConfig, PoolComputeProvider, RunnerInfo, RunnerStatus } from '../../../../core';
+import { bootTimeExceeded, type Ec2RunnerResourceOperations } from '../runners';
+import { createRunners, loadEc2ProviderConfig } from './runner-creation';
 
 const logger = createChildLogger('pool');
 
-async function listEc2PoolRunners({
-  environment,
-  runnerOwner,
-  runnerType,
-}: ListPoolRunnersInput): Promise<RunnerInfo[]> {
-  return await listEC2Runners({
-    environment,
-    runnerOwner,
-    runnerType,
-    statuses: ['running'],
-  });
-}
-
-async function createEc2PoolRunners(
-  { githubRunnerConfig, numberOfRunners, githubInstallationClient }: CreatePoolRunnersInput,
-  createStartRunnerConfig: CreateStartRunnerConfig,
-): Promise<string[]> {
-  const config = loadEc2ProviderConfig();
-
-  const { instances } = await createRunners(
-    githubRunnerConfig,
-    {
-      ec2instanceCriteria: config.ec2instanceCriteria,
-      environment: config.environment,
-      launchTemplateName: config.launchTemplateName,
-      subnets: config.subnets,
-      amiIdSsmParameterName: config.amiIdSsmParameterName,
-      tracingEnabled: config.tracingEnabled,
-      onDemandFailoverOnError: config.onDemandFailoverOnError,
-      scaleErrors: config.scaleErrors,
-    },
-    numberOfRunners,
-    githubInstallationClient,
-    createStartRunnerConfig,
-    'pool-lambda',
-  );
-  return instances;
-}
-
-export function createEc2PoolProvider(
-  createStartRunnerConfig: CreateStartRunnerConfig,
-): Omit<PoolComputeProvider<RunnerInfo>, 'type'> {
-  return {
-    listRunners: listEc2PoolRunners,
-    countAvailableRunners: calculateEc2PoolSize,
-    createRunners: (input) => createEc2PoolRunners(input, createStartRunnerConfig),
-  };
-}
-
-export function calculateEc2PoolSize(
+function countAvailableEc2PoolRunners(
   ec2runners: RunnerInfo[],
   runnerStatus: Map<string, RunnerStatus>,
   includeBusyRunners = false,
@@ -86,3 +30,42 @@ export function calculateEc2PoolSize(
   }
   return numberOfRunnersInPool;
 }
+
+export function createEc2PoolCapability(
+  ec2Operations: Ec2RunnerResourceOperations,
+  createStartRunnerConfig: CreateStartRunnerConfig,
+): Omit<PoolComputeProvider<RunnerInfo>, 'type'> {
+  return {
+    listRunners: ({ environment, runnerOwner, runnerType }) =>
+      ec2Operations.list({
+        environment,
+        runnerOwner,
+        runnerType,
+        statuses: ['running'],
+      }),
+    countAvailableRunners: countAvailableEc2PoolRunners,
+    createRunners: async ({ githubRunnerConfig, numberOfRunners, githubInstallationClient }) => {
+      const config = loadEc2ProviderConfig();
+
+      const { instances } = await createRunners(
+        ec2Operations,
+        githubRunnerConfig,
+        {
+          ec2instanceCriteria: config.ec2instanceCriteria,
+          environment: config.environment,
+          launchTemplateName: config.launchTemplateName,
+          subnets: config.subnets,
+          amiIdSsmParameterName: config.amiIdSsmParameterName,
+          tracingEnabled: config.tracingEnabled,
+          onDemandFailoverOnError: config.onDemandFailoverOnError,
+          scaleErrors: config.scaleErrors,
+        },
+        numberOfRunners,
+        githubInstallationClient,
+        createStartRunnerConfig,
+        'pool-lambda',
+      );
+      return instances;
+    },
+  };
+}