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

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

Author: Flarna

packages/opentelemetry-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);
   }
 
   /**

packages/opentelemetry-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',

packages/opentelemetry-context-async-hooks/src/AbstractAsyncHooksContextManager.ts

@@ -42,10 +42,12 @@ export abstract class AbstractAsyncHooksContextManager
   implements ContextManager {
   abstract active(): Context;
 
-  abstract with<T extends (...args: unknown[]) => ReturnType<T>>(
+  abstract 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>;
 
   abstract enable(): this;
 

packages/opentelemetry-context-async-hooks/src/AsyncHooksContextManager.ts

@@ -38,13 +38,15 @@ export class AsyncHooksContextManager extends AbstractAsyncHooksContextManager {
     return this._stack[this._stack.length - 1] ?? ROOT_CONTEXT;
   }
 
-  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> {
     this._enterContext(context);
     try {
-      return fn();
+      return fn.call(thisArg!, ...args);
     } finally {
       this._exitContext();
     }

packages/opentelemetry-context-async-hooks/src/AsyncLocalStorageContextManager.ts

@@ -30,11 +30,14 @@ export class AsyncLocalStorageContextManager extends AbstractAsyncHooksContextMa
     return this._asyncLocalStorage.getStore() ?? ROOT_CONTEXT;
   }
 
-  with<T extends (...args: unknown[]) => ReturnType<T>>(
+  with<A extends unknown[], F extends (...args: A) => ReturnType<F>>(
     context: Context,
-    fn: T
-  ): ReturnType<T> {
-    return this._asyncLocalStorage.run(context, fn) as ReturnType<T>;
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
+    const cb = thisArg == null ? fn : fn.bind(thisArg);
+    return this._asyncLocalStorage.run(context, cb as any, ...args);
   }
 
   enable(): this {

packages/opentelemetry-context-async-hooks/test/AsyncHooksContextManager.test.ts

@@ -109,6 +109,30 @@ for (const contextManagerClass of [
         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
+        );
+      });
+
       it('should finally restore an old context', done => {
         const ctx1 = ROOT_CONTEXT.setValue(key1, 'ctx1');
         const ctx2 = ROOT_CONTEXT.setValue(key1, 'ctx2');

packages/opentelemetry-context-base/src/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 {

packages/opentelemetry-context-base/src/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)

packages/opentelemetry-context-base/test/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()', () => {

packages/opentelemetry-context-zone-peer-dep/src/ZoneContextManager.ts

@@ -244,15 +244,19 @@ export class ZoneContextManager implements ContextManager {
    *     The context will be set as active
    * @param context A context (span) to be called with provided callback
    * @param fn Callback function
+   * @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 | null,
-    fn: () => ReturnType<T>
-  ): ReturnType<T> {
+    fn: F,
+    thisArg?: ThisParameterType<F>,
+    ...args: A
+  ): ReturnType<F> {
     const zoneName = this._createZoneName();
 
     const newZone = this._createZone(zoneName, context);
 
-    return newZone.run(fn, context);
+    return newZone.run(fn, thisArg, args);
   }
 }