fix: no longer yield during OTEL handleSignal

Author: chris-olszewski

packages/interceptors-opentelemetry/src/worker/index.ts

@@ -34,7 +34,7 @@ export class OpenTelemetryActivityInboundInterceptor implements ActivityInboundC
   }
 
   async execute(input: ActivityExecuteInput, next: Next<ActivityInboundCallsInterceptor, 'execute'>): Promise<unknown> {
-    const context = extractContextFromHeaders(input.headers);
+    const context = await Promise.resolve(extractContextFromHeaders(input.headers));
     const spanName = `${SpanName.ACTIVITY_EXECUTE}${SPAN_DELIMITER}${this.ctx.info.activityType}`;
     return await instrument({ tracer: this.tracer, spanName, fn: () => next(input), context });
   }

packages/interceptors-opentelemetry/src/workflow/index.ts

@@ -51,6 +51,7 @@ function getTracer(): otel.Tracer {
  */
 export class OpenTelemetryInboundInterceptor implements WorkflowInboundCallsInterceptor {
   protected readonly tracer = getTracer();
+  protected readonly maybeInjectYield = true;
 
   public async execute(
     input: WorkflowExecuteInput,

packages/test/src/helpers-integration.ts

@@ -6,10 +6,8 @@ import {
   WorkflowFailedError,
   WorkflowHandle,
   WorkflowHandleWithFirstExecutionRunId,
-  WorkflowHandleWithSignaledRunId,
   WorkflowStartOptions,
   WorkflowUpdateFailedError,
-  WorkflowSignalWithStartOptionsWithArgs,
 } from '@temporalio/client';
 import {
   LocalTestWorkflowEnvironmentOptions,
@@ -197,16 +195,6 @@ export interface Helpers {
     fn: T,
     opts: Omit<WorkflowStartOptions, 'taskQueue' | 'workflowId'> & Partial<Pick<WorkflowStartOptions, 'workflowId'>>
   ): Promise<WorkflowHandleWithFirstExecutionRunId<T>>;
-  signalWithStart<T extends workflow.Workflow, U extends any[]>(
-    fn: T,
-    signal: workflow.SignalDefinition<U>
-  ): Promise<WorkflowHandleWithSignaledRunId<T>>;
-  signalWithStart<T extends workflow.Workflow, U extends any[]>(
-    fn: T,
-    signal: workflow.SignalDefinition<U>,
-    opts: Omit<WorkflowSignalWithStartOptionsWithArgs<U>, 'taskQueue' | 'workflowId'> &
-      Partial<Pick<WorkflowStartOptions, 'workflowId'>>
-  ): Promise<WorkflowHandleWithSignaledRunId<T>>;
   assertWorkflowUpdateFailed(p: Promise<any>, causeConstructor: ErrorConstructor, message?: string): Promise<void>;
   assertWorkflowFailedError(p: Promise<any>, causeConstructor: ErrorConstructor, message?: string): Promise<void>;
   updateHasBeenAdmitted(handle: WorkflowHandle<workflow.Workflow>, updateId: string): Promise<boolean>;
@@ -262,21 +250,6 @@ export function configurableHelpers<T>(
         ...opts,
       });
     },
-
-    async signalWithStart(
-      fn: workflow.Workflow,
-      signal: workflow.SignalDefinition<any[]>,
-      opts?: Omit<WorkflowSignalWithStartOptionsWithArgs<any[]>, 'taskQueue' | 'workflowId'> &
-        Partial<Pick<WorkflowStartOptions, 'workflowId'>>
-    ): Promise<WorkflowHandleWithSignaledRunId<workflow.Workflow>> {
-      return await testEnv.client.workflow.signalWithStart(fn, {
-        signal,
-        taskQueue,
-        workflowId: randomUUID(),
-        ...opts,
-      });
-    },
-
     async assertWorkflowUpdateFailed(
       p: Promise<any>,
       causeConstructor: ErrorConstructor,

packages/test/src/test-integration-workflows.ts

@@ -1,6 +1,3 @@
-import * as opentelemetry from '@opentelemetry/sdk-node';
-import { SemanticResourceAttributes } from '@opentelemetry/semantic-conventions';
-import { ConsoleSpanExporter } from '@opentelemetry/sdk-trace-node';
 import { setTimeout as setTimeoutPromise } from 'timers/promises';
 import { randomUUID } from 'crypto';
 import asyncRetry from 'async-retry';
@@ -52,7 +49,7 @@ import {
 } from './helpers-integration';
 import { overrideSdkInternalFlag } from './mock-internal-flags';
 import { heartbeatCancellationDetailsActivity } from './activities/heartbeat-cancellation-details';
-import { loadHistory, RUN_TIME_SKIPPING_TESTS, saveHistory, waitUntil } from './helpers';
+import { loadHistory, RUN_TIME_SKIPPING_TESTS, waitUntil } from './helpers';
 
 const test = makeTestFunction({
   workflowsPath: __filename,
@@ -516,7 +513,6 @@ test("Worker doesn't request Eager Activity Dispatch if no activities are regist
 
 const unblockSignal = defineSignal('unblock');
 const getBuildIdQuery = defineQuery<string>('getBuildId');
-const startSignal = defineSignal('startSignal');
 
 export async function buildIdTester(): Promise<void> {
   let blocked = true;

packages/test/src/test-otel.ts

@@ -8,12 +8,7 @@ import { SpanStatusCode } from '@opentelemetry/api';
 import { ExportResultCode } from '@opentelemetry/core';
 import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-grpc';
 import * as opentelemetry from '@opentelemetry/sdk-node';
-import {
-  BasicTracerProvider,
-  ConsoleSpanExporter,
-  InMemorySpanExporter,
-  SimpleSpanProcessor,
-} from '@opentelemetry/sdk-trace-base';
+import { BasicTracerProvider, InMemorySpanExporter, SimpleSpanProcessor } from '@opentelemetry/sdk-trace-base';
 import { SemanticResourceAttributes } from '@opentelemetry/semantic-conventions';
 import test from 'ava';
 import { v4 as uuid4 } from 'uuid';
@@ -518,6 +513,7 @@ if (RUN_INTEGRATION_TESTS) {
 }
 
 test('Can replay otel history from 1.11.3', async (t) => {
+  /*
   const staticResource = new opentelemetry.resources.Resource({
     [SemanticResourceAttributes.SERVICE_NAME]: 'ts-test-otel-worker',
   });
@@ -543,7 +539,6 @@ test('Can replay otel history from 1.11.3', async (t) => {
   });
   const client = new WorkflowClient();
 
-  /*
   const result = await worker.runUntil(async () => {
     const handle = await client.signalWithStart(workflows.signalStartOtel, {
       signal: workflows.startSignal,
@@ -577,11 +572,10 @@ test('Can replay otel history from 1.11.3', async (t) => {
       hist
     );
   });
-  // t.is('abc', result);
-  t.pass();
 });
 
 test('Can replay otel history from 1.13.1', async (t) => {
+  /*
   const staticResource = new opentelemetry.resources.Resource({
     [SemanticResourceAttributes.SERVICE_NAME]: 'ts-test-otel-worker',
   });
@@ -607,7 +601,6 @@ test('Can replay otel history from 1.13.1', async (t) => {
   });
   const client = new WorkflowClient();
 
-  /*
   const result = await worker.runUntil(async () => {
     const handle = await client.signalWithStart(workflows.signalStartOtel, {
       signal: workflows.startSignal,
@@ -641,6 +634,4 @@ test('Can replay otel history from 1.13.1', async (t) => {
       hist
     );
   });
-  // t.is('ac', result);
-  t.pass();
 });

packages/workflow/src/flags.ts

@@ -44,8 +44,6 @@ export const SdkFlags = {
    *        to implicitely have this flag on.
    */
   ProcessWorkflowActivationJobsAsSingleBatch: defineFlag(2, true, [buildIdSdkVersionMatches(/1\.11\.[01]/)]),
-
-  OpenTelemetryInterceptorInsertYieldPoint: defineFlag(3, false, [({ info }) => false]),
 } as const;
 
 function defineFlag(id: number, def: boolean, alternativeConditions?: AltConditionFn[]): SdkFlag {

packages/workflow/src/internals.ts

@@ -980,19 +980,8 @@ export class Activator implements ActivationHandler {
 
     const signalExecutionNum = this.signalHandlerExecutionSeq++;
     this.inProgressSignals.set(signalExecutionNum, { name: signalName, unfinishedPolicy });
-    const injectYield = shouldInjectYield(this.sdkVersion);
-    const addedInterceptor: WorkflowInterceptors['inbound'] = injectYield
-      ? [
-          {
-            handleSignal: async (input, next) => {
-              await Promise.resolve();
-              return next(input);
-            },
-          },
-        ]
-      : [];
     const execute = composeInterceptors(
-      [...addedInterceptor, ...interceptors],
+      this.maybeInjectYieldForOtelHandler(interceptors),
       'handleSignal',
       this.signalWorkflowNextHandler.bind(this)
     );
@@ -1272,6 +1261,31 @@ export class Activator implements ActivationHandler {
   failureToError(failure: ProtoFailure): Error {
     return this.failureConverter.failureToError(failure, this.payloadConverter);
   }
+
+  private maybeInjectYieldForOtelHandler(
+    interceptors: NonNullable<WorkflowInterceptors['inbound']>
+  ): NonNullable<WorkflowInterceptors['inbound']> {
+    if (!this.info.unsafe.isReplaying || !shouldInjectYield(this.sdkVersion)) {
+      return [...interceptors];
+    }
+    const otelInboundInterceptorIndex = findOpenTelemetryInboundInterceptor(interceptors);
+    if (otelInboundInterceptorIndex === null) {
+      return [...interceptors];
+    }
+    // A handler that only serves the insert a yield point in the interceptor handlers
+    const yieldHandleSignalInterceptor: NonNullable<WorkflowInterceptors['inbound']>[number] = {
+      handleSignal: async (input, next) => {
+        await Promise.resolve();
+        return next(input);
+      },
+    };
+    // Insert the yield handler before the OTEL one to synthesize the yield point added in the affected versions of the handler
+    return [
+      ...interceptors.slice(0, otelInboundInterceptorIndex),
+      yieldHandleSignalInterceptor,
+      ...interceptors.slice(otelInboundInterceptorIndex),
+    ];
+  }
 }
 
 function getSeq<T extends { seq?: number | null }>(activation: T): number {
@@ -1323,22 +1337,47 @@ then you can disable this warning by passing an option when setting the handler:
   )}`;
 }
 
+// Should only get run on replay
 function shouldInjectYield(version?: string): boolean {
   if (!version) {
     return false;
   }
-  const [major, minor, patch] = version.split('.');
+  const [major, minor, patchAndTags] = version.split('.', 3);
   // 1.11.5 - 1.13.1: need to inject
   if (major !== '1') return false;
 
+  // patch might have some extra stuff that needs cleaning
+  // basically "takeWhile digit"
+  let patch;
+  try {
+    const patchDigits = /[0-9]+/.exec(patchAndTags)?.[0];
+    patch = patchDigits ? Number.parseInt(patchDigits) : null;
+  } catch {
+    patch = null;
+  }
+
   switch (minor) {
     case '11':
-      return patch === '5';
+      // 1.11.3 was the last release that didn't inject a yield point
+      return Boolean(patch && patch > 3);
     case '12':
+      // Every 1.12 release requires a yield
       return true;
     case '13':
-      return patch === '1';
+      // 1.13.2 will be the first release since 1.11.3 that doesn't have a yield point in `handleSignal`
+      return Boolean(patch && patch < 2);
     default:
       return false;
   }
 }
+
+function findOpenTelemetryInboundInterceptor(
+  interceptors: NonNullable<WorkflowInterceptors['inbound']>
+): number | null {
+  const index = interceptors.findIndex(
+    (interceptor) =>
+      // We use a marker instead of `instanceof` to avoid taking a dependency on @temporalio/interceptors-opentelemetry
+      (interceptor as NonNullable<WorkflowInterceptors['inbound']> & { maybeInjectYield: boolean }).maybeInjectYield
+  );
+  return index !== -1 ? index : null;
+}