fix(instrumentation-dataloder): Patch batchLoadFn without creating an instance

Author: onurtemizkan

plugins/node/instrumentation-dataloader/src/instrumentation.ts

@@ -90,64 +90,73 @@ export class DataloaderInstrumentation extends InstrumentationBase<DataloaderIns
     return `${MODULE_NAME}.${operation} ${dataloaderName}`;
   }
 
+  private _wrapBatchLoadFn(
+    batchLoadFn: Dataloader.BatchLoadFn<unknown, unknown>
+  ): Dataloader.BatchLoadFn<unknown, unknown> {
+    const instrumentation = this;
+
+    return function patchedBatchLoadFn(
+      this: DataloaderInternal,
+      ...args: Parameters<Dataloader.BatchLoadFn<unknown, unknown>>
+    ) {
+      if (
+        !instrumentation.isEnabled() ||
+        !instrumentation.shouldCreateSpans()
+      ) {
+        return batchLoadFn.call(this, ...args);
+      }
+
+      const parent = context.active();
+      const span = instrumentation.tracer.startSpan(
+        instrumentation.getSpanName(this, 'batch'),
+        { links: this._batch?.spanLinks as Link[] | undefined },
+        parent
+      );
+
+      return context.with(trace.setSpan(parent, span), () => {
+        return (batchLoadFn.apply(this, args) as Promise<unknown[]>)
+          .then(value => {
+            span.end();
+            return value;
+          })
+          .catch(err => {
+            span.recordException(err);
+            span.setStatus({
+              code: SpanStatusCode.ERROR,
+              message: err.message,
+            });
+            span.end();
+            throw err;
+          });
+      });
+    };
+  }
+
   private _getPatchedConstructor(
     constructor: typeof Dataloader
   ): typeof Dataloader {
-    const prototype = constructor.prototype;
     const instrumentation = this;
+    const prototype = constructor.prototype;
+
+    if (!instrumentation.isEnabled() || !instrumentation.shouldCreateSpans()) {
+      return constructor;
+    }
 
     function PatchedDataloader(
+      this: DataloaderInternal,
       ...args: ConstructorParameters<typeof constructor>
     ) {
-      const inst = new constructor(...args) as DataloaderInternal;
-
-      if (!instrumentation.isEnabled()) {
-        return inst;
-      }
-
-      if (isWrapped(inst._batchLoadFn)) {
-        instrumentation._unwrap(inst, '_batchLoadFn');
+      // BatchLoadFn is the first constructor argument
+      // https://github.com/graphql/dataloader/blob/77c2cd7ca97e8795242018ebc212ce2487e729d2/src/index.js#L47
+      if (isWrapped(args[0])) {
+        instrumentation._unwrap(args, 0);
       }
 
-      instrumentation._wrap(inst, '_batchLoadFn', original => {
-        return function patchedBatchLoadFn(
-          this: DataloaderInternal,
-          ...args: Parameters<Dataloader.BatchLoadFn<unknown, unknown>>
-        ) {
-          if (
-            !instrumentation.isEnabled() ||
-            !instrumentation.shouldCreateSpans()
-          ) {
-            return original.call(this, ...args);
-          }
-
-          const parent = context.active();
-          const span = instrumentation.tracer.startSpan(
-            instrumentation.getSpanName(inst, 'batch'),
-            { links: this._batch?.spanLinks as Link[] | undefined },
-            parent
-          );
-
-          return context.with(trace.setSpan(parent, span), () => {
-            return (original.apply(this, args) as Promise<unknown[]>)
-              .then(value => {
-                span.end();
-                return value;
-              })
-              .catch(err => {
-                span.recordException(err);
-                span.setStatus({
-                  code: SpanStatusCode.ERROR,
-                  message: err.message,
-                });
-                span.end();
-                throw err;
-              });
-          });
-        };
-      });
+      args[0] = instrumentation._wrapBatchLoadFn(
+        args[0]
+      ) as Dataloader.BatchLoadFn<unknown, unknown>;
 
-      return inst;
+      return constructor.apply(this, args);
     }
 
     PatchedDataloader.prototype = prototype;

plugins/node/instrumentation-dataloader/test/dataloader.test.ts

@@ -31,11 +31,25 @@ extraInstrumentation.disable();
 
 import * as assert from 'assert';
 import * as Dataloader from 'dataloader';
+import * as crypto from 'crypto';
+
+const getMd5HashFromIdx = (idx: number) =>
+  crypto.createHash('md5').update(String(idx)).digest('hex');
 
 describe('DataloaderInstrumentation', () => {
   let dataloader: Dataloader<string, number>;
   let contextManager: AsyncHooksContextManager;
 
+  class CustomDataLoader extends Dataloader<string, string> {
+    constructor() {
+      super(async keys => keys.map((_, idx) => getMd5HashFromIdx(idx)));
+    }
+
+    public async customLoad() {
+      return this.load('test');
+    }
+  }
+
   const memoryExporter = new InMemorySpanExporter();
   const provider = new NodeTracerProvider();
   const tracer = provider.getTracer('default');
@@ -334,4 +348,15 @@ describe('DataloaderInstrumentation', () => {
     assert.deepStrictEqual(await alternativeDataloader.loadMany(['test']), [1]);
     assert.strictEqual(memoryExporter.getFinishedSpans().length, 5);
   });
+
+  it('should not prune custom methods', async () => {
+    const customDataloader = new CustomDataLoader();
+    await customDataloader.customLoad();
+
+    assert.strictEqual(memoryExporter.getFinishedSpans().length, 2);
+    const [batchSpan, loadSpan] = memoryExporter.getFinishedSpans();
+
+    assert.strictEqual(loadSpan.name, 'dataloader.load');
+    assert.strictEqual(batchSpan.name, 'dataloader.batch');
+  });
 });