-
Notifications
You must be signed in to change notification settings - Fork 534
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
fix(instrumentation-dataloader): Patch batchLoadFn
without creating an instance
#2498
base: main
Are you sure you want to change the base?
Changes from 1 commit
60d8499
117ba0b
9b675f6
1554621
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -91,64 +91,73 @@ | |
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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There should be a guard that if (typeof args[0] !== 'function') {
return constructor.apply(this, args);
} or perhaps only do the unwrap and wrap if it is a function. The current behaviour breaks the guard in Dataloader itself:
Run that without and with instrumentation shows a changed behaviour: % node foo.js
/Users/trentm/tm/opentelemetry-js-contrib10/node_modules/dataloader/index.js:32
throw new TypeError('DataLoader must be constructed with a function which accepts ' + ("Array<key> and returns Promise<Array<value>>, but got: " + batchLoadFn + "."));
^
TypeError: DataLoader must be constructed with a function which accepts Array<key> and returns Promise<Array<value>>, but got: undefined.
at new DataLoader (/Users/trentm/tm/opentelemetry-js-contrib10/node_modules/dataloader/index.js:32:13)
at Object.<anonymous> (/Users/trentm/tm/opentelemetry-js-contrib10/plugins/node/instrumentation-dataloader/foo.js:14:12)
at Module._compile (node:internal/modules/cjs/loader:1364:14)
at Module._extensions..js (node:internal/modules/cjs/loader:1422:10)
at Module.load (node:internal/modules/cjs/loader:1203:32)
at Module._load (node:internal/modules/cjs/loader:1019:12)
at Function.executeUserEntryPoint [as runMain] (node:internal/modules/run_main:128:12)
at node:internal/main/run_main_module:28:49
Node.js v18.20.4
% INSTRUMENT_IT=true node foo.js
the end Please add a test case for this. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I added a guard to check if It fails with:
|
||
// 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; | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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'); | ||
} | ||
} | ||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: I think this class and the getMd5HashFromIdx function could move inside the test case to be closer to the only place they are used: --- 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<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');
@@ -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<string, string> {
+ constructor() {
+ super(async keys => keys.map((_, idx) => getMd5HashFromIdx(idx)));
+ }
+
+ public async customLoad() {
+ return this.load('test');
+ }
+ }
+
const customDataloader = new CustomDataLoader();
await customDataloader.customLoad();
|
||
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'); | ||
}); | ||
}); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This if-block is at init-time before any spans are being created. This
!instrumentation.shouldCreateSpans()
is only true by fluke becauseconfig.requireParentSpan
is undefined in the tests.