chore: Refactor mass history replayer API (#974)

Author: mjameswh

.eslintrc

@@ -12,6 +12,7 @@
     "prettier"
   ],
   "rules": {
+    "eqeqeq": ["error", "always", { "null": "ignore" }],
     "no-duplicate-imports": "error",
     "object-shorthand": ["error", "always"],
     "deprecation/deprecation": "warn",

package-lock.json

@@ -16,6 +16,7 @@
     "@grpc/grpc-js": "^1.6.7",
     "@temporalio/common": "file:../common",
     "@temporalio/proto": "file:../proto",
+    "abort-controller": "^3.0.0",
     "long": "^5.2.0",
     "uuid": "^8.3.2"
   },

packages/client/package.json

@@ -0,0 +1,116 @@
+import { EventEmitter, on, once } from 'node:events';
+import { AbortController } from 'abort-controller';
+
+export interface MapAsyncOptions {
+  /**
+   * How many items to map concurrently. If set to less than 2 (or not set), then items are not mapped concurrently.
+   *
+   * When items are mapped concurrently, mapped items are returned by the resulting iterator in the order they complete
+   * mapping, not the order in which the corresponding source items were obtained from the source iterator.
+   *
+   * @default 1 (ie. items are not mapped concurrently)
+   */
+  concurrency?: number;
+
+  /**
+   * Maximum number of mapped items to keep in buffer, ready for consumption.
+   *
+   * Ignored unless `concurrency > 1`. No limit applies if set to `undefined`.
+   *
+   * @default unlimited
+   */
+  bufferLimit?: number | undefined;
+}
+
+function toAsyncIterator<A>(iterable: AsyncIterable<A>): AsyncIterator<A> {
+  return iterable[Symbol.asyncIterator]();
+}
+
+/**
+ * Return an async iterable that transforms items from a source iterable by mapping each item
+ * through a mapping function.
+ *
+ * If `concurrency > 1`, then up to `concurrency` items may be mapped concurrently. In that case,
+ * items are returned by the resulting iterator in the order they complete processing, not the order
+ * in which the corresponding source items were obtained from the source iterator.
+ *
+ * @param source the source async iterable
+ * @param mapFn a mapping function to apply on every item of the source iterable
+ */
+export async function* mapAsyncIterable<A, B>(
+  source: AsyncIterable<A>,
+  mapFn: (val: A) => Promise<B>,
+  options?: MapAsyncOptions
+): AsyncIterable<B> {
+  const { concurrency, bufferLimit } = options ?? {};
+
+  if (!concurrency || concurrency < 2) {
+    for await (const x of source) {
+      yield mapFn(x);
+    }
+    return;
+  }
+
+  const sourceIterator = toAsyncIterator(source);
+
+  const emitter = new EventEmitter();
+  const controller = new AbortController();
+  const emitterEventsIterable: AsyncIterable<[B]> = on(emitter, 'result', { signal: controller.signal });
+  const emitterError: Promise<unknown[]> = once(emitter, 'error');
+
+  const bufferLimitSemaphore =
+    typeof bufferLimit === 'number'
+      ? (() => {
+          const releaseEvents: AsyncIterator<void> = toAsyncIterator(
+            on(emitter, 'released', { signal: controller.signal })
+          );
+          let value = bufferLimit + concurrency;
+
+          return {
+            acquire: async () => {
+              while (value <= 0) {
+                await Promise.race([releaseEvents.next(), emitterError]);
+              }
+              value--;
+            },
+            release: () => {
+              value++;
+              emitter.emit('released');
+            },
+          };
+        })()
+      : undefined;
+
+  const mapper = async () => {
+    for (;;) {
+      await bufferLimitSemaphore?.acquire();
+      const val = await Promise.race([sourceIterator.next(), emitterError]);
+
+      if (Array.isArray(val)) return;
+      if ((val as IteratorResult<[B]>)?.done) return;
+
+      emitter.emit('result', await mapFn(val.value));
+    }
+  };
+
+  const mappers = Array(concurrency)
+    .fill(mapper)
+    .map((f: typeof mapper) => f());
+
+  Promise.all(mappers).then(
+    () => controller.abort(),
+    (err) => emitter.emit('error', err)
+  );
+
+  try {
+    for await (const [res] of emitterEventsIterable) {
+      bufferLimitSemaphore?.release();
+      yield res;
+    }
+  } catch (err: unknown) {
+    if ((err as Error)?.name === 'AbortError') {
+      return;
+    }
+    throw err;
+  }
+}

packages/client/src/iterators-utils.ts

@@ -4,6 +4,7 @@ import {
   CancelledFailure,
   compileRetryPolicy,
   mapToPayloads,
+  HistoryAndWorkflowId,
   QueryDefinition,
   RetryState,
   searchAttributePayloadConverter,
@@ -18,6 +19,7 @@ import {
   WorkflowResultType,
 } from '@temporalio/common';
 import { composeInterceptors } from '@temporalio/common/lib/interceptors';
+import { History } from '@temporalio/common/lib/proto-utils';
 import {
   decodeArrayFromPayloads,
   decodeFromPayloadsAtIndex,
@@ -65,6 +67,7 @@ import {
   LoadedWithDefaults,
   WithDefaults,
 } from './base-client';
+import { mapAsyncIterable } from './iterators-utils';
 
 /**
  * A client side handle to a single Workflow instance.
@@ -131,6 +134,11 @@ export interface WorkflowHandle<T extends Workflow = Workflow> extends BaseWorkf
    */
   describe(): Promise<WorkflowExecutionDescription>;
 
+  /**
+   * Return a workflow execution's history
+   */
+  fetchHistory(): Promise<History>;
+
   /**
    * Readonly accessor to the underlying WorkflowClient
    */
@@ -245,6 +253,19 @@ interface WorkflowHandleOptions extends GetWorkflowHandleOptions {
   runIdForResult?: string;
 }
 
+/**
+ * An iterable list of WorkflowExecution, as returned by {@link WorkflowClient.list}.
+ */
+interface AsyncWorkflowListIterable extends AsyncIterable<WorkflowExecutionInfo> {
+  /**
+   * Return an iterable of histories corresponding to this iterable's WorkflowExecutions.
+   * Workflow histories will be fetched concurrently.
+   *
+   * Useful in batch replaying
+   */
+  intoHistories: (intoHistoriesOptions?: IntoHistoriesOptions) => AsyncIterable<HistoryAndWorkflowId>;
+}
+
 /**
  * Options for {@link WorkflowClient.list}
  */
@@ -261,6 +282,31 @@ export interface ListOptions {
   query?: string;
 }
 
+/**
+ * Options for {@link WorkflowClient.list().intoHistories()}
+ */
+export interface IntoHistoriesOptions {
+  /**
+   * Maximum number of workflow histories to download concurrently.
+   *
+   * @default 5
+   */
+  concurrency?: number;
+
+  /**
+   * Maximum number of workflow histories to buffer ahead, ready for consumption.
+   *
+   * It is recommended to set `bufferLimit` to a rasonnably low number if it is expected that the
+   * iterable may be stopped before reaching completion (for example, when implementing a fail fast
+   * bach replay test).
+   *
+   * Ignored unless `concurrency > 1`. No limit applies if set to `undefined`.
+   *
+   * @default unlimited
+   */
+  bufferLimit?: number;
+}
+
 /**
  * Client for starting Workflow executions and creating Workflow handles.
  *
@@ -799,6 +845,22 @@ export class WorkflowClient extends BaseClient {
         (info as unknown as WorkflowExecutionDescription).raw = raw;
         return info;
       },
+      async fetchHistory() {
+        let nextPageToken: Uint8Array | undefined = undefined;
+        const events = Array<temporal.api.history.v1.IHistoryEvent>();
+        for (;;) {
+          const response: temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse =
+            await this.client.workflowService.getWorkflowExecutionHistory({
+              nextPageToken,
+              namespace: this.client.options.namespace,
+              execution: { workflowId, runId },
+            });
+          events.push(...(response.history?.events ?? []));
+          nextPageToken = response.nextPageToken;
+          if (nextPageToken == null || nextPageToken.length === 0) break;
+        }
+        return temporal.api.history.v1.History.create({ events });
+      },
       async signal<Args extends any[]>(def: SignalDefinition<Args> | string, ...args: Args): Promise<void> {
         const next = this.client._signalWorkflowHandler.bind(this.client);
         const fn = interceptors.length ? composeInterceptors(interceptors, 'signal', next) : next;
@@ -858,15 +920,7 @@ export class WorkflowClient extends BaseClient {
     });
   }
 
-  /**
-   * List workflows by given `query`.
-   *
-   * ⚠️ To use advanced query functionality, as of the 1.18 server release, you must use Elasticsearch based visibility.
-   *
-   * More info on the concept of "visibility" and the query syntax on the Temporal documentation site:
-   * https://docs.temporal.io/visibility
-   */
-  public async *list(options?: ListOptions): AsyncIterable<WorkflowExecutionInfo> {
+  protected async *_list(options?: ListOptions): AsyncIterable<WorkflowExecutionInfo> {
     let nextPageToken: Uint8Array = Buffer.alloc(0);
     for (;;) {
       const response = await this.workflowService.listWorkflowExecutions({
@@ -882,10 +936,36 @@ export class WorkflowClient extends BaseClient {
         yield await executionInfoFromRaw(raw, this.dataConverter);
       }
       nextPageToken = response.nextPageToken;
-      if (nextPageToken == null || nextPageToken.length == 0) break;
+      if (nextPageToken == null || nextPageToken.length === 0) break;
     }
   }
 
+  /**
+   * List workflows by given `query`.
+   *
+   * ⚠️ To use advanced query functionality, as of the 1.18 server release, you must use Elasticsearch based visibility.
+   *
+   * More info on the concept of "visibility" and the query syntax on the Temporal documentation site:
+   * https://docs.temporal.io/visibility
+   */
+  public list(options?: ListOptions): AsyncWorkflowListIterable {
+    return {
+      [Symbol.asyncIterator]: () => this._list(options)[Symbol.asyncIterator](),
+      intoHistories: (intoHistoriesOptions?: IntoHistoriesOptions) => {
+        return mapAsyncIterable(
+          this._list(options),
+          async ({ workflowId, runId }) => ({
+            workflowId,
+            history: await this.getHandle(workflowId, runId)
+              .fetchHistory()
+              .catch((_) => undefined),
+          }),
+          { concurrency: intoHistoriesOptions?.concurrency ?? 5 }
+        );
+      },
+    };
+  }
+
   protected getOrMakeInterceptors(workflowId: string, runId?: string): WorkflowClientInterceptor[] {
     if (typeof this.options.interceptors === 'object' && 'calls' in this.options.interceptors) {
       // eslint-disable-next-line deprecation/deprecation

packages/client/src/workflow-client.ts

@@ -62,3 +62,11 @@ export type ActivityInterface = Record<string, ActivityFunction>;
  * Mapping of Activity name to function
  */
 export type UntypedActivities = Record<string, ActivityFunction>;
+
+/**
+ * A workflow's history and ID. Useful for replay.
+ */
+export interface HistoryAndWorkflowId {
+  workflowId: string;
+  history: temporal.api.history.v1.History | unknown | undefined;
+}