Skip to content

Commit

Permalink
fix(instrumentation-dataloder): Patch batchLoadFn without creating …
Browse files Browse the repository at this point in the history
…an instance
  • Loading branch information
onurtemizkan committed Oct 29, 2024
1 parent 9a20e15 commit 60d8499
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 47 deletions.
103 changes: 56 additions & 47 deletions plugins/node/instrumentation-dataloader/src/instrumentation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -91,64 +91,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;

Check warning on line 143 in plugins/node/instrumentation-dataloader/src/instrumentation.ts

View check run for this annotation

Codecov / codecov/patch

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

Added line #L143 was not covered by tests
}

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);

Check warning on line 153 in plugins/node/instrumentation-dataloader/src/instrumentation.ts

View check run for this annotation

Codecov / codecov/patch

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

Added line #L153 was not covered by tests
}

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;
Expand Down
25 changes: 25 additions & 0 deletions plugins/node/instrumentation-dataloader/test/dataloader.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand Down Expand Up @@ -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');
});
});

0 comments on commit 60d8499

Please sign in to comment.