feat: add support to forward args in context.with (open-telemetry#1883)

Co-authored-by: Daniel Dyla <[email protected]>

Author: Flarna

api/src/api/context.ts

@@ -78,12 +78,16 @@ export class ContextAPI {
    *
    * @param context context to be active during function execution
    * @param fn function to execute in a context
+   * @param thisArg optional receiver to be used for calling fn
+   * @param args optional arguments forwarded to fn
    */
-  public with<T extends (...args: unknown[]) => ReturnType<T>>(
+  public with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
     context: Context,
-    fn: T
-  ): ReturnType<T> {
-    return this._getContextManager().with(context, fn);
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    return this._getContextManager().with(context, fn, thisArg, ...args);
   }
 
   /**

api/src/context-base/NoopContextManager.ts

@@ -22,11 +22,13 @@ export class NoopContextManager implements types.ContextManager {
     return ROOT_CONTEXT;
   }
 
-  with<T extends (...args: unknown[]) => ReturnType<T>>(
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
     _context: types.Context,
-    fn: T
-  ): ReturnType<T> {
-    return fn();
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    return fn.call(thisArg, ...args);
   }
 
   bind<T>(target: T, _context?: types.Context): T {

api/src/context-base/types.ts

@@ -50,11 +50,15 @@ export interface ContextManager {
    * Run the fn callback with object set as the current active context
    * @param context Any object to set as the current active context
    * @param fn A callback to be immediately run within a specific context
+   * @param thisArg optional receiver to be used for calling fn
+   * @param args optional arguments forwarded to fn
    */
-  with<T extends (...args: unknown[]) => ReturnType<T>>(
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
     context: Context,
-    fn: T
-  ): ReturnType<T>;
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F>;
 
   /**
    * Bind an object as the current context (or a specific one)

api/test/api/api.test.ts

@@ -41,6 +41,26 @@ describe('API', () => {
     assert.strictEqual(typeof tracer, 'object');
   });
 
+  describe('Context', () => {
+    it('with should forward this, arguments and return value', () => {
+      function fnWithThis(this: string, a: string, b: number): string {
+        assert.strictEqual(this, 'that');
+        assert.strictEqual(arguments.length, 2);
+        assert.strictEqual(a, 'one');
+        assert.strictEqual(b, 2);
+        return 'done';
+      }
+
+      const res = context.with(ROOT_CONTEXT, fnWithThis, 'that', 'one', 2);
+      assert.strictEqual(res, 'done');
+
+      assert.strictEqual(
+        context.with(ROOT_CONTEXT, () => 3.14),
+        3.14
+      );
+    });
+  });
+
   describe('GlobalTracerProvider', () => {
     const spanContext = {
       traceId: 'd4cda95b652f4a1592b449d5929fda1b',

api/test/context-base/NoopContextManager.test.ts

@@ -70,6 +70,30 @@ describe('NoopContextManager', () => {
         return done();
       });
     });
+
+    it('should forward this, arguments and return value', () => {
+      function fnWithThis(this: string, a: string, b: number): string {
+        assert.strictEqual(this, 'that');
+        assert.strictEqual(arguments.length, 2);
+        assert.strictEqual(a, 'one');
+        assert.strictEqual(b, 2);
+        return 'done';
+      }
+
+      const res = contextManager.with(
+        ROOT_CONTEXT,
+        fnWithThis,
+        'that',
+        'one',
+        2
+      );
+      assert.strictEqual(res, 'done');
+
+      assert.strictEqual(
+        contextManager.with(ROOT_CONTEXT, () => 3.14),
+        3.14
+      );
+    });
   });
 
   describe('.active()', () => {