Different approach + wip on #1421

mjameswh · mjameswh · commit c4e0b53832f4 · 2024-08-08T03:06:21.000-04:00
diff --git a/packages/test/src/test-interceptors.ts b/packages/test/src/test-interceptors.ts
@@ -282,6 +282,13 @@ if (RUN_INTEGRATION_TESTS) {
         workflowId: uuid4(),
       })
     );
-    t.deepEqual(events, ['activate: 0', 'concludeActivation: 1', 'activate: 0', 'concludeActivation: 1']);
+    t.deepEqual(events, [
+      'activate: 0',
+      'activate: 1',
+      'concludeActivation: 1',
+      'activate: 0',
+      'activate: 1',
+      'concludeActivation: 1',
+    ]);
   });
 }
diff --git a/packages/test/src/test-workflows.ts b/packages/test/src/test-workflows.ts
@@ -122,6 +122,7 @@ async function createWorkflow(
     randomnessSeed: Long.fromInt(1337).toBytes(),
     now: startTime,
     patches: [],
+    sdkFlags: [],
     showStackTraceSources: true,
   })) as VMWorkflow;
   return workflow;
@@ -1955,15 +1956,19 @@ test('scopeCancelledWhileWaitingOnExternalWorkflowCancellation', async (t) => {
 test('query not found - successString', async (t) => {
   const { workflowType } = t.context;
   {
-    const completion = await activate(
+    const completion = await activate(t, makeActivation(undefined, makeStartWorkflowJob(workflowType)));
+    compareCompletion(
       t,
-      makeActivation(undefined, makeStartWorkflowJob(workflowType), makeQueryWorkflowJob('qid', 'not-found'))
+      completion,
+      makeSuccess([makeCompleteWorkflowExecution(defaultPayloadConverter.toPayload('success'))])
     );
+  }
+  {
+    const completion = await activate(t, makeActivation(undefined, makeQueryWorkflowJob('qid', 'not-found')));
     compareCompletion(
       t,
       completion,
       makeSuccess([
-        makeCompleteWorkflowExecution(defaultPayloadConverter.toPayload('success')),
         makeRespondToQueryCommand({
           queryId: 'qid',
           failed: {
@@ -2097,11 +2102,14 @@ test("Pending promises can't unblock between signals and updates - 1.11.0+ - sig
     compareCompletion(
       t,
       completion,
-      makeSuccess([
-        { updateResponse: { protocolInstanceId: '2', accepted: {} } },
-        { updateResponse: { protocolInstanceId: '2', completed: defaultPayloadConverter.toPayload(3) } },
-        { completeWorkflowExecution: { result: defaultPayloadConverter.toPayload(3) } },
-      ])
+      makeSuccess(
+        [
+          { updateResponse: { protocolInstanceId: '2', accepted: {} } },
+          { updateResponse: { protocolInstanceId: '2', completed: defaultPayloadConverter.toPayload(3) } },
+          { completeWorkflowExecution: { result: defaultPayloadConverter.toPayload(3) } },
+        ],
+        [SdkFlags.GroupUpdatesJobsWithSignals]
+      )
     );
   }
 });
diff --git a/packages/worker/src/utils.ts b/packages/worker/src/utils.ts
@@ -3,14 +3,6 @@ import { IllegalStateError, ParentWorkflowInfo } from '@temporalio/workflow';
 
 export const MiB = 1024 ** 2;
 
-// ts-prune-ignore-next (no idea why ts-prune is complaining all of a sudden)
-export function partition<T>(arr: T[], predicate: (x: T) => boolean): [T[], T[]] {
-  const truthy = Array<T>();
-  const falsy = Array<T>();
-  arr.forEach((v) => (predicate(v) ? truthy : falsy).push(v));
-  return [truthy, falsy];
-}
-
 export function toMB(bytes: number, fractionDigits = 2): string {
   return (bytes / 1024 / 1024).toFixed(fractionDigits);
 }
diff --git a/packages/worker/src/worker.ts b/packages/worker/src/worker.ts
@@ -1262,6 +1262,7 @@ export class Worker {
       randomnessSeed: randomnessSeed.toBytes(),
       now: tsToMs(activation.timestamp),
       patches,
+      sdkFlags: activation.availableInternalFlags ?? [],
       showStackTraceSources: this.options.showStackTraceSources,
     });
 
diff --git a/packages/worker/src/workflow/vm-shared.ts b/packages/worker/src/workflow/vm-shared.ts
@@ -7,8 +7,7 @@ import type { WorkflowInfo, StackTraceFileLocation } from '@temporalio/workflow'
 import { type SinkCall } from '@temporalio/workflow/lib/sinks';
 import * as internals from '@temporalio/workflow/lib/worker-interface';
 import { Activator } from '@temporalio/workflow/lib/internals';
-import { assertValidFlag, SdkFlag, SdkFlags } from '@temporalio/workflow/lib/flags';
-import { partition } from '../utils';
+import { SdkFlags } from '@temporalio/workflow/lib/flags';
 import { Workflow } from './interface';
 import { WorkflowBundleWithSourceMapAndFilename } from './workflow-worker-thread/input';
 
@@ -274,7 +273,6 @@ export type WorkflowModule = typeof internals;
  * A Workflow implementation using Node.js' built-in `vm` module
  */
 export abstract class BaseVMWorkflow implements Workflow {
-  private readonly knownFlags = new Set<number>();
   unhandledRejection: unknown;
 
   constructor(
@@ -289,7 +287,7 @@ export abstract class BaseVMWorkflow implements Workflow {
    * Send request to the Workflow runtime's worker-interface
    */
   async getAndResetSinkCalls(): Promise<SinkCall[]> {
-    return this.workflowModule.getAndResetSinkCalls();
+    return this.activator.getAndResetSinkCalls();
   }
 
   /**
@@ -301,48 +299,45 @@ export abstract class BaseVMWorkflow implements Workflow {
   public async activate(
     activation: coresdk.workflow_activation.IWorkflowActivation
   ): Promise<coresdk.workflow_completion.IWorkflowActivationCompletion> {
-    if (this.context === undefined) {
-      throw new IllegalStateError('Workflow isolate context uninitialized');
-    }
-    if (!activation.jobs) {
-      throw new Error('Expected workflow activation jobs to be defined');
+    if (this.context === undefined) throw new IllegalStateError('Workflow isolate context uninitialized');
+    if (!activation.jobs) throw new Error('Expected workflow activation jobs to be defined');
+
+    // Queries are particular in many ways; handle them separately
+    const [queries, nonQueries] = partition(activation.jobs, ({ queryWorkflow }) => queryWorkflow != null);
+    if (queries.length > 0) {
+      // Core guarantees that a single activation will not contain both queries and other jobs
+      if (nonQueries.length > 0) throw new TypeError('Got both queries and other jobs in a single activation');
+      return this.activateQueries(activation);
     }
 
-    this.addKnownFlags(activation.availableInternalFlags ?? []);
+    // Let's prepare the activation
+    const [patches, nonPatches] = partition(nonQueries, ({ notifyHasPatch }) => notifyHasPatch != null);
+    this.workflowModule.startActivation(
+      coresdk.workflow_activation.WorkflowActivation.fromObject({ ...activation, jobs: patches })
+    );
 
+    // Extract signals and updates
     const hasUpdates = activation.jobs.some(({ doUpdate }) => doUpdate != null);
-    const groupUpdatesWithSignals =
-      hasUpdates && this.hasFlag(SdkFlags.GroupUpdatesJobsWithSignals, activation.isReplaying ?? false);
-
-    // Job processing order
-    // 1. patch notifications
-    // 2. signals and updates (updates )
-    // 3. anything left except for queries
-    // 4. queries
-    const [patches, nonPatches] = partition(activation.jobs, ({ notifyHasPatch }) => notifyHasPatch != null);
-    // const [signals, nonSignals] = partition(nonPatches, ({ signalWorkflow }) => signalWorkflow != null);
-    const [signals, nonSignals] = partition(
+    const groupUpdatesWithSignals = hasUpdates && this.activator.hasFlag(SdkFlags.GroupUpdatesJobsWithSignals);
+
+    const [signals, rest] = partition(
       nonPatches,
       ({ signalWorkflow, doUpdate }) => signalWorkflow != null || (doUpdate != null && groupUpdatesWithSignals)
     );
-    const [queries, rest] = partition(nonSignals, ({ queryWorkflow }) => queryWorkflow != null);
-    let batchIndex = 0;
-
-    // Loop and invoke each batch and wait for microtasks to complete.
-    // This is done outside of the isolate because when we used isolated-vm we couldn't wait for microtasks from inside the isolate, not relevant anymore.
-    for (const jobs of [patches, signals, rest, queries]) {
-      if (jobs.length === 0) {
-        continue;
-      }
+
+    // Loop and invoke each batch, waiting for microtasks to complete after each batch.
+    let batchIndex = 1;
+    for (const jobs of [signals, rest]) {
+      if (jobs.length === 0) continue;
       this.workflowModule.activate(
         coresdk.workflow_activation.WorkflowActivation.fromObject({ ...activation, jobs }),
         batchIndex++
       );
-      if (internals.shouldUnblockConditions(jobs[0])) {
-        this.tryUnblockConditions();
-      }
+      this.tryUnblockMicrotasks();
     }
+
     const completion = this.workflowModule.concludeActivation();
+
     // Give unhandledRejection handler a chance to be triggered.
     await new Promise(setImmediate);
     if (this.unhandledRejection) {
@@ -354,6 +349,13 @@ export abstract class BaseVMWorkflow implements Workflow {
     return completion;
   }
 
+  private activateQueries(
+    activation: coresdk.workflow_activation.IWorkflowActivation
+  ): coresdk.workflow_completion.IWorkflowActivationCompletion {
+    this.workflowModule.activateQueries(activation);
+    return this.workflowModule.concludeActivation();
+  }
+
   /**
    * If called (by an external unhandledRejection handler), activations will fail with provided error.
    */
@@ -362,12 +364,12 @@ export abstract class BaseVMWorkflow implements Workflow {
   }
 
   /**
-   * Call into the Workflow context to attempt to unblock any blocked conditions.
+   * Call into the Workflow context to attempt to unblock any blocked conditions and microtasks.
    *
    * This is performed in a loop allowing microtasks to be processed between
    * each iteration until there are no more conditions to unblock.
    */
-  protected tryUnblockConditions(): void {
+  protected tryUnblockMicrotasks(): void {
     for (;;) {
       const numUnblocked = this.workflowModule.tryUnblockConditions();
       if (numUnblocked === 0) break;
@@ -378,34 +380,11 @@ export abstract class BaseVMWorkflow implements Workflow {
    * Do not use this Workflow instance after this method has been called.
    */
   public abstract dispose(): Promise<void>;
+}
 
-  /**
-   * Called early while handling an activation to register known flags
-   *
-   * This duplicates the functionality of `Activator.addKnownFlags`, for use outside of the sandbox.
-   */
-  private addKnownFlags(flags: number[]): void {
-    for (const flag of flags) {
-      assertValidFlag(flag);
-      this.knownFlags.add(flag);
-    }
-  }
-
-  /**
-   * Check if a flag is known to the Workflow Execution; if not, enable the flag if workflow
-   * is not replaying and the flag is configured to be enabled by default.
-   *
-   * This duplicates the functionality of `Activator.addKnownFlags`, for use outside of the sandbox.
-   * Note that we can't rely on WorkflowInfo.isReplaying here.
-   */
-  protected hasFlag(flag: SdkFlag, isReplaying: boolean): boolean {
-    if (this.knownFlags.has(flag.id)) {
-      return true;
-    }
-    if (!isReplaying && flag.default) {
-      this.knownFlags.add(flag.id);
-      return true;
-    }
-    return false;
-  }
+function partition<T>(arr: T[], predicate: (x: T) => boolean): [T[], T[]] {
+  const truthy = Array<T>();
+  const falsy = Array<T>();
+  arr.forEach((v) => (predicate(v) ? truthy : falsy).push(v));
+  return [truthy, falsy];
 }
diff --git a/packages/workflow/src/interfaces.ts b/packages/workflow/src/interfaces.ts
@@ -453,6 +453,7 @@ export interface WorkflowCreateOptions {
   randomnessSeed: number[];
   now: number;
   patches: string[];
+  sdkFlags: number[];
   showStackTraceSources: boolean;
 }
 
diff --git a/packages/workflow/src/internals.ts b/packages/workflow/src/internals.ts
@@ -129,15 +129,6 @@ export class Activator implements ActivationHandler {
    */
   readonly bufferedSignals = Array<coresdk.workflow_activation.ISignalWorkflow>();
 
-  /**
-   * Holds buffered query calls until a handler is registered.
-   *
-   * **IMPORTANT** queries are only buffered until workflow is started.
-   * This is required because async interceptors might block workflow function invocation
-   * which delays query handler registration.
-   */
-  protected readonly bufferedQueries = Array<coresdk.workflow_activation.IQueryWorkflow>();
-
   /**
    * Mapping of update name to handler and validator
    */
@@ -292,13 +283,6 @@ export class Activator implements ActivationHandler {
    */
   protected cancelled = false;
 
-  /**
-   * This is tracked to allow buffering queries until a workflow function is called.
-   * TODO(bergundy): I don't think this makes sense since queries run last in an activation and must be responded to in
-   * the same activation.
-   */
-  protected workflowFunctionWasCalled = false;
-
   /**
    * The next (incremental) sequence to assign when generating completable commands
    */
@@ -366,6 +350,7 @@ export class Activator implements ActivationHandler {
     showStackTraceSources,
     sourceMap,
     getTimeOfDay,
+    sdkFlags,
     randomnessSeed,
     patches,
     registeredActivityNames,
@@ -378,11 +363,8 @@ export class Activator implements ActivationHandler {
     this.random = alea(randomnessSeed);
     this.registeredActivityNames = registeredActivityNames;
 
-    if (info.unsafe.isReplaying) {
-      for (const patchId of patches) {
-        this.notifyHasPatch({ patchId });
-      }
-    }
+    this.addKnownPatches(patches.map((patch) => ({ patchId: patch })));
+    this.addKnownFlags(sdkFlags);
   }
 
   mutateWorkflowInfo(fn: (info: WorkflowInfo) => WorkflowInfo): void {
@@ -443,19 +425,7 @@ export class Activator implements ActivationHandler {
     if (workflow === undefined) {
       throw new IllegalStateError('Workflow uninitialized');
     }
-    let promise: Promise<any>;
-    try {
-      promise = workflow(...args);
-    } finally {
-      // Queries must be handled even if there was an exception when invoking the Workflow function.
-      this.workflowFunctionWasCalled = true;
-      // Empty the buffer
-      const buffer = this.bufferedQueries.splice(0);
-      for (const activation of buffer) {
-        this.queryWorkflow(activation);
-      }
-    }
-    return await promise;
+    return await workflow(...args);
   }
 
   public startWorkflow(activation: coresdk.workflow_activation.IStartWorkflow): void {
@@ -586,11 +556,6 @@ export class Activator implements ActivationHandler {
   }
 
   public queryWorkflow(activation: coresdk.workflow_activation.IQueryWorkflow): void {
-    if (!this.workflowFunctionWasCalled) {
-      this.bufferedQueries.push(activation);
-      return;
-    }
-
     const { queryType, queryId, headers } = activation;
     if (!(queryType && queryId)) {
       throw new TypeError('Missing query activation attributes');
@@ -845,12 +810,22 @@ export class Activator implements ActivationHandler {
   }
 
   public notifyHasPatch(activation: coresdk.workflow_activation.INotifyHasPatch): void {
-    if (!activation.patchId) {
-      throw new TypeError('Notify has patch missing patch name');
-    }
+    // Technically, notifyHasPatch jobs should not reach here, but just for the sake of completeness
+    if (!this.info.unsafe.isReplaying)
+      throw new IllegalStateError('Unexpected notifyHasPatch job on non-replay activation');
+    if (!activation.patchId) throw new TypeError('notifyHasPatch missing patch id');
     this.knownPresentPatches.add(activation.patchId);
   }
 
+  public addKnownPatches(patches: coresdk.workflow_activation.INotifyHasPatch[]): void {
+    if (patches.length > 0 && !this.info.unsafe.isReplaying)
+      throw new IllegalStateError('Unexpected notifyHasPatch job on non-replay activation');
+    for (const patch of patches) {
+      if (!patch.patchId) throw new TypeError('notifyHasPatch missing patch id');
+      this.knownPresentPatches.add(patch.patchId);
+    }
+  }
+
   public patchInternal(patchId: string, deprecated: boolean): boolean {
     if (this.workflow === undefined) {
       throw new IllegalStateError('Patches cannot be used before Workflow starts');
diff --git a/packages/workflow/src/worker-interface.ts b/packages/workflow/src/worker-interface.ts

Original file line number	Diff line number	Diff line change
`@@ -453,6 +453,7 @@ export interface WorkflowCreateOptions {`
`453`	`453`	`randomnessSeed: number[];`
`454`	`454`	`now: number;`
`455`	`455`	`patches: string[];`
	`456`	`+ sdkFlags: number[];`
`456`	`457`	`showStackTraceSources: boolean;`
`457`	`458`	`}`
`458`	`459`