From 60d8499cbf6ac75d5d5f6a7bf04c1c9b502791a6 Mon Sep 17 00:00:00 2001 From: Onur Temizkan Date: Wed, 23 Oct 2024 18:21:44 +0100 Subject: [PATCH 1/3] fix(instrumentation-dataloder): Patch `batchLoadFn` without creating an instance --- .../src/instrumentation.ts | 103 ++++++++++-------- .../test/dataloader.test.ts | 25 +++++ 2 files changed, 81 insertions(+), 47 deletions(-) diff --git a/plugins/node/instrumentation-dataloader/src/instrumentation.ts b/plugins/node/instrumentation-dataloader/src/instrumentation.ts index 21a192484f..7f6e51c58c 100644 --- a/plugins/node/instrumentation-dataloader/src/instrumentation.ts +++ b/plugins/node/instrumentation-dataloader/src/instrumentation.ts @@ -91,64 +91,73 @@ export class DataloaderInstrumentation extends InstrumentationBase + ): Dataloader.BatchLoadFn { + const instrumentation = this; + + return function patchedBatchLoadFn( + this: DataloaderInternal, + ...args: Parameters> + ) { + 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) + .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 ) { - 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> - ) { - 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) - .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; - return inst; + return constructor.apply(this, args); } PatchedDataloader.prototype = prototype; diff --git a/plugins/node/instrumentation-dataloader/test/dataloader.test.ts b/plugins/node/instrumentation-dataloader/test/dataloader.test.ts index 9686c21d2e..5a72827ed3 100644 --- a/plugins/node/instrumentation-dataloader/test/dataloader.test.ts +++ b/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; let contextManager: AsyncHooksContextManager; + class CustomDataLoader extends Dataloader { + 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'); + }); }); From 117ba0bff29e8067c0d088901598a263a7adf6c1 Mon Sep 17 00:00:00 2001 From: Onur Temizkan Date: Tue, 12 Nov 2024 09:11:35 +0000 Subject: [PATCH 2/3] Move `CustomDataLoader` class inside the test case. --- .../test/dataloader.test.ts | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/plugins/node/instrumentation-dataloader/test/dataloader.test.ts b/plugins/node/instrumentation-dataloader/test/dataloader.test.ts index 5a72827ed3..62f0b47337 100644 --- a/plugins/node/instrumentation-dataloader/test/dataloader.test.ts +++ b/plugins/node/instrumentation-dataloader/test/dataloader.test.ts @@ -33,23 +33,10 @@ 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; let contextManager: AsyncHooksContextManager; - class CustomDataLoader extends Dataloader { - 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'); @@ -350,6 +337,19 @@ describe('DataloaderInstrumentation', () => { }); it('should not prune custom methods', async () => { + const getMd5HashFromIdx = (idx: number) => + crypto.createHash('md5').update(String(idx)).digest('hex'); + + class CustomDataLoader extends Dataloader { + constructor() { + super(async keys => keys.map((_, idx) => getMd5HashFromIdx(idx))); + } + + public async customLoad() { + return this.load('test'); + } + } + const customDataloader = new CustomDataLoader(); await customDataloader.customLoad(); From 9b675f6ede31a4ad7a2ca2d370f12c5755f3bf0e Mon Sep 17 00:00:00 2001 From: Onur Temizkan Date: Tue, 12 Nov 2024 09:37:42 +0000 Subject: [PATCH 3/3] Address review comments. --- .../src/instrumentation.ts | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/plugins/node/instrumentation-dataloader/src/instrumentation.ts b/plugins/node/instrumentation-dataloader/src/instrumentation.ts index 7f6e51c58c..c97eca3edb 100644 --- a/plugins/node/instrumentation-dataloader/src/instrumentation.ts +++ b/plugins/node/instrumentation-dataloader/src/instrumentation.ts @@ -139,7 +139,7 @@ export class DataloaderInstrumentation extends InstrumentationBase; + args[0] = instrumentation._wrapBatchLoadFn( + args[0] + ) as Dataloader.BatchLoadFn; + } return constructor.apply(this, args); }