External Storage: add external_storage_runner to handle core extstore logic (#2094)

Author: cconstable

packages/common/src/internal-non-workflow/external-storage-runner.ts

@@ -0,0 +1,202 @@
+/**
+ * Engine that drives External Storage `store` and `retrieve` operations.
+ *
+ * @module
+ */
+import { temporal } from '@temporalio/proto';
+
+import { StorageDriverClaim } from '../converter/extstore';
+import type {
+  ExternalStorage,
+  StorageDriver,
+  StorageDriverRetrieveContext,
+  StorageDriverStoreContext,
+  StorageDriverTargetInfo,
+} from '../converter/extstore';
+import { ValueError } from '../errors';
+import type { Payload } from '../interfaces';
+import { decodeReferencePayload, encodeReferencePayload, isReferencePayload } from './extstore-helpers';
+
+const PayloadProto = temporal.api.common.v1.Payload;
+
+/** @internal @experimental */
+export interface ExternalStoreOptions {
+  /** Identity of the workflow or activity that produced the payloads. */
+  target?: StorageDriverTargetInfo;
+  /** Aborts the in-flight store operation. */
+  abortSignal?: AbortSignal;
+}
+
+/** @internal @experimental */
+export interface ExternalRetrieveOptions {
+  /** Aborts the in-flight retrieve operation. */
+  abortSignal?: AbortSignal;
+}
+
+/**
+ * Drives External Storage operations against a configured {@link ExternalStorage}.
+ *
+ * @internal
+ * @experimental
+ */
+export class ExternalStorageRunner {
+  constructor(private readonly externalStorage: ExternalStorage) {}
+
+  /**
+   * Replace each payload above the configured size threshold with a reference payload.
+   * Payloads below the threshold (or that the selector keeps inline) pass through
+   * unchanged. Order is preserved.
+   */
+  async store(payloads: Payload[], options: ExternalStoreOptions = {}): Promise<Payload[]> {
+    if (payloads.length === 0) return payloads;
+
+    const { driverSelector, payloadSizeThreshold } = this.externalStorage;
+    const { batchSignal, batchController } = makeBatchSignal(options.abortSignal);
+    const storeCtx: StorageDriverStoreContext = { abortSignal: batchSignal, target: options.target };
+
+    interface StoreItem {
+      index: number;
+      payload: Payload;
+      size: number;
+    }
+    const driverGroups = new Map<string, { driver: StorageDriver; items: StoreItem[] }>();
+
+    for (const [i, payload] of payloads.entries()) {
+      const size = payloadProtoSize(payload);
+      if (size < payloadSizeThreshold) continue;
+
+      const selected = driverSelector(storeCtx, payload);
+      if (selected === null) continue;
+      if (this.externalStorage.getDriver(selected.name) !== selected) {
+        throw new ValueError(
+          `Driver '${selected.name}' returned by driverSelector is not registered in ExternalStorage.drivers`
+        );
+      }
+
+      let group = driverGroups.get(selected.name);
+      if (group === undefined) {
+        group = { driver: selected, items: [] };
+        driverGroups.set(selected.name, group);
+      }
+      group.items.push({ index: i, payload, size });
+    }
+
+    if (driverGroups.size === 0) return payloads;
+
+    const result = payloads.slice();
+    await runWithAbortOnFirstError(batchController, [...driverGroups.values()], async (group) => {
+      const claims = await group.driver.store(
+        storeCtx,
+        group.items.map((it) => it.payload)
+      );
+      if (claims.length !== group.items.length) {
+        throw new ValueError(
+          `Driver '${group.driver.name}' returned ${claims.length} claims for ${group.items.length} payloads`
+        );
+      }
+      for (const [j, claim] of claims.entries()) {
+        const item = group.items[j]!;
+        result[item.index] = encodeReferencePayload({
+          driverName: group.driver.name,
+          claim,
+          sizeBytes: item.size,
+        });
+      }
+    });
+
+    return result;
+  }
+
+  /**
+   * Replace each reference payload in `payloads` with the payload bytes returned by the
+   * named driver. Non-reference payloads are passed through unchanged. Order is preserved.
+   */
+  async retrieve(payloads: Payload[], options: ExternalRetrieveOptions = {}): Promise<Payload[]> {
+    if (payloads.length === 0) return payloads;
+
+    const { batchSignal, batchController } = makeBatchSignal(options.abortSignal);
+    const retrieveCtx: StorageDriverRetrieveContext = { abortSignal: batchSignal };
+
+    interface RetrieveItem {
+      index: number;
+      claim: StorageDriverClaim;
+    }
+    const driverGroups = new Map<string, { driver: StorageDriver; items: RetrieveItem[] }>();
+
+    for (const [i, payload] of payloads.entries()) {
+      if (!isReferencePayload(payload)) continue;
+      const decoded = decodeReferencePayload(payload);
+      const driver = this.externalStorage.getDriver(decoded.driverName);
+      if (driver === null) {
+        throw new ValueError(`No driver registered with name '${decoded.driverName}'`);
+      }
+      let group = driverGroups.get(decoded.driverName);
+      if (group === undefined) {
+        group = { driver, items: [] };
+        driverGroups.set(decoded.driverName, group);
+      }
+      group.items.push({ index: i, claim: new StorageDriverClaim(decoded.claimData) });
+    }
+
+    if (driverGroups.size === 0) return payloads;
+
+    const result = payloads.slice();
+    await runWithAbortOnFirstError(batchController, [...driverGroups.values()], async (group) => {
+      const retrieved = await group.driver.retrieve(
+        retrieveCtx,
+        group.items.map((it) => it.claim)
+      );
+      if (retrieved.length !== group.items.length) {
+        throw new ValueError(
+          `Driver '${group.driver.name}' returned ${retrieved.length} payloads for ${group.items.length} claims`
+        );
+      }
+      for (const [j, retrievedPayload] of retrieved.entries()) {
+        const item = group.items[j]!;
+        result[item.index] = retrievedPayload;
+      }
+    });
+
+    return result;
+  }
+}
+
+// ============================================================================
+// Internal helpers
+// ============================================================================
+
+function payloadProtoSize(payload: Payload): number {
+  return PayloadProto.encode(payload).len;
+}
+
+/**
+ * Builds an internal controller composed with the caller's signal, so a failing
+ * driver call can abort its siblings.
+ */
+function makeBatchSignal(abortSignal?: AbortSignal): { batchSignal: AbortSignal; batchController: AbortController } {
+  const batchController = new AbortController();
+  const batchSignal = abortSignal ? AbortSignal.any([batchController.signal, abortSignal]) : batchController.signal;
+  return { batchSignal, batchController };
+}
+
+/**
+ * Run `task(item)` for each item in parallel. As soon as any task rejects,
+ * signal `controller.abort(reason)` so siblings can cancel mid-flight. Awaits
+ * all tasks regardless of outcome and re-throws the first rejection.
+ */
+async function runWithAbortOnFirstError<T>(
+  controller: AbortController,
+  items: T[],
+  task: (item: T) => Promise<void>
+): Promise<void> {
+  const promises = items.map(task);
+  for (const p of promises) {
+    p.catch((reason: unknown) => {
+      if (!controller.signal.aborted) controller.abort(reason);
+    });
+  }
+  const settled = await Promise.allSettled(promises);
+  for (const outcome of settled) {
+    if (outcome.status === 'rejected') throw outcome.reason;
+  }
+}

packages/common/src/internal-non-workflow/index.ts

@@ -7,6 +7,7 @@ export * from './codec-helpers';
 export * from './codec-types';
 export * from './data-converter-helpers';
 export * from './extstore-helpers';
+export * from './external-storage-runner';
 export * from './parse-host-uri';
 export * from './proxy-config';
 export * from './tls-config';

packages/test/src/extstore-fake-driver.ts

@@ -0,0 +1,57 @@
+/**
+ * Shared test double for {@link StorageDriver}. By default it stores and retrieves
+ * payloads via an in-memory `Map`, so a store→retrieve round-trip yields the original
+ * bytes. Tests that need failure or unusual-shape responses can override individual
+ * operations via `onStore` / `onRetrieve`.
+ *
+ * Every call is recorded on `storeCalls` / `retrieveCalls` (the recording happens
+ * before the override runs, so an override can introspect its own call).
+ */
+import type { Payload } from '@temporalio/common';
+import {
+  type StorageDriver,
+  StorageDriverClaim,
+  type StorageDriverRetrieveContext,
+  type StorageDriverStoreContext,
+} from '@temporalio/common/lib/converter/extstore';
+
+export interface FakeDriver extends StorageDriver {
+  storeCalls: { context: StorageDriverStoreContext; payloads: Payload[] }[];
+  retrieveCalls: { context: StorageDriverRetrieveContext; claims: StorageDriverClaim[] }[];
+}
+
+export interface FakeDriverOptions {
+  /** Driver name. Defaults to `'mem'`. */
+  name?: string;
+  /** Override default in-memory store. */
+  onStore?: (payloads: Payload[]) => StorageDriverClaim[] | Promise<StorageDriverClaim[]>;
+  /** Override default in-memory retrieve. */
+  onRetrieve?: (claims: StorageDriverClaim[]) => Payload[] | Promise<Payload[]>;
+}
+
+export function makeFakeDriver(opts: FakeDriverOptions = {}): FakeDriver {
+  const name = opts.name ?? 'mem';
+  const memory = new Map<string, Payload>();
+  let nextId = 0;
+  const driver: FakeDriver = {
+    name,
+    type: name,
+    storeCalls: [],
+    retrieveCalls: [],
+    async store(context, payloads) {
+      driver.storeCalls.push({ context, payloads });
+      if (opts.onStore) return opts.onStore(payloads);
+      return payloads.map((p) => {
+        const id = `${name}-${nextId++}`;
+        memory.set(id, p);
+        return new StorageDriverClaim({ id });
+      });
+    },
+    async retrieve(context, claims) {
+      driver.retrieveCalls.push({ context, claims });
+      if (opts.onRetrieve) return opts.onRetrieve(claims);
+      return claims.map((c) => memory.get(c.claimData.id!)!);
+    },
+  };
+  return driver;
+}