feat(storage): add shared DynamoDB runner storage

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

@@ -17,6 +17,13 @@ vi.mock('../github/auth', () => ({
   getStoredInstallationId: vi.fn().mockResolvedValue(undefined),
 }));
 
+vi.mock('@aws-github-runner/storage-providers', () => ({
+  createStorageProviders: vi.fn().mockReturnValue({
+    githubAppCredentials: { get: vi.fn() },
+  }),
+  getRunnerStateStore: vi.fn().mockReturnValue(undefined),
+}));
+
 vi.mock('../scale-runners/github-runner', () => ({
   createStartRunnerConfig: vi.fn(),
   getGitHubEnterpriseApiUrl: vi.fn(),

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

@@ -1,7 +1,12 @@
 import { Octokit } from '@octokit/rest';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import { resolveComputeProviderType } from '@aws-github-runner/compute-providers/provider-types';
-import { createStorageProviders, type StorageProviders } from '@aws-github-runner/storage-providers';
+import {
+  createStorageProviders,
+  getRunnerStateStore,
+  type RunnerStateRecord,
+  type StorageProviders,
+} from '@aws-github-runner/storage-providers';
 import yn from 'yn';
 
 import {
@@ -60,26 +65,43 @@ export async function adjust(event: PoolEvent): Promise<void> {
     runnerNamePrefix,
   );
 
-  // Look up the managed provider runners, but running does not mean idle.
-  const poolRunners = await computeProvider.listRunners({
+  const runnerStateStore = getRunnerStateStore();
+  let currentRunnerCount: number;
+  let numberOfRunnersInPool: number;
+  const providerRunners = await computeProvider.listRunners({
     environment,
     runnerOwner,
     runnerType: 'Org',
   });
-
-  const numberOfRunnersInPool = computeProvider.countAvailableRunners(poolRunners, runnerStatusses, includeBusyRunners);
+  if (runnerStateStore) {
+    const runnerStates = (await runnerStateStore.list({ computeProvider: computeProvider.type })).filter(
+      (record) => record.runnerOwner === runnerOwner && record.runnerType === 'Org',
+    );
+    const storedComputeResourceIds = new Set(runnerStates.map((record) => record.computeResourceId));
+    const untrackedProviderRunners = providerRunners.filter((runner) => !storedComputeResourceIds.has(runner.id));
+    // Inventory is canonical for tracked resources. Provider discovery contributes
+    // only untracked resources so a launch-before-state crash cannot over-provision.
+    currentRunnerCount = runnerStates.length + untrackedProviderRunners.length;
+    numberOfRunnersInPool =
+      countAvailableStoredRunners(runnerStates, runnerStatusses, includeBusyRunners) +
+      computeProvider.countAvailableRunners(untrackedProviderRunners, runnerStatusses, includeBusyRunners);
+  } else {
+    // Look up the managed provider runners, but running does not mean idle.
+    currentRunnerCount = providerRunners.length;
+    numberOfRunnersInPool = computeProvider.countAvailableRunners(providerRunners, runnerStatusses, includeBusyRunners);
+  }
   let topUp = event.poolSize - numberOfRunnersInPool;
 
   // The pool must never push the total number of runners (busy + idle) past the configured maximum.
-  // poolRunners contains every running runner for this type, so its length is the current total and no
-  // extra API call is needed. Without this clamp the pool keeps topping up against idle-only counts and
-  // can overshoot runners_maximum_count, while the scale-up lambda correctly refuses to launch.
+  // currentRunnerCount includes both canonical inventory and untracked provider recovery records. Without
+  // this clamp the pool keeps topping up against idle-only counts and can overshoot runners_maximum_count,
+  // while the scale-up lambda correctly refuses to launch.
   if (maximumRunners !== -1 && topUp > 0) {
-    const headroom = maximumRunners - poolRunners.length;
+    const headroom = maximumRunners - currentRunnerCount;
     if (topUp > headroom) {
       logger.info(
         `Capping pool top-up from ${topUp} to ${Math.max(headroom, 0)} to respect the maximum of ` +
-          `${maximumRunners} runners (currently ${poolRunners.length} running).`,
+          `${maximumRunners} runners (currently ${currentRunnerCount} running).`,
       );
       topUp = headroom;
     }
@@ -129,6 +151,33 @@ async function getInstallationId(
   ).data.id;
 }
 
+function countAvailableStoredRunners(
+  runnerStates: RunnerStateRecord[],
+  runnerStatuses: Map<string, RunnerStatus>,
+  includeBusyRunners: boolean,
+): number {
+  let available = 0;
+  for (const runner of runnerStates) {
+    if (runner.state === 'orphan' || runner.state === 'terminating') {
+      continue;
+    }
+
+    const status = runnerStatuses.get(runner.computeResourceId);
+    if ((status?.busy === false || includeBusyRunners) && status?.status === 'online') {
+      available++;
+    } else if (status === undefined && !runnerBootTimeExceeded(runner.createdAt)) {
+      available++;
+    }
+  }
+  return available;
+}
+
+function runnerBootTimeExceeded(createdAt: string): boolean {
+  const bootTimeMinutes = Number(process.env.RUNNER_BOOT_TIME_IN_MINUTES);
+  const launchTime = new Date(createdAt).getTime();
+  return launchTime + bootTimeMinutes * 60_000 < Date.now();
+}
+
 async function getGitHubRegisteredRunnnerStatusses(
   ghClient: Octokit,
   runnerOwner: string,

lambdas/functions/control-plane/src/scale-runners/github-runner.ts

@@ -1,11 +1,16 @@
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import {
   createStorageProviders,
+  getRunnerConfigStore,
+  getRunnerStateStore,
   type RunnerConfigMetadata,
+  type RunnerConfigRecord,
   type RunnerConfigStore,
   type GitHubAppCredentialsStore,
   type RunnerGroupCacheStore,
+  type RunnerStateStore,
 } from '@aws-github-runner/storage-providers';
+import type { ComputeProviderType } from '@aws-github-runner/compute-providers/provider-types';
 import { Octokit } from '@octokit/rest';
 
 import { getStoredInstallationId } from '../github/auth';
@@ -22,6 +27,8 @@ export interface GitHubRunnerMetadata {
 export interface StartRunnerConfigOptions {
   runnerConfigStore?: RunnerConfigStore;
   runnerGroupCacheStore?: RunnerGroupCacheStore;
+  computeProvider?: ComputeProviderType;
+  getRunnerConfigAccessScope?: (runnerId: string) => string | undefined;
   getRunnerConfigMetadata?: (runnerId: string) => RunnerConfigMetadata[];
   onJitConfigCreated?: (runnerId: string, metadata: GitHubRunnerMetadata) => Promise<void>;
 }
@@ -207,11 +214,22 @@ export async function createStartRunnerConfig(
   ghClient: Octokit,
   options: StartRunnerConfigOptions = {},
 ): Promise<string[]> {
-  const runnerConfigStore = options.runnerConfigStore ?? createStorageProviders().runnerConfig;
+  const runnerConfigStore = options.runnerConfigStore ?? getRunnerConfigStore();
+  const runnerStateStore = getRunnerStateStore();
+  if (runnerStateStore && options.computeProvider === undefined) {
+    throw new Error('A compute provider is required when runner state storage is enabled');
+  }
   if (githubRunnerConfig.enableJitConfig && githubRunnerConfig.ephemeral) {
-    return await createJitConfig(githubRunnerConfig, runnerIds, ghClient, runnerConfigStore, options);
+    return await createJitConfig(githubRunnerConfig, runnerIds, ghClient, runnerConfigStore, runnerStateStore, options);
   } else {
-    return await createRegistrationTokenConfig(githubRunnerConfig, runnerIds, ghClient, runnerConfigStore, options);
+    return await createRegistrationTokenConfig(
+      githubRunnerConfig,
+      runnerIds,
+      ghClient,
+      runnerConfigStore,
+      runnerStateStore,
+      options,
+    );
   }
 }
 
@@ -226,13 +244,15 @@ function addDelay(runnerIds: string[], runnerConfigStore: RunnerConfigStore) {
 /**
  * Creates registration token configuration for non-ephemeral runners.
  *
- * @returns Empty array (this configuration method does not have failure cases)
+ * @returns Runner IDs whose durable configuration failed. The legacy SSM path
+ * still throws on the first failure so existing all-or-nothing cleanup is preserved.
  */
 async function createRegistrationTokenConfig(
   githubRunnerConfig: CreateGitHubRunnerConfig,
   runnerIds: string[],
   ghClient: Octokit,
   runnerConfigStore: RunnerConfigStore,
+  runnerStateStore: RunnerStateStore | undefined,
   options: StartRunnerConfigOptions,
 ): Promise<string[]> {
   const { isDelay, delay, delayMilliseconds } = addDelay(runnerIds, runnerConfigStore);
@@ -243,18 +263,53 @@ async function createRegistrationTokenConfig(
     runner_service_config: removeTokenFromLogging(runnerServiceConfig),
   });
 
+  if (!runnerStateStore) {
+    for (const runnerId of runnerIds) {
+      const metadata = options.getRunnerConfigMetadata?.(runnerId);
+      await runnerConfigStore.create(createRunnerConfigRecord(runnerId, runnerServiceConfig.join(' '), options), {
+        metadata,
+      });
+      if (isDelay) {
+        // Delay to stay within the selected store's maximum write throughput.
+        await delay(delayMilliseconds);
+      }
+    }
+    return [];
+  }
+
+  const failedRunnerIds: string[] = [];
   for (const runnerId of runnerIds) {
-    await runnerConfigStore.create(
-      { runnerId, value: runnerServiceConfig.join(' ') },
-      { metadata: options.getRunnerConfigMetadata?.(runnerId) },
-    );
-    if (isDelay) {
-      // Delay to stay within the selected store's maximum write throughput.
-      await delay(delayMilliseconds);
+    try {
+      const metadata = options.getRunnerConfigMetadata?.(runnerId);
+      await createProvisioningRunnerState(runnerStateStore, githubRunnerConfig, runnerId, options, metadata);
+      await runnerConfigStore.create(createRunnerConfigRecord(runnerId, runnerServiceConfig.join(' '), options), {
+        metadata,
+      });
+      await runnerStateStore.activate(runnerId, { metadata });
+      if (isDelay) {
+        // Delay to stay within the selected store's maximum write throughput.
+        await delay(delayMilliseconds);
+      }
+    } catch (error) {
+      failedRunnerIds.push(runnerId);
+      logger.warn('Failed to configure runner, continuing with remaining instances', {
+        instance: runnerId,
+        error: error instanceof Error ? error.message : String(error),
+        retryable: true,
+      });
     }
   }
 
-  return [];
+  if (failedRunnerIds.length > 0) {
+    logger.error('Failed to configure some runner instances', {
+      failedInstances: failedRunnerIds,
+      totalInstances: runnerIds.length,
+      successfulInstances: runnerIds.length - failedRunnerIds.length,
+      retryable: true,
+    });
+  }
+
+  return failedRunnerIds;
 }
 
 /**
@@ -268,6 +323,7 @@ async function createJitConfig(
   runnerIds: string[],
   ghClient: Octokit,
   runnerConfigStore: RunnerConfigStore,
+  runnerStateStore: RunnerStateStore | undefined,
   options: StartRunnerConfigOptions,
 ): Promise<string[]> {
   const runnerGroupId = await getRunnerGroupId(githubRunnerConfig, ghClient, options.runnerGroupCacheStore);
@@ -285,6 +341,17 @@ async function createJitConfig(
         runnerGroupId: runnerGroupId,
         runnerLabels: runnerLabels,
       };
+      // Durable inventory needs metadata before JIT creation so failures leave a
+      // complete provisioning record. The legacy path retains its original timing.
+      let metadata = runnerStateStore ? options.getRunnerConfigMetadata?.(runnerId) : undefined;
+      await createProvisioningRunnerState(
+        runnerStateStore,
+        githubRunnerConfig,
+        runnerId,
+        options,
+        metadata,
+        ephemeralRunnerConfig,
+      );
       logger.debug(`Runner name: ${ephemeralRunnerConfig.runnerName}`);
       const runnerConfig =
         githubRunnerConfig.runnerType === 'Org'
@@ -304,18 +371,47 @@ async function createJitConfig(
 
       metricGitHubAppRateLimit(runnerConfig.headers, githubRunnerConfig.appIndex);
 
-      await options.onJitConfigCreated?.(runnerId, {
+      const githubRunnerMetadata = {
         githubRunnerId: runnerConfig.data.runner.id.toString(),
         runnerLabels,
-      });
+      };
+      if (runnerStateStore) {
+        try {
+          await runnerStateStore.recordGitHubIdentity(runnerId, {
+            ...githubRunnerMetadata,
+            runnerName: ephemeralRunnerConfig.runnerName,
+            metadata,
+          });
+        } catch (error) {
+          await deregisterJitRunnerAfterIdentityWriteFailure(
+            githubRunnerConfig,
+            ghClient,
+            runnerConfig.data.runner.id,
+            runnerId,
+          );
+          throw error;
+        }
+      }
+      await options.onJitConfigCreated?.(runnerId, githubRunnerMetadata);
 
       logger.debug('Runner JIT config for ephemeral runner generated.', {
         instance: runnerId,
       });
+      if (!runnerStateStore) {
+        metadata = options.getRunnerConfigMetadata?.(runnerId);
+      }
       await runnerConfigStore.create(
-        { runnerId, value: runnerConfig.data.encoded_jit_config },
-        { metadata: options.getRunnerConfigMetadata?.(runnerId) },
+        createRunnerConfigRecord(runnerId, runnerConfig.data.encoded_jit_config, options),
+        {
+          metadata,
+        },
       );
+      await runnerStateStore?.activate(runnerId, {
+        githubRunnerId: githubRunnerMetadata.githubRunnerId,
+        runnerLabels,
+        runnerName: ephemeralRunnerConfig.runnerName,
+        metadata,
+      });
       if (isDelay) {
         // Delay to stay within the selected store's maximum write throughput.
         await delay(delayMilliseconds);
@@ -342,6 +438,73 @@ async function createJitConfig(
   return failedRunnerIds;
 }
 
+async function deregisterJitRunnerAfterIdentityWriteFailure(
+  githubRunnerConfig: CreateGitHubRunnerConfig,
+  ghClient: Octokit,
+  githubRunnerId: number,
+  runnerId: string,
+): Promise<void> {
+  try {
+    if (githubRunnerConfig.runnerType === 'Org') {
+      await ghClient.actions.deleteSelfHostedRunnerFromOrg({
+        org: githubRunnerConfig.runnerOwner,
+        runner_id: githubRunnerId,
+      });
+    } else {
+      const [owner, repo] = githubRunnerConfig.runnerOwner.split('/');
+      await ghClient.actions.deleteSelfHostedRunnerFromRepo({ owner, repo, runner_id: githubRunnerId });
+    }
+    logger.info('De-registered JIT GitHub runner after its inventory identity could not be stored', {
+      instance: runnerId,
+      githubRunnerId,
+    });
+  } catch (cleanupError) {
+    logger.error('Failed to de-register JIT GitHub runner after its inventory identity could not be stored', {
+      instance: runnerId,
+      githubRunnerId,
+      error: cleanupError instanceof Error ? cleanupError.message : String(cleanupError),
+      retryable: true,
+    });
+  }
+}
+
+function createRunnerConfigRecord(
+  runnerId: string,
+  value: string,
+  options: StartRunnerConfigOptions,
+): RunnerConfigRecord {
+  const accessScope = options.getRunnerConfigAccessScope?.(runnerId);
+  return {
+    runnerId,
+    value,
+    ...(accessScope === undefined ? {} : { accessScope }),
+  };
+}
+
+async function createProvisioningRunnerState(
+  runnerStateStore: RunnerStateStore | undefined,
+  githubRunnerConfig: CreateGitHubRunnerConfig,
+  runnerId: string,
+  options: StartRunnerConfigOptions,
+  metadata: RunnerConfigMetadata[] | undefined,
+  jitConfig?: EphemeralRunnerConfig,
+): Promise<void> {
+  if (!runnerStateStore) {
+    return;
+  }
+
+  await runnerStateStore.create({
+    runnerId,
+    computeProvider: options.computeProvider!,
+    computeResourceId: runnerId,
+    runnerOwner: githubRunnerConfig.runnerOwner,
+    runnerType: githubRunnerConfig.runnerType,
+    runnerName: jitConfig?.runnerName,
+    runnerLabels: jitConfig?.runnerLabels,
+    metadata,
+  });
+}
+
 export function getGitHubEnterpriseApiUrl() {
   const ghesBaseUrl = process.env.GHES_URL;
   let ghesApiUrl = '';