feat(gax): o11y tracer helper updates

core/packages/gax/src/index.ts

@@ -120,7 +120,7 @@ export {
   checkTelemetryEnabled,
 } from './util';
 
-export {StaticTraceContext} from './observability/TracerHelper';
+export {StaticTraceContext, traceAttempt} from './observability/TracerHelper';
 
 export {ServiceError, ChannelCredentials} from '@grpc/grpc-js';
 export {warn} from './warnings';

core/packages/gax/src/observability/TracerHelper.ts

@@ -14,7 +14,9 @@
  * limitations under the License.
  */
 
+import {EventEmitter} from 'events';
 import {Span, trace, Tracer} from '@opentelemetry/api';
+import {GaxCallResult} from '../apitypes';
 
 /**
  * Static metadata about the Google Cloud client library used to populate
@@ -67,23 +69,165 @@ export function getGaxTracer(): Tracer {
   return trace.getTracer('google-gax');
 }
 
+/**
+ * Checks if a value behaves like a Promise or Thenable.
+ *
+ * Note: It is not sufficient to check `result instanceof Promise` because:
+ * 1. Custom classes implementing `CancellablePromise` or Thenables may not
+ *    inherit directly from the native JavaScript `Promise` prototype.
+ * 2. GAX callers or custom callers may return objects such as `OngoingCallPromise`
+ *    that hold the actual promise on a `.promise` property.
+ * 3. Promises originating from different execution realms (such as Node.js vm
+ *    contexts or different package bundles) fail `instanceof Promise` checks.
+ */
+function isPromiseLike<T = unknown>(value: unknown): value is PromiseLike<T> {
+  return (
+    value instanceof Promise ||
+    (value !== null &&
+      (typeof value === 'object' || typeof value === 'function') &&
+      typeof (value as {then?: unknown}).then === 'function')
+  );
+}
+
+/**
+ * Extracts a PromiseLike target from a value, supporting native Promises,
+ * custom Thenables, classes implementing CancellablePromise, and OngoingCallPromise wrappers.
+ */
+function getPromiseTarget<T = unknown>(value: unknown): PromiseLike<T> | null {
+  if (isPromiseLike<T>(value)) {
+    return value;
+  }
+  if (
+    value !== null &&
+    (typeof value === 'object' || typeof value === 'function') &&
+    'promise' in (value as object) &&
+    isPromiseLike<T>((value as {promise?: unknown}).promise)
+  ) {
+    return (value as {promise: PromiseLike<T>}).promise;
+  }
+  return null;
+}
+
+/**
+ * Manages span lifecycle for Promise-based operations.
+ *
+ * @template T
+ * @param {T} promise - The promise returned from the traced operation.
+ * @param {function} recordError - Callback to record errors on the span.
+ * @param {function} endSpan - Callback to end the span idempotently.
+ */
+export function handlePromise<T>(
+  promise: T,
+  recordError: (err: unknown) => void,
+  endSpan: () => void,
+): void {
+  let spanEnded = false;
+  const endSpanOnce = () => {
+    if (!spanEnded) {
+      spanEnded = true;
+      endSpan();
+    }
+  };
+
+  const target = getPromiseTarget(promise) ?? promise;
+  Promise.resolve(target)
+    .then(() => {
+      endSpanOnce();
+      return null;
+    })
+    .catch(err => {
+      if (!spanEnded) {
+        recordError(err);
+        endSpanOnce();
+      }
+    });
+}
+
+/**
+ * Manages span lifecycle for Stream-based operations and cleans up event listeners.
+ *
+ * @param {EventEmitter} stream - The stream returned from the traced operation.
+ * @param {function} recordError - Callback to record errors on the span.
+ * @param {function} endSpan - Callback to end the span idempotently.
+ */
+export function handleStream(
+  stream: EventEmitter,
+  recordError: (err: unknown) => void,
+  endSpan: () => void,
+): void {
+  let spanEnded = false;
+
+  const cleanup = () => {
+    stream.removeListener('error', onError);
+    stream.removeListener('end', onEnd);
+    stream.removeListener('close', onClose);
+  };
+
+  const endSpanOnce = () => {
+    if (!spanEnded) {
+      spanEnded = true;
+      cleanup();
+      endSpan();
+    }
+  };
+
+  const onError = (err: unknown) => {
+    if (!spanEnded) {
+      recordError(err);
+      endSpanOnce();
+    }
+  };
+
+  const onEnd = () => {
+    endSpanOnce();
+  };
+
+  const onClose = () => {
+    endSpanOnce();
+  };
+
+  stream.on('error', onError);
+  stream.on('end', onEnd);
+  stream.on('close', onClose);
+}
+
 /**
  * Executes a function within an active OpenTelemetry span, populating standard
  * GCP telemetry attributes and recording errors/exceptions if thrown.
  *
  * @template T
  * @param {DynamicTraceContext} dynamicArgs - Dynamic trace context for the RPC call.
  * @param {StaticTraceContext} staticArgs - Static trace context for the client library.
- * @param {() => Promise<T>} fn - The asynchronous operation to trace.
- * @returns {Promise<T>} The result of the traced operation.
+ * @param {function} fn - The operation to trace.
+ * @param {boolean} [isStreamCall=false] - Whether the operation is a stream call (true) or promise call (false).
+ * @returns {T} The result of the traced operation.
  */
-export async function traceAttempt<T = unknown>(
+export function traceAttempt(
+  dynamicArgs: DynamicTraceContext,
+  staticArgs: StaticTraceContext,
+  fn: () => GaxCallResult,
+  isStreamCall?: boolean,
+): GaxCallResult;
+export function traceAttempt<T extends EventEmitter>(
+  dynamicArgs: DynamicTraceContext,
+  staticArgs: StaticTraceContext,
+  fn: () => T,
+  isStreamCall: true,
+): T;
+export function traceAttempt<T>(
   dynamicArgs: DynamicTraceContext,
   staticArgs: StaticTraceContext,
-  fn: () => Promise<T>,
-): Promise<T> {
+  fn: () => T,
+  isStreamCall?: false,
+): T;
+export function traceAttempt(
+  dynamicArgs: DynamicTraceContext,
+  staticArgs: StaticTraceContext,
+  fn: () => GaxCallResult,
+  isStreamCall = false,
+): GaxCallResult {
   const spanName = `${dynamicArgs.clientName}.${dynamicArgs.methodName}`;
-  return getGaxTracer().startActiveSpan(spanName, {}, async (span: Span) => {
+  return getGaxTracer().startActiveSpan(spanName, {}, (span: Span) => {
     span.setAttributes({
       'gcp.client.service': staticArgs.gcpClientService,
       'gcp.client.version': staticArgs.gcpVersion,
@@ -93,22 +237,51 @@ export async function traceAttempt<T = unknown>(
       'gcp.method.type': dynamicArgs.rpcType,
     });
 
+    let spanEnded = false;
+    const endSpan = () => {
+      if (!spanEnded) {
+        spanEnded = true;
+        span.end();
+      }
+    };
+
+    const recordError = (e: unknown) => {
+      if (e instanceof Error) {
+        span.setAttributes({
+          'error.message': e.message,
+          'error.type': e.constructor?.name ?? e.name,
+        });
+        span.recordException(e);
+        if (e.name) {
+          span.setAttribute('exception.type', e.name);
+        }
+      } else {
+        const message = String(e);
+        span.setAttributes({
+          'error.message': message,
+        });
+        span.recordException(message);
+      }
+    };
+
     try {
-      const result = await fn();
+      const result = fn();
+      // Use getPromiseTarget instead of `result instanceof Promise` to ensure custom
+      // thenables, CancellablePromise implementations, and OngoingCallPromise wrappers
+      // are properly tracked rather than leaving spans unclosed or ending them prematurely.
+      const promiseTarget = !isStreamCall ? getPromiseTarget(result) : null;
+      if (isStreamCall && result instanceof EventEmitter) {
+        handleStream(result, recordError, endSpan);
+      } else if (promiseTarget) {
+        handlePromise(promiseTarget, recordError, endSpan);
+      } else {
+        endSpan();
+      }
       return result;
     } catch (e) {
-      const err = e as Error;
-      span.setAttributes({
-        'error.message': err.message,
-        'error.type': err.constructor?.name ?? err.name,
-      });
-      span.recordException(err);
-      if (err.name) {
-        span.setAttribute('exception.type', err.name);
-      }
+      recordError(e);
+      endSpan();
       throw e;
-    } finally {
-      span.end();
     }
   });
 }