feat(microvm): add lifecycle hook service

lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.test.ts

@@ -0,0 +1,230 @@
+import { ConditionalCheckFailedException, DeleteItemCommand, type DynamoDBClient } from '@aws-sdk/client-dynamodb';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+
+import {
+  AwsSdkDynamoDbRunnerConfigApi,
+  createAwsDynamoDbRunnerConfigConsumer,
+  type AwsDynamoDbRunnerConfigApi,
+} from './runner-config-consumer';
+
+const dynamoDbEnvironment = {
+  RUNNER_CONFIG_STORAGE_PROVIDER: 'aws_dynamodb',
+  RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME: 'runner-state',
+} as const;
+
+function namedError(name: string, message = 'provider detail'): Error {
+  const error = new Error(message);
+  error.name = name;
+  return error;
+}
+
+describe('AWS SDK DynamoDB runner config API', () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it('atomically deletes an unexpired composite-key item and returns its previous value', async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z'));
+    const send = vi.fn().mockResolvedValue({ Attributes: { value: { S: 'encoded-jit' } } });
+    const api = new AwsSdkDynamoDbRunnerConfigApi({ send } as unknown as DynamoDBClient);
+    const signal = new AbortController().signal;
+
+    await expect(api.deleteItem('runner-state', 'microvm-123', signal)).resolves.toBe('encoded-jit');
+
+    expect(send.mock.calls[0][0]).toBeInstanceOf(DeleteItemCommand);
+    expect(send.mock.calls[0][0].input).toEqual({
+      TableName: 'runner-state',
+      Key: {
+        scope: { S: 'microvm-123' },
+        id: { S: 'config' },
+      },
+      ConditionExpression: 'attribute_exists(#expires_at) AND #expires_at > :now',
+      ExpressionAttributeNames: { '#expires_at': 'expires_at' },
+      ExpressionAttributeValues: { ':now': { N: '1767225600' } },
+      ReturnValues: 'ALL_OLD',
+    });
+    expect(send.mock.calls[0][1]).toEqual({ abortSignal: signal });
+  });
+
+  it('treats missing or expired records as unavailable after the conditional delete', async () => {
+    const conditional = new ConditionalCheckFailedException({
+      $metadata: {},
+      message: 'record is absent or expired',
+    });
+    const api = new AwsSdkDynamoDbRunnerConfigApi({
+      send: vi.fn().mockRejectedValue(conditional),
+    } as unknown as DynamoDBClient);
+
+    await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).resolves.toBeUndefined();
+  });
+
+  it('supports a deserialized conditional error without exposing its message', async () => {
+    const api = new AwsSdkDynamoDbRunnerConfigApi({
+      send: vi.fn().mockRejectedValue(namedError('ConditionalCheckFailedException', 'expired-secret-detail')),
+    } as unknown as DynamoDBClient);
+
+    await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).resolves.toBeUndefined();
+  });
+
+  it('propagates non-conditional provider errors', async () => {
+    const error = namedError('AccessDeniedException');
+    const api = new AwsSdkDynamoDbRunnerConfigApi({
+      send: vi.fn().mockRejectedValue(error),
+    } as unknown as DynamoDBClient);
+
+    await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).rejects.toBe(error);
+  });
+
+  it('returns undefined when no item was deleted', async () => {
+    const api = new AwsSdkDynamoDbRunnerConfigApi({
+      send: vi.fn().mockResolvedValue({}),
+    } as unknown as DynamoDBClient);
+
+    await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).resolves.toBeUndefined();
+  });
+
+  it.each([{ Attributes: {} }, { Attributes: { value: { N: '1' } } }, { Attributes: { value: { S: '' } } }])(
+    'rejects a deleted item without a string value %#',
+    async (response) => {
+      const api = new AwsSdkDynamoDbRunnerConfigApi({
+        send: vi.fn().mockResolvedValue(response),
+      } as unknown as DynamoDBClient);
+
+      await expect(api.deleteItem('runner-state', 'microvm-123', new AbortController().signal)).rejects.toThrow(
+        'runner configuration record has an invalid value',
+      );
+    },
+  );
+});
+
+describe('DynamoDB runner config consumer', () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it('polls until an atomic delete returns the stored configuration', async () => {
+    const deleteItem = vi
+      .fn<AwsDynamoDbRunnerConfigApi['deleteItem']>()
+      .mockResolvedValueOnce(undefined)
+      .mockResolvedValueOnce('encoded-jit');
+    const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, {
+      api: { deleteItem },
+      callTimeoutMs: 100,
+      configTimeoutMs: 500,
+      pollIntervalMs: 1,
+    });
+
+    await expect(
+      consumer.consume('microvm-123', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).resolves.toBe('encoded-jit');
+    expect(deleteItem).toHaveBeenCalledTimes(2);
+    expect(deleteItem).toHaveBeenCalledWith('runner-state', 'microvm-123', expect.any(AbortSignal));
+  });
+
+  it('retries transient provider failures', async () => {
+    const deleteItem = vi
+      .fn<AwsDynamoDbRunnerConfigApi['deleteItem']>()
+      .mockRejectedValueOnce(namedError('ProvisionedThroughputExceededException'))
+      .mockResolvedValueOnce('encoded-jit');
+    const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, {
+      api: { deleteItem },
+      callTimeoutMs: 100,
+      configTimeoutMs: 500,
+      pollIntervalMs: 1,
+    });
+
+    await expect(
+      consumer.consume('microvm-123', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).resolves.toBe('encoded-jit');
+    expect(deleteItem).toHaveBeenCalledTimes(2);
+  });
+
+  it('sanitizes non-retryable provider failures', async () => {
+    const api: AwsDynamoDbRunnerConfigApi = {
+      deleteItem: vi.fn().mockRejectedValue(namedError('AccessDeniedException', 'encoded-jit-secret')),
+    };
+    const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, {
+      api,
+      callTimeoutMs: 100,
+      configTimeoutMs: 100,
+      pollIntervalMs: 1,
+    });
+
+    const pending = consumer.consume('microvm-123', {
+      deadlineMs: Date.now() + 1_000,
+      signal: new AbortController().signal,
+    });
+    await expect(pending).rejects.toThrow('failed to consume runner configuration from DynamoDB');
+    await expect(pending).rejects.not.toThrow('encoded-jit-secret');
+  });
+
+  it('validates the complete composite key before calling DynamoDB', async () => {
+    const api: AwsDynamoDbRunnerConfigApi = { deleteItem: vi.fn() };
+    const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, {
+      api,
+      callTimeoutMs: 100,
+      configTimeoutMs: 100,
+      pollIntervalMs: 1,
+    });
+
+    await expect(
+      consumer.consume('invalid/scope', {
+        deadlineMs: Date.now() + 1_000,
+        signal: new AbortController().signal,
+      }),
+    ).rejects.toThrow('runnerId is invalid');
+    expect(api.deleteItem).not.toHaveBeenCalled();
+  });
+
+  it('times out while missing or expired items remain unavailable', async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z'));
+    const api: AwsDynamoDbRunnerConfigApi = { deleteItem: vi.fn().mockResolvedValue(undefined) };
+    const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, {
+      api,
+      callTimeoutMs: 10,
+      configTimeoutMs: 20,
+      pollIntervalMs: 5,
+    });
+    const pending = consumer.consume('microvm-123', {
+      deadlineMs: Date.now() + 100,
+      signal: new AbortController().signal,
+    });
+    const rejection = expect(pending).rejects.toThrow(
+      'runner configuration did not become available before the deadline',
+    );
+
+    await vi.runAllTimersAsync();
+
+    await rejection;
+    expect(api.deleteItem).toHaveBeenCalled();
+  });
+
+  it('stops a provider call immediately when the caller aborts', async () => {
+    const api: AwsDynamoDbRunnerConfigApi = {
+      deleteItem: vi.fn().mockReturnValue(new Promise(() => undefined)),
+    };
+    const controller = new AbortController();
+    const consumer = createAwsDynamoDbRunnerConfigConsumer(dynamoDbEnvironment, {
+      api,
+      callTimeoutMs: 10_000,
+      configTimeoutMs: 10_000,
+      pollIntervalMs: 1,
+    });
+    const pending = consumer.consume('microvm-123', {
+      deadlineMs: Date.now() + 10_000,
+      signal: controller.signal,
+    });
+
+    controller.abort();
+
+    await expect(pending).rejects.toThrow('runner configuration consumption was cancelled');
+  });
+});

lambdas/libs/storage-providers/aws/dynamodb/runner-config-consumer.ts

@@ -0,0 +1,147 @@
+import {
+  ConditionalCheckFailedException,
+  DeleteItemCommand,
+  DynamoDBClient,
+  type DeleteItemCommandOutput,
+} from '@aws-sdk/client-dynamodb';
+
+import type {
+  AwsDynamoDbRunnerConfigStorageEnvironment,
+  RunnerConfigConsumer,
+  RunnerConfigConsumeOptions,
+} from '../../core';
+import {
+  delay,
+  isRetryableProviderError,
+  resolvePollingOptions,
+  throwIfCancelled,
+  validateConsumeOptions,
+  validateDynamoDbRunnerConfigKey,
+  validateDynamoDbTableName,
+  withCallDeadline,
+  type RunnerConfigPollingOptions,
+} from '../../runner-config-consumer-common';
+import { EXPIRES_AT_ATTRIBUTE, ID_ATTRIBUTE, RUNNER_CONFIG_ID, SCOPE_ATTRIBUTE, VALUE_ATTRIBUTE } from './keys';
+
+export interface AwsDynamoDbRunnerConfigApi {
+  deleteItem(tableName: string, scope: string, signal: AbortSignal): Promise<string | undefined>;
+}
+
+export class AwsSdkDynamoDbRunnerConfigApi implements AwsDynamoDbRunnerConfigApi {
+  private client?: DynamoDBClient;
+
+  public constructor(client?: DynamoDBClient) {
+    this.client = client;
+  }
+
+  private getClient(): DynamoDBClient {
+    // Do not use the Lambda tracing wrapper here: lifecycle hooks run inside
+    // the runner image and may be snapshotted before their first request.
+    this.client ??= new DynamoDBClient({ maxAttempts: 1 });
+    return this.client;
+  }
+
+  public async deleteItem(tableName: string, scope: string, signal: AbortSignal): Promise<string | undefined> {
+    let response: DeleteItemCommandOutput;
+    try {
+      response = await this.getClient().send(
+        new DeleteItemCommand({
+          TableName: tableName,
+          Key: {
+            [SCOPE_ATTRIBUTE]: { S: scope },
+            [ID_ATTRIBUTE]: { S: RUNNER_CONFIG_ID },
+          },
+          ConditionExpression: 'attribute_exists(#expires_at) AND #expires_at > :now',
+          ExpressionAttributeNames: { '#expires_at': EXPIRES_AT_ATTRIBUTE },
+          ExpressionAttributeValues: { ':now': { N: Math.floor(Date.now() / 1_000).toString() } },
+          ReturnValues: 'ALL_OLD',
+        }),
+        { abortSignal: signal },
+      );
+    } catch (error) {
+      if (
+        error instanceof ConditionalCheckFailedException ||
+        (error !== null &&
+          typeof error === 'object' &&
+          'name' in error &&
+          error.name === 'ConditionalCheckFailedException')
+      ) {
+        return undefined;
+      }
+      throw error;
+    }
+    if (response.Attributes === undefined) {
+      return undefined;
+    }
+    const value = response.Attributes[VALUE_ATTRIBUTE];
+    if (value?.S === undefined || value.S.length === 0) {
+      throw new Error('runner configuration record has an invalid value');
+    }
+    return value.S;
+  }
+}
+
+export interface AwsDynamoDbRunnerConfigConsumerOptions extends RunnerConfigPollingOptions {
+  api?: AwsDynamoDbRunnerConfigApi;
+}
+
+export function createAwsDynamoDbRunnerConfigConsumer(
+  environment: AwsDynamoDbRunnerConfigStorageEnvironment,
+  options: AwsDynamoDbRunnerConfigConsumerOptions = {},
+): RunnerConfigConsumer {
+  return new AwsDynamoDbRunnerConfigConsumer(
+    environment.RUNNER_CONFIG_DYNAMODB_RUNNER_STATE_TABLE_NAME,
+    options.api ?? new AwsSdkDynamoDbRunnerConfigApi(),
+    options,
+  );
+}
+
+class AwsDynamoDbRunnerConfigConsumer implements RunnerConfigConsumer {
+  private readonly callTimeoutMs: number;
+  private readonly configTimeoutMs: number;
+  private readonly pollIntervalMs: number;
+
+  public constructor(
+    private readonly tableName: string,
+    private readonly api: AwsDynamoDbRunnerConfigApi,
+    options: AwsDynamoDbRunnerConfigConsumerOptions,
+  ) {
+    const polling = resolvePollingOptions(options);
+    this.callTimeoutMs = polling.callTimeoutMs;
+    this.configTimeoutMs = polling.configTimeoutMs;
+    this.pollIntervalMs = polling.pollIntervalMs;
+  }
+
+  public async consume(runnerId: string, options: RunnerConfigConsumeOptions): Promise<string> {
+    validateConsumeOptions(options);
+    const tableName = validateDynamoDbTableName(this.tableName);
+    validateDynamoDbRunnerConfigKey(runnerId, RUNNER_CONFIG_ID);
+    const pollDeadline = Math.min(Date.now() + this.configTimeoutMs, options.deadlineMs);
+
+    while (Date.now() < pollDeadline) {
+      throwIfCancelled(options.signal);
+      try {
+        const runnerConfig = await withCallDeadline(options.signal, pollDeadline, this.callTimeoutMs, (callSignal) =>
+          this.api.deleteItem(tableName, runnerId, callSignal),
+        );
+        if (runnerConfig !== undefined) {
+          return runnerConfig;
+        }
+      } catch (error) {
+        if (options.signal.aborted) {
+          throw new Error('runner configuration consumption was cancelled');
+        }
+        if (!isRetryableProviderError(error)) {
+          throw new Error('failed to consume runner configuration from DynamoDB');
+        }
+      }
+
+      const remaining = pollDeadline - Date.now();
+      if (remaining > 0) {
+        await delay(Math.min(this.pollIntervalMs, remaining), options.signal);
+      }
+    }
+
+    throw new Error('runner configuration did not become available before the deadline');
+  }
+}