feat: expose withSpanAsync on NodeTracer open-telemetry#752

Author: vmarchaud

packages/opentelemetry-node/src/NodeTracer.ts

@@ -0,0 +1,71 @@
+/*!
+ * Copyright 2020, OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*!
+ * Copyright 2019, OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as api from '@opentelemetry/api';
+import { setActiveSpan } from '@opentelemetry/core';
+import {
+  BasicTracerProvider,
+  Tracer,
+  TracerConfig,
+} from '@opentelemetry/tracing';
+import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
+
+/**
+ * This class represents a nodejs-specific tracer.
+ */
+export class NodeTracer extends Tracer {
+  /**
+   * Constructs a new NodeTracer instance.
+   */
+  constructor(config: TracerConfig, _tracerProvider: BasicTracerProvider) {
+    super(config, _tracerProvider);
+  }
+
+  async withSpanAsync<
+    T extends Promise<any>,
+    U extends (...args: unknown[]) => T
+  >(span: api.Span, fn: U): Promise<T> {
+    // @ts-ignore
+    const contextManager = api.context._getContextManager();
+    if (contextManager instanceof AsyncHooksContextManager) {
+      return contextManager.withAsync(
+        setActiveSpan(api.context.active(), span),
+        fn
+      );
+    } else {
+      this.logger.warn(
+        `Using withAsync without AsyncHookContextManager doesn't work, please refer to`
+      );
+      return fn();
+    }
+  }
+}

packages/opentelemetry-node/src/NodeTracerProvider.ts

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2019, OpenTelemetry Authors
+ * Copyright 2020, OpenTelemetry Authors
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,9 +18,11 @@ import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
 import {
   BasicTracerProvider,
   SDKRegistrationConfig,
+  TracerConfig,
 } from '@opentelemetry/tracing';
 import { DEFAULT_INSTRUMENTATION_PLUGINS, NodeTracerConfig } from './config';
 import { PluginLoader } from './instrumentation/PluginLoader';
+import { NodeTracer } from './NodeTracer';
 
 /**
  * Register this TracerProvider for use with the OpenTelemetry API.
@@ -31,15 +33,16 @@ import { PluginLoader } from './instrumentation/PluginLoader';
  */
 export class NodeTracerProvider extends BasicTracerProvider {
   private readonly _pluginLoader: PluginLoader;
+  private readonly _nodeTracers: Map<string, NodeTracer> = new Map();
 
   /**
    * Constructs a new Tracer instance.
    */
-  constructor(config: NodeTracerConfig = {}) {
+  constructor(config?: NodeTracerConfig) {
     super(config);
 
     this._pluginLoader = new PluginLoader(this, this.logger);
-    this._pluginLoader.load(config.plugins || DEFAULT_INSTRUMENTATION_PLUGINS);
+    this._pluginLoader.load(config?.plugins ?? DEFAULT_INSTRUMENTATION_PLUGINS);
   }
 
   stop() {
@@ -51,7 +54,15 @@ export class NodeTracerProvider extends BasicTracerProvider {
       config.contextManager = new AsyncHooksContextManager();
       config.contextManager.enable();
     }
-
     super.register(config);
   }
+
+  getTracer(name: string, version = '*', config?: TracerConfig): NodeTracer {
+    const key = `${name}@${version}`;
+    if (!this._nodeTracers.has(key)) {
+      this._nodeTracers.set(key, new NodeTracer(config ?? this._config, this));
+    }
+
+    return this._nodeTracers.get(key)!;
+  }
 }

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

@@ -186,7 +186,7 @@ describe('NodeTracerProvider', () => {
   });
 
   describe('.withSpan()', () => {
-    it('should run context with AsyncHooksContextManager context manager', done => {
+    it('should run context with AsyncHooksContextManager', done => {
       provider = new NodeTracerProvider({});
       const span = provider.getTracer('default').startSpan('my-span');
       provider.getTracer('default').withSpan(span, () => {
@@ -202,7 +202,7 @@ describe('NodeTracerProvider', () => {
       );
     });
 
-    it('should run context with AsyncHooksContextManager context manager with multiple spans', done => {
+    it('should run context with AsyncHooksContextManager with multiple spans', done => {
       provider = new NodeTracerProvider({});
       const span = provider.getTracer('default').startSpan('my-span');
       provider.getTracer('default').withSpan(span, () => {
@@ -254,7 +254,7 @@ describe('NodeTracerProvider', () => {
   });
 
   describe('.bind()', () => {
-    it('should bind context with AsyncHooksContextManager context manager', done => {
+    it('should bind context with AsyncHooksContextManager', done => {
       const provider = new NodeTracerProvider({});
       const span = provider.getTracer('default').startSpan('my-span');
       const fn = () => {
@@ -268,4 +268,72 @@ describe('NodeTracerProvider', () => {
       return patchedFn();
     });
   });
+
+  describe('.withSpanAsync()', () => {
+    it('should run async context with AsyncHooksContextManager', done => {
+      provider = new NodeTracerProvider({});
+      const span = provider.getTracer('default').startSpan('my-span');
+      void provider.getTracer('default').withSpanAsync(span, async () => {
+        assert.deepStrictEqual(
+          provider.getTracer('default').getCurrentSpan(),
+          span
+        );
+        return done();
+      });
+      assert.deepStrictEqual(
+        provider.getTracer('default').getCurrentSpan(),
+        undefined
+      );
+    });
+
+    it('should run context with AsyncHooksContextManager with multiple spans', async () => {
+      provider = new NodeTracerProvider({});
+      let nestedHasBeenRun: boolean = false;
+      const span = provider.getTracer('default').startSpan('my-span');
+      await provider.getTracer('default').withSpanAsync(span, async () => {
+        assert.deepStrictEqual(
+          provider.getTracer('default').getCurrentSpan(),
+          span
+        );
+
+        const span1 = provider.getTracer('default').startSpan('my-span1');
+
+        await provider.getTracer('default').withSpanAsync(span1, async () => {
+          assert.deepStrictEqual(
+            provider.getTracer('default').getCurrentSpan(),
+            span1
+          );
+          assert.deepStrictEqual(
+            span1.context().traceId,
+            span.context().traceId
+          );
+          nestedHasBeenRun = true;
+        });
+      });
+      assert.deepStrictEqual(
+        provider.getTracer('default').getCurrentSpan(),
+        undefined
+      );
+      assert(nestedHasBeenRun);
+    });
+
+    it('should find correct context with promises', async () => {
+      provider = new NodeTracerProvider();
+      const span = provider.getTracer('default').startSpan('my-span');
+      await provider.getTracer('default').withSpanAsync(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
+      );
+    });
+  });
 });

packages/opentelemetry-tracing/src/BasicTracerProvider.ts

@@ -27,7 +27,7 @@ import { Resource } from '@opentelemetry/resources';
  * This class represents a basic tracer provider which platform libraries can extend
  */
 export class BasicTracerProvider implements api.TracerProvider {
-  private readonly _config: TracerConfig;
+  protected readonly _config: TracerConfig;
   private readonly _registeredSpanProcessors: SpanProcessor[] = [];
   private readonly _tracers: Map<string, Tracer> = new Map();