refactor(runners): introduce provider plugin framework

lambdas/functions/control-plane/package.json

@@ -32,7 +32,7 @@
   "dependencies": {
     "@aws-github-runner/aws-powertools-util": "*",
     "@aws-github-runner/aws-ssm-util": "*",
-    "@aws-github-runner/runner-provider": "*",
+    "@aws-github-runner/runner-providers": "*",
     "@aws-lambda-powertools/parameters": "^2.31.0",
     "@aws-sdk/client-ec2": "^3.1009.0",
     "@aws-sdk/client-sqs": "^3.1009.0",

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

@@ -0,0 +1,8 @@
+import { createControlPlaneProviderRegistry } from '@aws-github-runner/runner-providers/control-plane';
+import { runnerProviderTypes } from '@aws-github-runner/runner-providers/provider-types';
+
+import { createStartRunnerConfig } from './scale-runners/github-runner';
+
+export const controlPlaneProviderRegistry = createControlPlaneProviderRegistry(createStartRunnerConfig);
+
+export { runnerProviderTypes };

lambdas/functions/control-plane/src/pool/pool-contract.test.ts

@@ -0,0 +1,74 @@
+import type { Octokit } from '@octokit/rest';
+import type { RunnerProviderType } from '@aws-github-runner/runner-providers/provider-types';
+import { beforeEach, vi } from 'vitest';
+
+import { definePoolContractTests } from '../test/runner-provider-contracts/pool';
+import { providerTypes } from '../test/runner-provider-contracts/provider-types';
+import * as ghAuth from '../github/auth';
+import { controlPlaneProviderRegistry } from '../control-plane-providers';
+import * as githubRunner from '../scale-runners/github-runner';
+import { adjust } from './pool';
+import type { PoolRunnerProvider } from './pool-provider';
+
+vi.mock('../github/auth', () => ({
+  createGithubAppAuth: vi.fn(),
+  createGithubInstallationAuth: vi.fn(),
+  createOctokitClient: vi.fn(),
+}));
+
+vi.mock('../scale-runners/github-runner', () => ({
+  createStartRunnerConfig: vi.fn(),
+  getGitHubEnterpriseApiUrl: vi.fn(),
+  validateSsmParameterStoreTags: vi.fn(),
+}));
+
+const mockedAppAuth = vi.mocked(ghAuth.createGithubAppAuth);
+const mockedInstallationAuth = vi.mocked(ghAuth.createGithubInstallationAuth);
+const mockedCreateClient = vi.mocked(ghAuth.createOctokitClient);
+const mockedResolveCapability = vi.spyOn(controlPlaneProviderRegistry, 'capability');
+
+const githubClient = {
+  actions: { listSelfHostedRunnersForOrg: vi.fn() },
+  apps: { getOrgInstallation: vi.fn() },
+  paginate: vi.fn(),
+} as unknown as Octokit;
+
+const cleanEnv = process.env;
+
+const lanes = providerTypes.map((type) => ({
+  provider: {
+    type,
+    listRunners: vi.fn(),
+    countAvailableRunners: vi.fn(),
+    createRunners: vi.fn(),
+  } satisfies PoolRunnerProvider,
+}));
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  process.env = { ...cleanEnv };
+
+  mockedAppAuth.mockResolvedValue({ type: 'app', token: 'app-token', appId: 1, expiresAt: 'some-date' });
+  mockedInstallationAuth.mockResolvedValue({
+    type: 'token',
+    tokenType: 'installation',
+    token: 'installation-token',
+    createdAt: 'some-date',
+    expiresAt: 'some-date',
+    permissions: {},
+    repositorySelection: 'selected',
+    installationId: 2,
+  });
+  mockedCreateClient.mockResolvedValue(githubClient);
+  vi.mocked(githubRunner.getGitHubEnterpriseApiUrl).mockReturnValue({ ghesApiUrl: '', ghesBaseUrl: '' });
+  vi.mocked(githubRunner.validateSsmParameterStoreTags).mockReturnValue([]);
+  vi.mocked(githubClient.apps.getOrgInstallation).mockResolvedValue({ data: { id: 2 } } as never);
+  vi.mocked(githubClient.paginate).mockResolvedValue([]);
+});
+
+definePoolContractTests<RunnerProviderType>({
+  adjust,
+  githubInstallationClient: githubClient,
+  lanes,
+  resolveCapability: mockedResolveCapability,
+});

lambdas/functions/control-plane/src/pool/pool-provider.ts

@@ -1,31 +1,6 @@
-import type { Octokit } from '@octokit/rest';
-import type { RunnerProvider } from '@aws-github-runner/runner-provider';
-
-import type { CreateGitHubRunnerConfig, GitHubRunnerType } from '../scale-runners/types';
-
-export interface RunnerStatus {
-  busy: boolean;
-  status: string;
-}
-
-export interface ListPoolRunnersInput {
-  environment: string;
-  runnerOwner: string;
-  runnerType: GitHubRunnerType;
-}
-
-export interface CreatePoolRunnersInput {
-  githubRunnerConfig: CreateGitHubRunnerConfig;
-  numberOfRunners: number;
-  githubInstallationClient: Octokit;
-}
-
-export interface PoolRunnerProvider<TRunner = unknown> extends RunnerProvider {
-  listRunners(input: ListPoolRunnersInput): Promise<TRunner[]>;
-  countAvailableRunners(
-    runners: TRunner[],
-    runnerStatus: Map<string, RunnerStatus>,
-    includeBusyRunners: boolean,
-  ): number;
-  createRunners(input: CreatePoolRunnersInput): Promise<string[]>;
-}
+export type {
+  CreatePoolRunnersInput,
+  ListPoolRunnersInput,
+  PoolRunnerProvider,
+  RunnerStatus,
+} from '@aws-github-runner/runner-providers/core';

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

@@ -2,9 +2,9 @@ import { Octokit } from '@octokit/rest';
 import moment from 'moment-timezone';
 import * as nock from 'nock';
 
-import { listEC2Runners } from '../aws/ec2-runners';
+import { createRunners } from '@aws-github-runner/runner-providers/aws/ec2/control-plane/runner-config';
+import { listEC2Runners } from '@aws-github-runner/runner-providers/aws/ec2/control-plane/runners';
 import * as ghAuth from '../github/auth';
-import { createRunners } from '../scale-runners/ec2';
 import { getGitHubEnterpriseApiUrl } from '../scale-runners/github-runner';
 import { adjust } from './pool';
 import { describe, it, expect, beforeEach, vi, MockedClass } from 'vitest';
@@ -26,7 +26,7 @@ vi.mock('@octokit/rest', () => ({
   }),
 }));
 
-vi.mock('./../aws/ec2-runners', async () => ({
+vi.mock('@aws-github-runner/runner-providers/aws/ec2/control-plane/runners', async () => ({
   listEC2Runners: vi.fn(),
   // Include any other functions from the module that might be used
   bootTimeExceeded: vi.fn(),
@@ -37,12 +37,13 @@ vi.mock('./../github/auth', async () => ({
   createOctokitClient: vi.fn(),
 }));
 
-vi.mock('../scale-runners/ec2', async (importOriginal) => ({
-  ...(await importOriginal<typeof import('../scale-runners/ec2')>()),
+vi.mock('@aws-github-runner/runner-providers/aws/ec2/control-plane/runner-config', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('@aws-github-runner/runner-providers/aws/ec2/control-plane/runner-config')>()),
   createRunners: vi.fn(),
 }));
 
 vi.mock('../scale-runners/github-runner', async () => ({
+  createStartRunnerConfig: vi.fn(),
   getGitHubEnterpriseApiUrl: vi.fn().mockReturnValue({
     ghesApiUrl: '',
     ghesBaseUrl: '',
@@ -206,6 +207,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         1,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -223,6 +225,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         8,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -234,6 +237,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         8,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -323,6 +327,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         3,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -344,6 +349,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         3,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -400,6 +406,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         1,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -437,6 +444,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         2,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -451,6 +459,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         1,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -464,6 +473,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         8,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });
@@ -488,6 +498,7 @@ describe('Test simple pool.', () => {
         expect.anything(),
         2,
         expect.anything(),
+        expect.anything(),
         'pool-lambda',
       );
     });

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

@@ -1,10 +1,10 @@
 import { Octokit } from '@octokit/rest';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
-import { resolveRunnerProviderType } from '@aws-github-runner/runner-provider';
+import { resolveRunnerProviderType } from '@aws-github-runner/runner-providers/provider-types';
 import yn from 'yn';
 
 import { createGithubAppAuth, createGithubInstallationAuth, createOctokitClient } from '../github/auth';
-import { createPoolRunnerProvider } from '../runner-provider-registry';
+import { controlPlaneProviderRegistry } from '../control-plane-providers';
 import { getGitHubEnterpriseApiUrl, validateSsmParameterStoreTags } from '../scale-runners/github-runner';
 import type { RunnerStatus } from './pool-provider';
 
@@ -17,7 +17,10 @@ export interface PoolEvent {
 
 export async function adjust(event: PoolEvent): Promise<void> {
   const runnerProviderType = resolveRunnerProviderType(event.type);
-  const runnerProvider = createPoolRunnerProvider(runnerProviderType);
+  const runnerProvider = {
+    ...controlPlaneProviderRegistry.capability(runnerProviderType, 'pool')(),
+    type: runnerProviderType,
+  };
   logger.info(`Checking current ${runnerProvider.type} pool size against pool of size: ${event.poolSize}`);
   const runnerLabels = process.env.RUNNER_LABELS || '';
   const runnerGroup = process.env.RUNNER_GROUP_NAME || '';

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

@@ -0,0 +1,34 @@
+import type { RunnerProviderType } from '@aws-github-runner/runner-providers/provider-types';
+import { beforeEach, vi } from 'vitest';
+
+import { providerTypes } from '../test/runner-provider-contracts/provider-types';
+import { defineScaleDownContractTests } from '../test/runner-provider-contracts/scale-down';
+import { controlPlaneProviderRegistry } from '../control-plane-providers';
+import { scaleDown } from './scale-down';
+import type { ScaleDownRunnerProvider } from './scale-down-provider';
+
+const mockedResolveCapability = vi.spyOn(controlPlaneProviderRegistry, 'capability');
+
+const cleanEnv = process.env;
+
+const lanes = providerTypes.map((type) => ({
+  provider: {
+    type,
+    list: vi.fn(),
+    bootTimeExceeded: vi.fn(),
+    markOrphan: vi.fn(),
+    unmarkOrphan: vi.fn(),
+    terminate: vi.fn(),
+  } satisfies ScaleDownRunnerProvider,
+}));
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  process.env = { ...cleanEnv };
+});
+
+defineScaleDownContractTests<RunnerProviderType>({
+  lanes,
+  resolveCapability: mockedResolveCapability,
+  scaleDown,
+});