fix(gax): unify dynamicArgs, preserve outer span context in logging-only mode, and test stream error logs

Author: shivanee-p

core/packages/gax/src/createApiCall.ts

@@ -45,7 +45,6 @@ import {logCall, logNetworkAttempt} from './observability/LoggerHelper';
 import {checkLoggingEnabled, checkTracingEnabled} from './observability/util';
 import {
   StaticTraceContext,
-  DynamicTraceContext,
   AttemptTraceContext,
   resolveStaticTraceContext,
 } from './observability/contextUtil';
@@ -92,19 +91,12 @@ export function createApiCall(
       : undefined;
   const serviceName = settings.apiName?.split('.').pop() ?? '';
   const isFallback = Boolean(_fallback);
-  const dynamicArgs: DynamicTraceContext | undefined =
+  const dynamicArgs: AttemptTraceContext | undefined =
     tracingEnabled || loggingEnabled
       ? {
           clientName: serviceName ? `${serviceName}Client` : '',
           methodName: settings.otherArgs?.internalMethodName ?? '',
           rpcType: isFallback ? 'http' : 'grpc',
-        }
-      : undefined;
-  // Context for per-attempt low level network CLIENT spans.
-  const attemptDynamicArgs: AttemptTraceContext | undefined =
-    (tracingEnabled || loggingEnabled) && dynamicArgs
-      ? {
-          ...dynamicArgs,
           apiName: settings.apiName ?? '',
         }
       : undefined;
@@ -114,8 +106,10 @@ export function createApiCall(
     callOptions?: CallOptions,
     callback?: APICallback,
   ) => {
-    // Capture the active client request span context to parent async low level network attempt spans.
-    const parentContext = tracingEnabled ? context.active() : undefined;
+    // Capture the active OpenTelemetry context to parent async attempt spans and preserve
+    // outer span correlation on attempt logs when context does not propagate across async gaps.
+    const parentContext =
+      tracingEnabled || loggingEnabled ? context.active() : undefined;
     let currentApiCaller = apiCaller;
 
     let thisSettings: CallSettings;
@@ -185,11 +179,7 @@ export function createApiCall(
           ?.streaming;
 
         // Wrap the transport call so each attempt (initial send and retries) emits a low level network CLIENT span.
-        if (
-          (tracingEnabled || loggingEnabled) &&
-          attemptDynamicArgs &&
-          staticArgs
-        ) {
+        if ((tracingEnabled || loggingEnabled) && dynamicArgs && staticArgs) {
           const callerWrappedFunc = wrappedFunc;
           wrappedFunc = (
             argument: {},
@@ -205,18 +195,26 @@ export function createApiCall(
                 cb,
               ]);
 
-            // If logging is enabled, wrap with L4 attempt error logging
+            // If logging is enabled, wrap with L4 attempt error logging.
+            // When tracing is disabled, restore `parentContext` around the attempt so
+            // `logNetworkAttempt` preserves the caller's outer OpenTelemetry span even
+            // without an async ContextManager across `funcPromise.then` or retry timers.
             if (loggingEnabled) {
               const originalAttemptExecution = attemptExecution;
-              attemptExecution = (cb: APICallback) =>
-                logNetworkAttempt(
-                  attemptDynamicArgs,
-                  staticArgs,
-                  loggedCallback =>
-                    originalAttemptExecution(loggedCallback ?? cb),
-                  Boolean(streaming),
-                  cb,
-                );
+              attemptExecution = (cb: APICallback) => {
+                const runLoggedAttempt = () =>
+                  logNetworkAttempt(
+                    dynamicArgs,
+                    staticArgs,
+                    loggedCallback =>
+                      originalAttemptExecution(loggedCallback ?? cb),
+                    Boolean(streaming),
+                    cb,
+                  );
+                return !tracingEnabled && parentContext
+                  ? context.with(parentContext, runLoggedAttempt)
+                  : runLoggedAttempt();
+              };
             }
 
             if (tracingEnabled) {
@@ -228,7 +226,7 @@ export function createApiCall(
               const originalAttemptExecution = attemptExecution;
               attemptExecution = (cb: APICallback) =>
                 traceAttempt(
-                  {...attemptDynamicArgs, resendCount},
+                  {...dynamicArgs, resendCount},
                   staticArgs,
                   tracedAttemptCallback =>
                     originalAttemptExecution(tracedAttemptCallback ?? cb),
@@ -310,12 +308,7 @@ export function createApiCall(
     return currentApiCaller.result(ongoingCall);
   };
 
-  if (
-    (tracingEnabled || loggingEnabled) &&
-    dynamicArgs &&
-    attemptDynamicArgs &&
-    staticArgs
-  ) {
+  if ((tracingEnabled || loggingEnabled) && dynamicArgs && staticArgs) {
     const isStreamingCall = apiCaller instanceof StreamingApiCaller;
     return (
       request: RequestType,
@@ -330,7 +323,7 @@ export function createApiCall(
         const originalCallExecution = callExecution;
         callExecution = (cb?: APICallback) =>
           logCall(
-            attemptDynamicArgs,
+            dynamicArgs,
             staticArgs,
             loggedCallback => originalCallExecution(loggedCallback ?? cb),
             isStreamingCall,

core/packages/gax/test/unit/apiCallable.ts

@@ -17,7 +17,16 @@
 import assert from 'assert';
 import {PassThrough} from 'stream';
 import {status} from '@grpc/grpc-js';
-import {SpanKind, SpanStatusCode} from '@opentelemetry/api';
+import {
+  Context,
+  ContextManager,
+  ROOT_CONTEXT,
+  Span,
+  SpanKind,
+  SpanStatusCode,
+  context,
+  trace,
+} from '@opentelemetry/api';
 import {GoogleAuth, PassThroughClient} from 'google-auth-library';
 import * as loggingUtils from 'google-logging-utils';
 import {afterEach, beforeEach, describe, it} from 'mocha';
@@ -402,6 +411,7 @@ describe('createApiCall', () => {
         clientName: 'EchoClient',
         methodName: 'Echo',
         rpcType: 'grpc',
+        apiName: 'google.example.v1.Echo',
       });
       assert.deepStrictEqual(staticArgs, {
         gcpClientService: 'example',
@@ -468,6 +478,7 @@ describe('createApiCall', () => {
         clientName: '',
         methodName: '',
         rpcType: 'grpc',
+        apiName: '',
       });
     });
 
@@ -2130,6 +2141,206 @@ describe('createApiCall', () => {
       assert.ok(spans.find(s => s.name === 'google.example.v1.Echo/Echo'));
       assert.ok(spans.find(s => s.name === 'EchoClient.Echo'));
     });
+
+    it('preserves outer OpenTelemetry span correlation across funcPromise and retries in logging-only mode without async context propagation', async () => {
+      // Replace OtelHarness's AsyncLocalStorageContextManager with a synchronous
+      // stack-only ContextManager so context.active() resets to ROOT_CONTEXT across
+      // funcPromise.then(...) microtasks and retry timers.
+      harness.teardown();
+      class StackOnlyContextManager implements ContextManager {
+        private stack: Context[] = [];
+        active(): Context {
+          return this.stack[this.stack.length - 1] ?? ROOT_CONTEXT;
+        }
+        with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
+          ctx: Context,
+          fn: F,
+          thisArg?: ThisParameterType<F>,
+          ...args: A
+        ): ReturnType<F> {
+          this.stack.push(ctx);
+          try {
+            return fn.apply(thisArg, args);
+          } finally {
+            this.stack.pop();
+          }
+        }
+        bind<T>(_ctx: Context, target: T): T {
+          return target;
+        }
+        enable(): this {
+          return this;
+        }
+        disable(): this {
+          this.stack = [];
+          return this;
+        }
+      }
+      context.setGlobalContextManager(new StackOnlyContextManager().enable());
+      trace.setGlobalTracerProvider(harness.provider);
+
+      let attempts = 0;
+      function failingRetriedFunc(
+        argument: {},
+        metadata: {},
+        options: {},
+        callback: (err: GoogleError | null, resp?: unknown) => void,
+      ) {
+        attempts++;
+        setTimeout(() => {
+          const error = new GoogleError(
+            attempts === 1 ? 'retryable failure' : 'terminal failure',
+          );
+          error.code = attempts === 1 ? FAKE_STATUS_CODE_1 : FAKE_STATUS_CODE_2;
+          callback(error);
+        }, 0);
+        return {cancel: () => {}};
+      }
+
+      const settings = new gax.CallSettings({
+        apiName: 'google.example.v1.Echo',
+        enableTelemetryTracing: false,
+        enableTelemetryLogging: true,
+        timeout: 0,
+        retry: utils.createRetryOptions(1, 1, 1, 0, 0, 0, 500),
+        otherArgs: {
+          internalMethodName: 'Echo',
+        },
+      });
+
+      const apiCall = gaxCreateApiCall(
+        Promise.resolve(failingRetriedFunc),
+        settings,
+      );
+
+      let outerSpan: Span | undefined;
+      const callPromise = trace
+        .getTracer('external-app')
+        .startActiveSpan('outer-request', span => {
+          outerSpan = span;
+          const p = apiCall({}, undefined);
+          span.end();
+          return p;
+        });
+
+      await assert.rejects(async () => {
+        await callPromise;
+      });
+
+      assert.ok(outerSpan);
+      assert.strictEqual(attempts, 2);
+      // 2 L4 DEBUG attempt logs + 1 L3 WARNING call log, all correlated to outerSpan.
+      assert.strictEqual(sink.logs.length, 3);
+      assert.deepStrictEqual(
+        sink.logs.map(l => l.fields.severity),
+        ['DEBUG', 'DEBUG', 'WARNING'],
+      );
+      for (const logEntry of sink.logs) {
+        assert.strictEqual(
+          logEntry.fields.telemetryTraceId,
+          outerSpan.spanContext().traceId,
+        );
+        assert.strictEqual(
+          logEntry.fields.telemetrySpanId,
+          outerSpan.spanContext().spanId,
+        );
+      }
+    });
+
+    it('emits one L4 DEBUG log and one L3 WARNING log when a streaming API call emits an error', async () => {
+      const settings = new gax.CallSettings({
+        apiName: 'google.example.v1.Echo',
+        enableTelemetryTracing: true,
+        enableTelemetryLogging: true,
+        otherArgs: {
+          internalMethodName: 'Echo',
+        },
+      });
+
+      const streamError = Object.assign(
+        new GoogleError('streaming test failure'),
+        {
+          code: status.UNAVAILABLE,
+        },
+      );
+      const spy = sinon.spy(() => {
+        const s = new PassThrough({objectMode: true});
+        setImmediate(() => {
+          s.emit('error', streamError);
+        });
+        return Object.assign(s, {cancel: () => {}});
+      });
+
+      const apiCall = gaxCreateApiCall(
+        spy as unknown as GRPCCall,
+        settings,
+        new StreamDescriptor(StreamType.SERVER_STREAMING, true),
+      );
+      const stream = apiCall({}, undefined) as CancellableStream;
+      assert.strictEqual(sink.logs.length, 0);
+
+      const caught = await new Promise<GoogleError>(resolve => {
+        stream.on('error', (err: GoogleError) => resolve(err));
+      });
+      // Allow StreamProxy / downstream 'error' listeners installed after caller listener (if any) to settle.
+      await new Promise<void>(resolve => setImmediate(resolve));
+
+      assert.strictEqual(caught, streamError);
+      assert.strictEqual(sink.logs.length, 2);
+
+      const debugLogs = sink.logs.filter(l => l.fields.severity === 'DEBUG');
+      const warnLogs = sink.logs.filter(l => l.fields.severity === 'WARNING');
+      assert.strictEqual(debugLogs.length, 1);
+      assert.strictEqual(warnLogs.length, 1);
+
+      const attemptLog = debugLogs[0];
+      const callLog = warnLogs[0];
+      assert.deepStrictEqual(attemptLog.args, [
+        'RPC attempt failed: %s',
+        'streaming test failure',
+      ]);
+      assert.deepStrictEqual(callLog.args, [
+        'Client API call failed: %s',
+        'streaming test failure',
+      ]);
+
+      const attemptOther = attemptLog.fields.other as loggerHelper.LogPayload;
+      const callOther = callLog.fields.other as loggerHelper.LogPayload;
+      assert.strictEqual(
+        attemptOther['rpc.method'],
+        'google.example.v1.Echo/Echo',
+      );
+      assert.strictEqual(
+        attemptOther['rpc.response.status_code'],
+        'UNAVAILABLE',
+      );
+      assert.strictEqual('exception.type' in attemptOther, false);
+      assert.strictEqual(
+        callOther['rpc.method'],
+        'google.example.v1.Echo/Echo',
+      );
+      assert.strictEqual(callOther['rpc.response.status_code'], 'UNAVAILABLE');
+      assert.strictEqual(callOther['exception.type'], 'GoogleError');
+      assert.strictEqual(
+        callOther['exception.message'],
+        'streaming test failure',
+      );
+
+      const spans = harness.getSpans('google-gax');
+      assert.strictEqual(spans.length, 2);
+      const attemptSpan = spans.find(
+        s => s.name === 'google.example.v1.Echo/Echo',
+      )!;
+      const clientRequestSpan = spans.find(s => s.name === 'EchoClient.Echo')!;
+      assert.strictEqual(
+        attemptLog.fields.telemetrySpanId,
+        attemptSpan.spanContext().spanId,
+      );
+      assert.strictEqual(
+        callLog.fields.telemetrySpanId,
+        clientRequestSpan.spanContext().spanId,
+      );
+    });
   });
 });