feat: add tracer.startActiveSpan() (#54)

Co-authored-by: Gerhard Stöbich <deb2001-github@yahoo.de>
Co-authored-by: Daniel Dyla <dyladan@users.noreply.github.com>

Author: 3 people

src/trace/NoopTracer.ts

@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-import { getSpanContext } from '../trace/context-utils';
+import { context } from '../';
 import { Context } from '../context/types';
+import { getSpanContext, setSpan } from '../trace/context-utils';
 import { NonRecordingSpan } from './NonRecordingSpan';
 import { Span } from './span';
 import { isSpanContextValid } from './spancontext-utils';
@@ -45,6 +46,45 @@ export class NoopTracer implements Tracer {
       return new NonRecordingSpan();
     }
   }
+
+  startActiveSpan<F extends (span: Span) => ReturnType<F>>(
+    name: string,
+    arg2: F | SpanOptions,
+    arg3?: F | Context,
+    arg4?: F
+  ): ReturnType<F> | undefined {
+    let fn: F | undefined,
+      options: SpanOptions | undefined,
+      activeContext: Context | undefined;
+    if (arguments.length === 2 && typeof arg2 === 'function') {
+      fn = arg2;
+    } else if (
+      arguments.length === 3 &&
+      typeof arg2 === 'object' &&
+      typeof arg3 === 'function'
+    ) {
+      options = arg2;
+      fn = arg3;
+    } else if (
+      arguments.length === 4 &&
+      typeof arg2 === 'object' &&
+      typeof arg3 === 'object' &&
+      typeof arg4 === 'function'
+    ) {
+      options = arg2;
+      activeContext = arg3;
+      fn = arg4;
+    }
+
+    const parentContext = activeContext ?? context.active();
+    const span = this.startSpan(name, options, parentContext);
+    const contextWithSpanSet = setSpan(parentContext, span);
+
+    if (fn) {
+      return context.with(contextWithSpanSet, fn, undefined, span);
+    }
+    return;
+  }
 }
 
 function isSpanContext(spanContext: any): spanContext is SpanContext {

src/trace/ProxyTracer.ts

@@ -38,6 +38,16 @@ export class ProxyTracer implements Tracer {
     return this._getTracer().startSpan(name, options, context);
   }
 
+  startActiveSpan<F extends (span: Span) => unknown>(
+    _name: string,
+    _options: F | SpanOptions,
+    _context?: F | Context,
+    _fn?: F
+  ): ReturnType<F> {
+    const tracer = this._getTracer();
+    return Reflect.apply(tracer.startActiveSpan, tracer, arguments);
+  }
+
   /**
    * Try to get a tracer from the proxy tracer provider.
    * If the proxy tracer provider has no delegate, return a noop tracer.

src/trace/tracer.ts

@@ -37,4 +37,63 @@ export interface Tracer {
    *     span.end();
    */
   startSpan(name: string, options?: SpanOptions, context?: Context): Span;
+
+  /**
+   * Starts a new {@link Span} and calls the given function passing it the
+   * created span as first argument.
+   * Additionally the new span gets set in context and this context is activated
+   * for the duration of the function call.
+   *
+   * @param name The name of the span
+   * @param [options] SpanOptions used for span creation
+   * @param [context] Context to use to extract parent
+   * @param fn function called in the context of the span and receives the newly created span as an argument
+   * @returns return value of fn
+   * @example
+   *   const something = tracer.startActiveSpan('op', span => {
+   *     try {
+   *       do some work
+   *       span.setStatus({code: SpanStatusCode.OK});
+   *       return something;
+   *     } catch (err) {
+   *       span.setStatus({
+   *         code: SpanStatusCode.ERROR,
+   *         message: err.message,
+   *       });
+   *       throw err;
+   *     } finally {
+   *       span.end();
+   *     }
+   *   });
+   * @example
+   *   const span = tracer.startActiveSpan('op', span => {
+   *     try {
+   *       do some work
+   *       return span;
+   *     } catch (err) {
+   *       span.setStatus({
+   *         code: SpanStatusCode.ERROR,
+   *         message: err.message,
+   *       });
+   *       throw err;
+   *     }
+   *   });
+   *   do some more work
+   *   span.end();
+   */
+  startActiveSpan<F extends (span: Span) => unknown>(
+    name: string,
+    fn: F
+  ): ReturnType<F>;
+  startActiveSpan<F extends (span: Span) => unknown>(
+    name: string,
+    options: SpanOptions,
+    fn: F
+  ): ReturnType<F>;
+  startActiveSpan<F extends (span: Span) => unknown>(
+    name: string,
+    options: SpanOptions,
+    context: Context,
+    fn: F
+  ): ReturnType<F>;
 }

test/noop-implementations/noop-tracer.test.ts

@@ -16,12 +16,13 @@
 
 import * as assert from 'assert';
 import {
+  context,
   NoopTracer,
+  Span,
   SpanContext,
   SpanKind,
-  TraceFlags,
-  context,
   trace,
+  TraceFlags,
 } from '../../src';
 import { NonRecordingSpan } from '../../src/trace/NonRecordingSpan';
 
@@ -56,4 +57,28 @@ describe('NoopTracer', () => {
     assert(span.spanContext().spanId === parent.spanId);
     assert(span.spanContext().traceFlags === parent.traceFlags);
   });
+
+  it('should accept 2 to 4 args and start an active span', () => {
+    const tracer = new NoopTracer();
+    const name = 'span-name';
+    const fn = (span: Span) => {
+      try {
+        return 1;
+      } finally {
+        span.end();
+      }
+    };
+    const opts = { attributes: { foo: 'bar' } };
+    const ctx = context.active();
+
+    const a = tracer.startActiveSpan(name, fn);
+    assert.strictEqual(a, 1);
+
+    const b = tracer.startActiveSpan(name, opts, fn);
+
+    assert.strictEqual(b, 1);
+
+    const c = tracer.startActiveSpan(name, opts, ctx, fn);
+    assert.strictEqual(c, 1);
+  });
 });

test/proxy-implementations/proxy-tracer.test.ts

@@ -17,18 +17,18 @@
 import * as assert from 'assert';
 import * as sinon from 'sinon';
 import {
-  ProxyTracerProvider,
-  SpanKind,
-  TracerProvider,
-  ProxyTracer,
-  Tracer,
-  Span,
+  context,
   NoopTracer,
+  ProxyTracer,
+  ProxyTracerProvider,
   ROOT_CONTEXT,
+  Span,
+  SpanKind,
   SpanOptions,
+  Tracer,
+  TracerProvider,
 } from '../../src';
 import { NonRecordingSpan } from '../../src/trace/NonRecordingSpan';
-
 describe('ProxyTracer', () => {
   let provider: ProxyTracerProvider;
   const sandbox = sinon.createSandbox();
@@ -96,6 +96,10 @@ describe('ProxyTracer', () => {
         startSpan() {
           return delegateSpan;
         },
+
+        startActiveSpan() {
+          // stubbed
+        },
       };
 
       tracer = provider.getTracer('test');
@@ -114,6 +118,34 @@ describe('ProxyTracer', () => {
       assert.strictEqual(span, delegateSpan);
     });
 
+    it('should create active spans using the delegate tracer', () => {
+      // sinon types are broken with overloads, hence the any
+      // https://github.com/DefinitelyTyped/DefinitelyTyped/issues/36436
+      const startActiveSpanStub = sinon.stub<Tracer, any>(
+        delegateTracer,
+        'startActiveSpan'
+      );
+
+      const name = 'span-name';
+      const fn = (span: Span) => {
+        try {
+          return 1;
+        } finally {
+          span.end();
+        }
+      };
+      const opts = { attributes: { foo: 'bar' } };
+      const ctx = context.active();
+
+      startActiveSpanStub.withArgs(name, fn).returns(1);
+      startActiveSpanStub.withArgs(name, opts, fn).returns(2);
+      startActiveSpanStub.withArgs(name, opts, ctx, fn).returns(3);
+
+      assert.strictEqual(tracer.startActiveSpan(name, fn), 1);
+      assert.strictEqual(tracer.startActiveSpan(name, opts, fn), 2);
+      assert.strictEqual(tracer.startActiveSpan(name, opts, ctx, fn), 3);
+    });
+
     it('should pass original arguments to DelegateTracer#startSpan', () => {
       const startSpanStub = sandbox.stub(delegateTracer, 'startSpan');
 
@@ -130,6 +162,7 @@ describe('ProxyTracer', () => {
       assert.deepStrictEqual(Object.getOwnPropertyNames(NoopTracer.prototype), [
         'constructor',
         'startSpan',
+        'startActiveSpan',
       ]);
       sandbox.assert.calledOnceWithExactly(startSpanStub, name, options, ctx);
     });