From 2816c4cc2b81d1870176716f49f92974a700c25f Mon Sep 17 00:00:00 2001
From: Gerhard Stoebich <18708370+Flarna@users.noreply.github.com>
Date: Sat, 20 Feb 2021 01:30:07 +0100
Subject: [PATCH] feat: add withSpan helper

Add withSpan helper to execute a function on a context holding a given span.
---
 src/context/context.ts       | 24 +++++++++++++++-
 test/context/context.test.ts | 55 ++++++++++++++++++++++++++++++++++++
 test/internal/global.test.ts |  6 ++++
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/src/context/context.ts b/src/context/context.ts
index 6e59831b..7eeeeaf9 100644
--- a/src/context/context.ts
+++ b/src/context/context.ts
@@ -15,7 +15,7 @@
  */
 
 import { Context } from './types';
-import { Baggage, Span, SpanContext } from '../';
+import { Baggage, Span, SpanContext, context } from '../';
 import { NoopSpan } from '../trace/NoopSpan';
 
 /**
@@ -55,6 +55,28 @@ export function setSpan(context: Context, span: Span): Context {
   return context.setValue(SPAN_KEY, span);
 }
 
+/**
+ * Helper to execute the given function on a new context created from
+ * active context with given span set.
+ * Note: The global context manager is used to get active context.
+ * @param span span to set on context
+ * @param fn function to execute in new context
+ * @param thisArg optional receiver to be used for calling fn
+ * @param args optional arguments forwarded to fn
+ */
+export function withSpan<
+  A extends unknown[],
+  F extends (...args: A) => ReturnType<F>
+>(
+  span: Span,
+  fn: F,
+  thisArg?: ThisParameterType<F>,
+  ...args: A
+): ReturnType<F> {
+  const ctx = setSpan(context.active(), span);
+  return context.with(ctx, fn, thisArg, ...args);
+}
+
 /**
  * Wrap span context in a NoopSpan and set as span in a new
  * context
diff --git a/test/context/context.test.ts b/test/context/context.test.ts
index dfea20b9..bb2e7995 100644
--- a/test/context/context.test.ts
+++ b/test/context/context.test.ts
@@ -17,16 +17,44 @@
 import * as assert from 'assert';
 import {
   createContextKey,
+  getSpan,
   isInstrumentationSuppressed,
   ROOT_CONTEXT,
   suppressInstrumentation,
   unsuppressInstrumentation,
+  withSpan,
 } from '../../src/context/context';
+import { NoopSpan } from '../../src/trace/NoopSpan';
+import { Context, context, NoopContextManager } from '../../src';
 
 const SUPPRESS_INSTRUMENTATION_KEY = createContextKey(
   'OpenTelemetry Context Key SUPPRESS_INSTRUMENTATION'
 );
 
+// simple ContextManager supporting sync calls (NoopContextManager.active returns always ROOT_CONTEXT)
+class SimpleContextManager extends NoopContextManager {
+  active(): Context {
+    return this._activeContext;
+  }
+
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
+    context: Context,
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    const prevContext = this._activeContext;
+    try {
+      this._activeContext = context;
+      return fn.call(thisArg, ...args);
+    } finally {
+      this._activeContext = prevContext;
+    }
+  }
+
+  private _activeContext = ROOT_CONTEXT;
+}
+
 describe('Context Helpers', () => {
   describe('suppressInstrumentation', () => {
     it('should set suppress to true', () => {
@@ -78,4 +106,31 @@ describe('Context Helpers', () => {
       });
     });
   });
+
+  describe('withSpan', () => {
+    before(() => {
+      const mgr = new SimpleContextManager();
+      context.setGlobalContextManager(mgr);
+    });
+
+    after(() => {
+      context.disable();
+    });
+
+    it('should run callback with span on context', () => {
+      const span = new NoopSpan();
+
+      function fnWithThis(this: string, a: string, b: number): string {
+        assert.strictEqual(getSpan(context.active()), span);
+        assert.strictEqual(this, 'that');
+        assert.strictEqual(arguments.length, 2);
+        assert.strictEqual(a, 'one');
+        assert.strictEqual(b, 2);
+        return 'done';
+      }
+
+      const res = withSpan(span, fnWithThis, 'that', 'one', 2);
+      assert.strictEqual(res, 'done');
+    });
+  });
 });
diff --git a/test/internal/global.test.ts b/test/internal/global.test.ts
index fc1f0a76..847cb485 100644
--- a/test/internal/global.test.ts
+++ b/test/internal/global.test.ts
@@ -46,6 +46,12 @@ describe('Global Utils', () => {
     delete _globalThis[Symbol.for('io.opentelemetry.js.api.0')];
   });
 
+  afterEach(() => {
+    api1.context.disable();
+    api1.propagation.disable();
+    api1.trace.disable();
+  });
+
   it('should change the global context manager', () => {
     const original = api1.context['_getContextManager']();
     const newContextManager = new NoopContextManager();