Unwrap, Masswrap & Massunwrap

experimental/packages/opentelemetry-instrumentation/src/instrumentation.ts

@@ -61,28 +61,7 @@ export abstract class InstrumentationAbstract<T = any>
   }
 
   /* Api to wrap instrumented method */
-  protected _wrap = (
-    moduleExports: any,
-    name: string,
-    wrapper: (originalFn: any) => any
-  ) => {
-    try {
-      return shimmer.wrap(moduleExports, name, wrapper);
-    } catch (e) {
-      // shimmer doesn't handle Proxy objects well
-      // if there is an error from import in the middle providing
-      // a Proxy of a Module we have to pass it to shimmer as a regular object
-      const wrapped: any = shimmer.wrap(
-        Object.assign({}, moduleExports),
-        name,
-        wrapper
-      );
-      Object.defineProperty(moduleExports, name, {
-        value: wrapped,
-      });
-      return moduleExports;
-    }
-  };
+  protected _wrap = shimmer.wrap;
   /* Api to unwrap instrumented methods */
   protected _unwrap = shimmer.unwrap;
   /* Api to mass wrap instrumented method */

experimental/packages/opentelemetry-instrumentation/src/platform/node/instrumentation.ts

@@ -16,7 +16,9 @@
 
 import * as types from '../../types';
 import * as path from 'path';
+import * as util from 'util';
 import { satisfies } from 'semver';
+import * as shimmer from 'shimmer';
 import { InstrumentationAbstract } from '../../instrumentation';
 import {
   RequireInTheMiddleSingleton,
@@ -69,6 +71,90 @@ export abstract class InstrumentationBase<T = any>
     }
   }
 
+  protected override _wrap: typeof shimmer.wrap = (
+    moduleExports,
+    name,
+    wrapper
+  ) => {
+    if (!util.types.isProxy(moduleExports)) {
+      return shimmer.wrap(moduleExports, name, wrapper);
+    } else {
+      return this._wrapEsm(moduleExports, name, wrapper);
+    }
+  };
+
+  protected override _unwrap: typeof shimmer.unwrap = (moduleExports, name) => {
+    if (!util.types.isProxy(moduleExports)) {
+      return shimmer.unwrap(moduleExports, name);
+    } else {
+      return this._unwrapEsm(moduleExports, name);
+    }
+  };
+
+  private _wrapEsm: typeof shimmer.wrap = (moduleExports, name, wrapper) => {
+    const wrapped = shimmer.wrap(
+      Object.assign({}, moduleExports),
+      name,
+      wrapper
+    );
+    Object.defineProperty(moduleExports, name, {
+      value: wrapped,
+    });
+  };
+
+  private _unwrapEsm: typeof shimmer.unwrap = (moduleExports, name) => {
+    Object.defineProperty(moduleExports, name, {
+      value: moduleExports[name],
+    });
+  };
+
+  protected override _massWrap: typeof shimmer.massWrap = (
+    moduleExportsArray,
+    names,
+    wrapper
+  ) => {
+    if (!moduleExportsArray) {
+      diag.error('must provide one or more modules to patch');
+      return;
+    } else if (!Array.isArray(moduleExportsArray)) {
+      moduleExportsArray = [moduleExportsArray];
+    }
+
+    if (!(names && Array.isArray(names))) {
+      diag.error('must provide one or more functions to wrap on modules');
+      return;
+    }
+
+    moduleExportsArray.forEach(moduleExports => {
+      names.forEach(name => {
+        this._wrap(moduleExports, name, wrapper);
+      });
+    });
+  };
+
+  protected override _massUnwrap: typeof shimmer.massUnwrap = (
+    moduleExportsArray,
+    names
+  ) => {
+    if (!moduleExportsArray) {
+      diag.error('must provide one or more modules to patch');
+      return;
+    } else if (!Array.isArray(moduleExportsArray)) {
+      moduleExportsArray = [moduleExportsArray];
+    }
+
+    if (!(names && Array.isArray(names))) {
+      diag.error('must provide one or more functions to wrap on modules');
+      return;
+    }
+
+    moduleExportsArray.forEach(moduleExports => {
+      names.forEach(name => {
+        this._unwrap(moduleExports, name);
+      });
+    });
+  };
+
   private _warnOnPreloadedModules(): void {
     this._modules.forEach((module: InstrumentationModuleDefinition<T>) => {
       const { name } = module;

experimental/packages/opentelemetry-instrumentation/test/node/EsmInstrumentation.test.mjs

@@ -44,6 +44,36 @@ class TestInstrumentationWrapFn extends InstrumentationBase {
   }
 }
 
+class TestInstrumentationMasswrapFn extends InstrumentationBase {
+  constructor(config) {
+    super('test-esm-instrumentation', '0.0.1', config);
+  }
+  init() {
+    console.log('test-esm-instrumentation initialized!');
+    return new InstrumentationNodeModuleDefinition(
+      'test-esm-module',
+      ['*'],
+      moduleExports => {
+        this._massWrap(
+          [moduleExports],
+          ['testFunction', 'secondTestFunction'],
+          () => {
+            return () => 'patched';
+          }
+        );
+        return moduleExports;
+      },
+      moduleExports => {
+        this._massUnwrap(
+          [moduleExports],
+          ['testFunction', 'secondTestFunction']
+        );
+        return moduleExports;
+      }
+    );
+  }
+}
+
 class TestInstrumentationSimple extends InstrumentationBase {
   constructor(config) {
     super('test-esm-instrumentation', '0.0.1', config);
@@ -78,11 +108,31 @@ describe('when loading esm module', () => {
     assert.deepEqual(exported.testFunction(), 'patched');
   });
 
-  // it('should unwrap a patched function', async () => {
-  //   // disable to trigger unwrap
-  //   const exported = await import('test-esm-module');
-  //   instrumentationWrap.enable();
-  //   instrumentationWrap.disable();
-  //   assert.deepEqual(exported.testFunction(), 'original');
-  // });
+  it('should unwrap a patched function', async () => {
+    instrumentationWrap.enable();
+    // disable to trigger unwrap
+    instrumentationWrap.disable();
+    assert.deepEqual(exported.testFunction(), 'original');
+  });
+
+  it('should wrap multiple functions with masswrap', () => {
+    const instrumentation = new TestInstrumentationMasswrapFn({
+      enabled: false,
+    });
+
+    instrumentation.enable();
+    assert.deepEqual(exported.testFunction(), 'patched');
+    assert.deepEqual(exported.secondTestFunction(), 'patched');
+  });
+
+  it('should unwrap multiple functions with massunwrap', () => {
+    const instrumentation = new TestInstrumentationMasswrapFn({
+      enabled: false,
+    });
+
+    instrumentation.enable();
+    instrumentation.disable();
+    assert.deepEqual(exported.testFunction(), 'original');
+    assert.deepEqual(exported.secondTestFunction(), 'original');
+  });
 });