Skip to content
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.

Sign up for GitHub

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

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix(instrumentation-dataloader): Patch batchLoadFn without creating an instance #2498

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
60d8499
fix(instrumentation-dataloder): Patch `batchLoadFn` without creating …
onurtemizkan Oct 23, 2024
117ba0b
Move `CustomDataLoader` class inside the test case.
onurtemizkan Nov 12, 2024
9b675f6
Address review comments.
onurtemizkan Nov 12, 2024
1554621
Merge branch 'main' into dataloader-custom-methods
onurtemizkan Nov 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 58 additions & 47 deletions plugins/node/instrumentation-dataloader/src/instrumentation.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,64 +91,75 @@
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()) {
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;
}
// BatchLoadFn is the first constructor argument
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There should be a guard that args[0] is actually a function before wrapping it. Something like (untested):

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:

foo.js

if (process.env.INSTRUMENT_IT) {
  const { DataloaderInstrumentation } = require('./');
  const { NodeSDK, tracing } = require('@opentelemetry/sdk-node');
  const instr = new DataloaderInstrumentation();
  const sdk = new NodeSDK({
    spanProcessor: new tracing.SimpleSpanProcessor(new tracing.ConsoleSpanExporter()),
    serviceName: 'foo',
    instrumentations: [instr]
  });
  sdk.start();
}

const Dataloader = require('dataloader');
const dl = new Dataloader(); // oops, no batchLoadFn passed in

console.log('the end');

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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a guard to check if batchLoadFn is a function. However, I could not reproduce the problem where Dataloader instance is created without a batchLoadFn (with or without the instrumentation).

It fails with:

TypeError: DataLoader must be constructed with a function which accepts Array<key> and returns Promise<Array<value>>, but got: undefined.
for both cases.

// https://github.com/graphql/dataloader/blob/77c2cd7ca97e8795242018ebc212ce2487e729d2/src/index.js#L47
if (typeof args[0] === 'function') {
if (isWrapped(args[0])) {
instrumentation._unwrap(args, 0);

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

View check run for this annotation

Codecov / codecov/patch

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

Added line #L154 was not covered by tests
}

if (isWrapped(inst._batchLoadFn)) {
instrumentation._unwrap(inst, '_batchLoadFn');
args[0] = instrumentation._wrapBatchLoadFn(
args[0]
) as Dataloader.BatchLoadFn<unknown, unknown>;
}

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

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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ extraInstrumentation.disable();

import * as assert from 'assert';
import * as Dataloader from 'dataloader';
import * as crypto from 'crypto';

describe('DataloaderInstrumentation', () => {
let dataloader: Dataloader<string, number>;
Expand Down Expand Up @@ -334,4 +335,28 @@ describe('DataloaderInstrumentation', () => {
assert.deepStrictEqual(await alternativeDataloader.loadMany(['test']), [1]);
assert.strictEqual(memoryExporter.getFinishedSpans().length, 5);
});

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

assert.strictEqual(memoryExporter.getFinishedSpans().length, 2);
const [batchSpan, loadSpan] = memoryExporter.getFinishedSpans();

assert.strictEqual(loadSpan.name, 'dataloader.load');
assert.strictEqual(batchSpan.name, 'dataloader.batch');
});
});