chore: rebase on ah-context refactor

Author: vmarchaud

packages/opentelemetry-node/src/NodeTracer.ts

@@ -19,31 +19,34 @@ import { setActiveSpan } from '@opentelemetry/core';
 import { Tracer } from '@opentelemetry/tracing';
 import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
 
+type UnPromisify<T> = T extends Promise<infer U> ? U : T;
+
 /**
  * This class represents a nodejs-specific tracer.
  */
 export class NodeTracer extends Tracer {
   /**
-   * Execute the provided function with the given span in the current context.
+   * Execute the provided function with the given span set in the current context.
    *
    * **NOTE**: This function is experimental, refer to to
    *  https://github.com/open-telemetry/opentelemetry-js/tree/master/packages/opentelemetry-node#withspanasync
    */
-  async withSpanAsync<T extends any, U extends () => Promise<T>>(
-    span: api.Span,
-    fn: U
-  ): Promise<void> {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  async withSpanAsync<
+    T extends () => Promise<any>,
+    U = UnPromisify<ReturnType<T>>
+  >(span: api.Span, fn: T): Promise<U> {
     const contextManager = api.context.getContextManager();
     if (contextManager instanceof AsyncHooksContextManager) {
-      await contextManager.withAsync(
+      return await contextManager.withAsync(
         setActiveSpan(api.context.active(), span),
         fn
       );
     } else {
       this.logger.warn(
-        `Using withAsync without AsyncHookContextManager doesn't work, please refer to https://github.com/open-telemetry/opentelemetry-js/tree/master/packages/opentelemetry-node#withspanasync`
+        "Using withAsync without AsyncHookContextManager doesn't work, please refer to https://github.com/open-telemetry/opentelemetry-js/tree/master/packages/opentelemetry-node#withspanasync"
       );
-      await fn();
+      return await fn();
     }
   }
 }

packages/opentelemetry-node/test/NodeTracerProvider.test.ts

@@ -232,25 +232,6 @@ describe('NodeTracerProvider', () => {
         undefined
       );
     });
-
-    it('should find correct context with promises', async () => {
-      provider = new NodeTracerProvider();
-      const span = provider.getTracer('default').startSpan('my-span');
-      await provider.getTracer('default').withSpan(span, async () => {
-        for (let i = 0; i < 3; i++) {
-          await sleep(5).then(() => {
-            assert.deepStrictEqual(
-              provider.getTracer('default').getCurrentSpan(),
-              span
-            );
-          });
-        }
-      });
-      assert.deepStrictEqual(
-        provider.getTracer('default').getCurrentSpan(),
-        undefined
-      );
-    });
   });
 
   describe('.bind()', () => {
@@ -288,7 +269,7 @@ describe('NodeTracerProvider', () => {
 
     it('should run context with AsyncHooksContextManager with multiple spans', async () => {
       provider = new NodeTracerProvider({});
-      let nestedHasBeenRun: boolean = false;
+      let nestedHasBeenRun = false;
       const span = provider.getTracer('default').startSpan('my-span');
       await provider.getTracer('default').withSpanAsync(span, async () => {
         assert.deepStrictEqual(

packages/opentelemetry-node/test/registration.test.ts

@@ -38,7 +38,7 @@ describe('API registration', () => {
     tracerProvider.register();
 
     assert.ok(
-      context['_getContextManager']() instanceof AsyncHooksContextManager
+      context.getContextManager() instanceof AsyncHooksContextManager
     );
     assert.ok(
       propagation['_getGlobalPropagator']() instanceof HttpTraceContext
@@ -57,7 +57,7 @@ describe('API registration', () => {
       propagator,
     });
 
-    assert.ok(context['_getContextManager']() === contextManager);
+    assert.ok(context.getContextManager() === contextManager);
     assert.ok(propagation['_getGlobalPropagator']() === propagator);
 
     assert.ok(trace.getTracerProvider() === tracerProvider);
@@ -69,7 +69,7 @@ describe('API registration', () => {
       contextManager: null,
     });
 
-    assert.ok(context['_getContextManager']() instanceof NoopContextManager);
+    assert.ok(context.getContextManager() instanceof NoopContextManager);
 
     assert.ok(
       propagation['_getGlobalPropagator']() instanceof HttpTraceContext
@@ -88,7 +88,7 @@ describe('API registration', () => {
     );
 
     assert.ok(
-      context['_getContextManager']() instanceof AsyncHooksContextManager
+      context.getContextManager() instanceof AsyncHooksContextManager
     );
     assert.ok(trace.getTracerProvider() === tracerProvider);
   });