From 4cbbf8622f3908325c19ec0d9f23a80a8ee8e9be Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 3 Jun 2025 10:45:50 -0400 Subject: [PATCH 01/13] remove async storage from mongo plugins --- .../src/mongodb-core.js | 73 ++++++++----------- .../datadog-instrumentations/src/mongodb.js | 22 +++--- .../datadog-plugin-mongodb-core/src/index.js | 8 +- 3 files changed, 47 insertions(+), 56 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongodb-core.js b/packages/datadog-instrumentations/src/mongodb-core.js index dd3ec6b2d6b..caa7177e070 100644 --- a/packages/datadog-instrumentations/src/mongodb-core.js +++ b/packages/datadog-instrumentations/src/mongodb-core.js @@ -2,8 +2,7 @@ const { channel, - addHook, - AsyncResource + addHook } = require('./helpers/instrument') const shimmer = require('../../datadog-shimmer') @@ -58,22 +57,6 @@ addHook({ name: 'mongodb-core', versions: ['~3.1.10'], file: 'lib/wireprotocol/2 return WireProtocol }) -addHook({ name: 'mongodb', versions: ['>=3.5.4 <4.11.0'], file: 'lib/utils.js' }, util => { - shimmer.wrap(util, 'maybePromise', maybePromise => function (parent, callback, fn) { - const asyncResource = new AsyncResource('bound-anonymous-fn') - const callbackIndex = arguments.length - 2 - - callback = arguments[callbackIndex] - - if (typeof callback === 'function') { - arguments[callbackIndex] = asyncResource.bind(callback) - } - - return maybePromise.apply(this, arguments) - }) - return util -}) - function wrapWp (wp) { shimmer.wrap(wp, 'command', command => wrapUnifiedCommand(command, 'command')) shimmer.wrap(wp, 'insert', insert => wrapUnifiedCommand(insert, 'insert', 'insert')) @@ -148,61 +131,69 @@ function wrapCommand (command, operation, name) { return wrapped } -function instrument (operation, command, ctx, args, server, ns, ops, options = {}) { +function instrument (operation, command, mongoCtx, args, server, ns, ops, options = {}) { const name = options.name || (ops && Object.keys(ops)[0]) const index = args.length - 1 let callback = args[index] - if (typeof callback !== 'function') return command.apply(ctx, args) + if (typeof callback !== 'function') return command.apply(mongoCtx, args) const serverInfo = server && server.s && server.s.options - const callbackResource = new AsyncResource('bound-anonymous-fn') - const asyncResource = new AsyncResource('bound-anonymous-fn') - - callback = callbackResource.bind(callback) - - return asyncResource.runInAsyncScope(() => { - startCh.publish({ ns, ops, options: serverInfo, name }) - args[index] = shimmer.wrapFunction(callback, callback => asyncResource.bind(function (err, res) { + const ctx = { + ns, + ops, + options: serverInfo, + name + } + return startCh.runStores(ctx, () => { + args[index] = shimmer.wrapFunction(callback, callback => function (err, res) { if (err) { - errorCh.publish(err) + ctx.error = err + errorCh.publish(ctx) } - finishCh.publish() + finishCh.publish(ctx) if (callback) { return callback.apply(this, arguments) } - })) + }) try { - return command.apply(ctx, args) + return command.apply(mongoCtx, args) } catch (err) { - errorCh.publish(err) + ctx.error = err + errorCh.publish(ctx) throw err } }) } -function instrumentPromise (operation, command, ctx, args, server, ns, ops, options = {}) { +function instrumentPromise (operation, command, mongoCtx, args, server, ns, ops, options = {}) { const name = options.name || (ops && Object.keys(ops)[0]) const serverInfo = server && server.s && server.s.options - const asyncResource = new AsyncResource('bound-anonymous-fn') - return asyncResource.runInAsyncScope(() => { - startCh.publish({ ns, ops, options: serverInfo, name }) + const ctx = { + ns, + ops, + options: serverInfo, + name + } - const promise = command.apply(ctx, args) + return startCh.runStores(ctx, () => { + const promise = command.apply(mongoCtx, args) return promise.then(function (res) { - finishCh.publish() + ctx.result = res + finishCh.publish(ctx) return res }, function (err) { - errorCh.publish(err) - finishCh.publish() + ctx.error = err + errorCh.publish(ctx) + finishCh.publish(ctx) throw err }) diff --git a/packages/datadog-instrumentations/src/mongodb.js b/packages/datadog-instrumentations/src/mongodb.js index f73aa21652a..dcc1b2ca96b 100644 --- a/packages/datadog-instrumentations/src/mongodb.js +++ b/packages/datadog-instrumentations/src/mongodb.js @@ -4,8 +4,7 @@ require('./mongodb-core') const { channel, - addHook, - AsyncResource + addHook } = require('./helpers/instrument') const shimmer = require('../../datadog-shimmer') @@ -41,19 +40,16 @@ addHook({ name: 'mongodb', versions: ['>=3.3 <5', '5', '>=6'] }, mongodb => { return method.apply(this, arguments) } - const asyncResource = new AsyncResource('bound-anonymous-fn') - - return asyncResource.runInAsyncScope(() => { - const filters = [arguments[0]] - if (useTwoArguments) { - filters.push(arguments[1]) - } + const ctx = { + filters: [arguments[0]], + methodName + } - startCh.publish({ - filters, - methodName - }) + if (useTwoArguments) { + ctx.filters.push(arguments[1]) + } + return startCh.runStores(ctx, () => { return method.apply(this, arguments) }) } diff --git a/packages/datadog-plugin-mongodb-core/src/index.js b/packages/datadog-plugin-mongodb-core/src/index.js index cd4b56109ea..efbbb8c4872 100644 --- a/packages/datadog-plugin-mongodb-core/src/index.js +++ b/packages/datadog-plugin-mongodb-core/src/index.js @@ -23,7 +23,9 @@ class MongodbCorePlugin extends DatabasePlugin { ) } - start ({ ns, ops, options = {}, name }) { + bindStart (ctx) { + const { ns, ops, options = {}, name } = ctx + // heartbeat commands can be disabled if this.config.heartbeatEnabled is false if (!this.config.heartbeatEnabled && isHeartbeat(ops, this.config)) { return @@ -43,11 +45,13 @@ class MongodbCorePlugin extends DatabasePlugin { 'out.host': options.host, 'out.port': options.port } - }) + }, ctx) const comment = this.injectDbmComment(span, ops.comment, service) if (comment) { ops.comment = comment } + + return ctx.currentStore } getPeerService (tags) { From fcc2b18ece0809e8ddc258f1d45f23431e2bf06b Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 3 Jun 2025 11:04:26 -0400 Subject: [PATCH 02/13] change mongoose to use asynstorage --- .../src/mongodb-core.js | 2 +- .../datadog-instrumentations/src/mongoose.js | 27 +++++++------------ 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongodb-core.js b/packages/datadog-instrumentations/src/mongodb-core.js index caa7177e070..e4d8df91acc 100644 --- a/packages/datadog-instrumentations/src/mongodb-core.js +++ b/packages/datadog-instrumentations/src/mongodb-core.js @@ -134,7 +134,7 @@ function wrapCommand (command, operation, name) { function instrument (operation, command, mongoCtx, args, server, ns, ops, options = {}) { const name = options.name || (ops && Object.keys(ops)[0]) const index = args.length - 1 - let callback = args[index] + const callback = args[index] if (typeof callback !== 'function') return command.apply(mongoCtx, args) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index 8116e38380e..4c2b9012310 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -2,7 +2,6 @@ const { addHook, channel } = require('./helpers/instrument') const { wrapThen } = require('./helpers/promise') -const { AsyncResource } = require('./helpers/instrument') const shimmer = require('../../datadog-shimmer') function wrapAddQueue (addQueue) { @@ -68,27 +67,21 @@ addHook({ return method.apply(this, arguments) } - const asyncResource = new AsyncResource('bound-anonymous-fn') - const filters = [arguments[0]] if (useTwoArguments) { filters.push(arguments[1]) } - const finish = asyncResource.bind(function () { - finishCh.publish() - }) - let callbackWrapped = false - const wrapCallbackIfExist = (args) => { + const wrapCallbackIfExist = (args, ctx) => { const lastArgumentIndex = args.length - 1 if (typeof args[lastArgumentIndex] === 'function') { // is a callback, wrap it to execute finish() shimmer.wrap(args, lastArgumentIndex, originalCb => { return function () { - finish() + finishCh.publish(ctx) return originalCb.apply(this, arguments) } @@ -98,13 +91,13 @@ addHook({ } } - wrapCallbackIfExist(arguments) + const ctx = { + filters, + methodName + } - return asyncResource.runInAsyncScope(() => { - startCh.publish({ - filters, - methodName - }) + return startCh.runStores(ctx, () => { + wrapCallbackIfExist(arguments, ctx) const res = method.apply(this, arguments) @@ -129,7 +122,7 @@ addHook({ const reject = arguments[1] arguments[0] = shimmer.wrapFunction(resolve, resolve => function wrappedResolve () { - finish() + finishCh.publish(ctx) if (resolve) { return resolve.apply(this, arguments) @@ -137,7 +130,7 @@ addHook({ }) arguments[1] = shimmer.wrapFunction(reject, reject => function wrappedReject () { - finish() + finishCh.publish(ctx) if (reject) { return reject.apply(this, arguments) From 455e222ee378d86afcd8d7dda5320d137b553509 Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 3 Jun 2025 14:54:41 -0400 Subject: [PATCH 03/13] try something --- .../datadog-instrumentations/src/mongoose.js | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index 4c2b9012310..329c0a27f9d 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -4,17 +4,17 @@ const { addHook, channel } = require('./helpers/instrument') const { wrapThen } = require('./helpers/promise') const shimmer = require('../../datadog-shimmer') -function wrapAddQueue (addQueue) { - return function addQueueWithTrace (name) { - if (typeof name === 'function') { - arguments[0] = AsyncResource.bind(name) - } else if (typeof this[name] === 'function') { - arguments[0] = AsyncResource.bind((...args) => this[name](...args)) - } - - return addQueue.apply(this, arguments) - } -} +// function wrapAddQueue (addQueue) { +// return function addQueueWithTrace (name) { +// if (typeof name === 'function') { +// arguments[0] = AsyncResource.bind(name) +// } else if (typeof this[name] === 'function') { +// arguments[0] = AsyncResource.bind((...args) => this[name](...args)) +// } + +// return addQueue.apply(this, arguments) +// } +// } addHook({ name: 'mongoose', @@ -25,7 +25,7 @@ addHook({ shimmer.wrap(mongoose.Promise.prototype, 'then', wrapThen) } - shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) + // shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) return mongoose }) From 7dd69ec9f569116d0f92578f62f5158461342727 Mon Sep 17 00:00:00 2001 From: William Conti Date: Mon, 23 Jun 2025 14:34:02 -0400 Subject: [PATCH 04/13] a few more fixes --- packages/datadog-instrumentations/src/mongodb-core.js | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongodb-core.js b/packages/datadog-instrumentations/src/mongodb-core.js index e4d8df91acc..8739993b957 100644 --- a/packages/datadog-instrumentations/src/mongodb-core.js +++ b/packages/datadog-instrumentations/src/mongodb-core.js @@ -153,11 +153,7 @@ function instrument (operation, command, mongoCtx, args, server, ns, ops, option errorCh.publish(ctx) } - finishCh.publish(ctx) - - if (callback) { - return callback.apply(this, arguments) - } + return callback ? finishCh.runStores(ctx, callback, this, ...arguments) : finishCh.publish(ctx) }) try { @@ -188,8 +184,9 @@ function instrumentPromise (operation, command, mongoCtx, args, server, ns, ops, return promise.then(function (res) { ctx.result = res - finishCh.publish(ctx) - return res + return finishCh.runStores(ctx, () => { + return res + }) }, function (err) { ctx.error = err errorCh.publish(ctx) From 6e560df456ff776fd799952eafddbfbf74b8345c Mon Sep 17 00:00:00 2001 From: William Conti Date: Wed, 9 Jul 2025 11:41:07 -0400 Subject: [PATCH 05/13] revert mongoose --- .../datadog-instrumentations/src/mongoose.js | 49 +++++++++++-------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index 329c0a27f9d..8116e38380e 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -2,19 +2,20 @@ const { addHook, channel } = require('./helpers/instrument') const { wrapThen } = require('./helpers/promise') +const { AsyncResource } = require('./helpers/instrument') const shimmer = require('../../datadog-shimmer') -// function wrapAddQueue (addQueue) { -// return function addQueueWithTrace (name) { -// if (typeof name === 'function') { -// arguments[0] = AsyncResource.bind(name) -// } else if (typeof this[name] === 'function') { -// arguments[0] = AsyncResource.bind((...args) => this[name](...args)) -// } +function wrapAddQueue (addQueue) { + return function addQueueWithTrace (name) { + if (typeof name === 'function') { + arguments[0] = AsyncResource.bind(name) + } else if (typeof this[name] === 'function') { + arguments[0] = AsyncResource.bind((...args) => this[name](...args)) + } -// return addQueue.apply(this, arguments) -// } -// } + return addQueue.apply(this, arguments) + } +} addHook({ name: 'mongoose', @@ -25,7 +26,7 @@ addHook({ shimmer.wrap(mongoose.Promise.prototype, 'then', wrapThen) } - // shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) + shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) return mongoose }) @@ -67,21 +68,27 @@ addHook({ return method.apply(this, arguments) } + const asyncResource = new AsyncResource('bound-anonymous-fn') + const filters = [arguments[0]] if (useTwoArguments) { filters.push(arguments[1]) } + const finish = asyncResource.bind(function () { + finishCh.publish() + }) + let callbackWrapped = false - const wrapCallbackIfExist = (args, ctx) => { + const wrapCallbackIfExist = (args) => { const lastArgumentIndex = args.length - 1 if (typeof args[lastArgumentIndex] === 'function') { // is a callback, wrap it to execute finish() shimmer.wrap(args, lastArgumentIndex, originalCb => { return function () { - finishCh.publish(ctx) + finish() return originalCb.apply(this, arguments) } @@ -91,13 +98,13 @@ addHook({ } } - const ctx = { - filters, - methodName - } + wrapCallbackIfExist(arguments) - return startCh.runStores(ctx, () => { - wrapCallbackIfExist(arguments, ctx) + return asyncResource.runInAsyncScope(() => { + startCh.publish({ + filters, + methodName + }) const res = method.apply(this, arguments) @@ -122,7 +129,7 @@ addHook({ const reject = arguments[1] arguments[0] = shimmer.wrapFunction(resolve, resolve => function wrappedResolve () { - finishCh.publish(ctx) + finish() if (resolve) { return resolve.apply(this, arguments) @@ -130,7 +137,7 @@ addHook({ }) arguments[1] = shimmer.wrapFunction(reject, reject => function wrappedReject () { - finishCh.publish(ctx) + finish() if (reject) { return reject.apply(this, arguments) From 680d0e1d4afa02ff2d2ecd2d17bd09c58d0aaa49 Mon Sep 17 00:00:00 2001 From: William Conti Date: Mon, 14 Jul 2025 13:46:38 -0400 Subject: [PATCH 06/13] fix mongoose --- .../datadog-instrumentations/src/mongoose.js | 82 +++++++++++++------ 1 file changed, 57 insertions(+), 25 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index 8116e38380e..14e1d1bc920 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -2,18 +2,36 @@ const { addHook, channel } = require('./helpers/instrument') const { wrapThen } = require('./helpers/promise') -const { AsyncResource } = require('./helpers/instrument') const shimmer = require('../../datadog-shimmer') +const { storage } = require('../../datadog-core') -function wrapAddQueue (addQueue) { - return function addQueueWithTrace (name) { - if (typeof name === 'function') { - arguments[0] = AsyncResource.bind(name) - } else if (typeof this[name] === 'function') { - arguments[0] = AsyncResource.bind((...args) => this[name](...args)) +function wrap (fn, store) { + if (typeof fn !== 'function') return fn + return function () { + storage('legacy').enterWith(store) + return fn.apply(this, arguments) + } +} + +function wrapExec (exec) { + return function (op, callback) { + if (typeof op === 'function') { + callback = op + op = undefined + } + + if (typeof callback === 'function') { + const store = storage('legacy').getStore() + const bound = wrap(callback, store) + + if (op === undefined) { + return exec.call(this, bound) + } else { + return exec.call(this, op, bound) + } } - return addQueue.apply(this, arguments) + return exec.apply(this, arguments) } } @@ -26,7 +44,27 @@ addHook({ shimmer.wrap(mongoose.Promise.prototype, 'then', wrapThen) } - shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) + if (mongoose.Query) { + shimmer.wrap(mongoose.Query.prototype, 'exec', wrapExec) + } + + if (mongoose.Aggregate) { + shimmer.wrap(mongoose.Aggregate.prototype, 'exec', wrapExec) + } + + // The `addQueue` function is used internally to execute buffered operations. + // We need to wrap it to make sure the context is propagated to the functions + // it executes. + if (mongoose.Collection) { + shimmer.wrap(mongoose.Collection.prototype, 'addQueue', function (addQueue) { + return function (name, args, options) { + if (typeof name === 'function') { + arguments[0] = wrap(name, storage('legacy').getStore()) + } + return addQueue.apply(this, arguments) + } + }) + } return mongoose }) @@ -68,27 +106,21 @@ addHook({ return method.apply(this, arguments) } - const asyncResource = new AsyncResource('bound-anonymous-fn') - const filters = [arguments[0]] if (useTwoArguments) { filters.push(arguments[1]) } - const finish = asyncResource.bind(function () { - finishCh.publish() - }) - let callbackWrapped = false - const wrapCallbackIfExist = (args) => { + const wrapCallbackIfExist = (args, ctx) => { const lastArgumentIndex = args.length - 1 if (typeof args[lastArgumentIndex] === 'function') { // is a callback, wrap it to execute finish() shimmer.wrap(args, lastArgumentIndex, originalCb => { return function () { - finish() + finishCh.publish(ctx) return originalCb.apply(this, arguments) } @@ -98,13 +130,13 @@ addHook({ } } - wrapCallbackIfExist(arguments) + const ctx = { + filters, + methodName + } - return asyncResource.runInAsyncScope(() => { - startCh.publish({ - filters, - methodName - }) + return startCh.runStores(ctx, () => { + wrapCallbackIfExist(arguments, ctx) const res = method.apply(this, arguments) @@ -129,7 +161,7 @@ addHook({ const reject = arguments[1] arguments[0] = shimmer.wrapFunction(resolve, resolve => function wrappedResolve () { - finish() + finishCh.publish(ctx) if (resolve) { return resolve.apply(this, arguments) @@ -137,7 +169,7 @@ addHook({ }) arguments[1] = shimmer.wrapFunction(reject, reject => function wrappedReject () { - finish() + finishCh.publish(ctx) if (reject) { return reject.apply(this, arguments) From 197c7469f7241a2ce1ec27b079d0099e549b8d60 Mon Sep 17 00:00:00 2001 From: William Conti Date: Mon, 14 Jul 2025 13:52:38 -0400 Subject: [PATCH 07/13] keep mongodb wrap queue function using async resource --- .../src/mongodb-core.js | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/packages/datadog-instrumentations/src/mongodb-core.js b/packages/datadog-instrumentations/src/mongodb-core.js index 8739993b957..628a3a54832 100644 --- a/packages/datadog-instrumentations/src/mongodb-core.js +++ b/packages/datadog-instrumentations/src/mongodb-core.js @@ -2,7 +2,8 @@ const { channel, - addHook + addHook, + AsyncResource } = require('./helpers/instrument') const shimmer = require('../../datadog-shimmer') @@ -57,6 +58,22 @@ addHook({ name: 'mongodb-core', versions: ['~3.1.10'], file: 'lib/wireprotocol/2 return WireProtocol }) +addHook({ name: 'mongodb', versions: ['>=3.5.4 <4.11.0'], file: 'lib/utils.js' }, util => { + shimmer.wrap(util, 'maybePromise', maybePromise => function (parent, callback, fn) { + const asyncResource = new AsyncResource('bound-anonymous-fn') + const callbackIndex = arguments.length - 2 + + callback = arguments[callbackIndex] + + if (typeof callback === 'function') { + arguments[callbackIndex] = asyncResource.bind(callback) + } + + return maybePromise.apply(this, arguments) + }) + return util +}) + function wrapWp (wp) { shimmer.wrap(wp, 'command', command => wrapUnifiedCommand(command, 'command')) shimmer.wrap(wp, 'insert', insert => wrapUnifiedCommand(insert, 'insert', 'insert')) From e8245f5d03cbfa80568e8d59076b12408b9a04e3 Mon Sep 17 00:00:00 2001 From: William Conti Date: Mon, 14 Jul 2025 19:58:16 -0400 Subject: [PATCH 08/13] fix lint --- packages/datadog-instrumentations/src/mongoose.js | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index 14e1d1bc920..f8ebb1268b7 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -24,11 +24,7 @@ function wrapExec (exec) { const store = storage('legacy').getStore() const bound = wrap(callback, store) - if (op === undefined) { - return exec.call(this, bound) - } else { - return exec.call(this, op, bound) - } + return op === undefined ? exec.call(this, bound) : exec.call(this, op, bound) } return exec.apply(this, arguments) From 5de74fdabb686ad758ad54fc05976b618489c2a0 Mon Sep 17 00:00:00 2001 From: William Conti Date: Mon, 14 Jul 2025 20:10:02 -0400 Subject: [PATCH 09/13] better naming --- packages/datadog-instrumentations/src/mongodb-core.js | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongodb-core.js b/packages/datadog-instrumentations/src/mongodb-core.js index 628a3a54832..51accdb1a2b 100644 --- a/packages/datadog-instrumentations/src/mongodb-core.js +++ b/packages/datadog-instrumentations/src/mongodb-core.js @@ -148,12 +148,12 @@ function wrapCommand (command, operation, name) { return wrapped } -function instrument (operation, command, mongoCtx, args, server, ns, ops, options = {}) { +function instrument (operation, command, instance, args, server, ns, ops, options = {}) { const name = options.name || (ops && Object.keys(ops)[0]) const index = args.length - 1 const callback = args[index] - if (typeof callback !== 'function') return command.apply(mongoCtx, args) + if (typeof callback !== 'function') return command.apply(instance, args) const serverInfo = server && server.s && server.s.options @@ -174,7 +174,7 @@ function instrument (operation, command, mongoCtx, args, server, ns, ops, option }) try { - return command.apply(mongoCtx, args) + return command.apply(instance, args) } catch (err) { ctx.error = err errorCh.publish(ctx) @@ -184,7 +184,7 @@ function instrument (operation, command, mongoCtx, args, server, ns, ops, option }) } -function instrumentPromise (operation, command, mongoCtx, args, server, ns, ops, options = {}) { +function instrumentPromise (operation, command, instance, args, server, ns, ops, options = {}) { const name = options.name || (ops && Object.keys(ops)[0]) const serverInfo = server && server.s && server.s.options @@ -197,7 +197,7 @@ function instrumentPromise (operation, command, mongoCtx, args, server, ns, ops, } return startCh.runStores(ctx, () => { - const promise = command.apply(mongoCtx, args) + const promise = command.apply(instance, args) return promise.then(function (res) { ctx.result = res From 0a871c4eee191d419af303c7cb82e341d0ca0ab0 Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 15 Jul 2025 10:37:58 -0400 Subject: [PATCH 10/13] fix reviewer comment --- packages/datadog-instrumentations/src/mongodb-core.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/datadog-instrumentations/src/mongodb-core.js b/packages/datadog-instrumentations/src/mongodb-core.js index 51accdb1a2b..4a27e919a3a 100644 --- a/packages/datadog-instrumentations/src/mongodb-core.js +++ b/packages/datadog-instrumentations/src/mongodb-core.js @@ -170,7 +170,7 @@ function instrument (operation, command, instance, args, server, ns, ops, option errorCh.publish(ctx) } - return callback ? finishCh.runStores(ctx, callback, this, ...arguments) : finishCh.publish(ctx) + return finishCh.runStores(ctx, callback, this, ...arguments) }) try { From b0bbf3b203c38eb4326bf6f682346a091f5f766b Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 15 Jul 2025 10:59:27 -0400 Subject: [PATCH 11/13] don't use store in instrumentation code --- .../datadog-instrumentations/src/mongoose.js | 37 ++++++++----------- 1 file changed, 15 insertions(+), 22 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index f8ebb1268b7..c37c581f5bb 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -3,13 +3,19 @@ const { addHook, channel } = require('./helpers/instrument') const { wrapThen } = require('./helpers/promise') const shimmer = require('../../datadog-shimmer') -const { storage } = require('../../datadog-core') -function wrap (fn, store) { - if (typeof fn !== 'function') return fn - return function () { - storage('legacy').enterWith(store) - return fn.apply(this, arguments) +const startCh = channel('datadog:mongoose:model:filter:start') +const finishCh = channel('datadog:mongoose:model:filter:finish') +// this channel is purely for wrapping the callback of exec methods and handling store context, it doesn't have any subscribers +const callbackCh = channel('datadog:mongoose:model:exec:callback') + + +function wrapAddQueue (addQueue) { + return function (name, args, options) { + if (typeof name === 'function') { + arguments[0] = callbackCh.runStores(ctx, addQueue, this, ...arguments) + } + return addQueue.apply(this, arguments) } } @@ -21,8 +27,8 @@ function wrapExec (exec) { } if (typeof callback === 'function') { - const store = storage('legacy').getStore() - const bound = wrap(callback, store) + const ctx = {} + const bound = callbackCh.runStores(ctx, callback, this, ...arguments) return op === undefined ? exec.call(this, bound) : exec.call(this, op, bound) } @@ -48,26 +54,13 @@ addHook({ shimmer.wrap(mongoose.Aggregate.prototype, 'exec', wrapExec) } - // The `addQueue` function is used internally to execute buffered operations. - // We need to wrap it to make sure the context is propagated to the functions - // it executes. if (mongoose.Collection) { - shimmer.wrap(mongoose.Collection.prototype, 'addQueue', function (addQueue) { - return function (name, args, options) { - if (typeof name === 'function') { - arguments[0] = wrap(name, storage('legacy').getStore()) - } - return addQueue.apply(this, arguments) - } - }) + shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) } return mongoose }) -const startCh = channel('datadog:mongoose:model:filter:start') -const finishCh = channel('datadog:mongoose:model:filter:finish') - const collectionMethodsWithFilter = [ 'count', 'countDocuments', From 6bb8d537ca23d974b9fcd3c76e600647a9bc03a1 Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 15 Jul 2025 11:14:16 -0400 Subject: [PATCH 12/13] fix lint --- packages/datadog-instrumentations/src/mongoose.js | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index c37c581f5bb..ea9220e01db 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -6,11 +6,10 @@ const shimmer = require('../../datadog-shimmer') const startCh = channel('datadog:mongoose:model:filter:start') const finishCh = channel('datadog:mongoose:model:filter:finish') -// this channel is purely for wrapping the callback of exec methods and handling store context, it doesn't have any subscribers +// this channel is for wrapping the callback of exec methods and handling store context, it doesn't have any subscribers const callbackCh = channel('datadog:mongoose:model:exec:callback') - -function wrapAddQueue (addQueue) { +function wrapAddQueue (addQueue, ctx) { return function (name, args, options) { if (typeof name === 'function') { arguments[0] = callbackCh.runStores(ctx, addQueue, this, ...arguments) @@ -55,7 +54,7 @@ addHook({ } if (mongoose.Collection) { - shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) + shimmer.wrap(mongoose.Collection.prototype, 'addQueue', addQueue => wrapAddQueue(addQueue, ctx)) } return mongoose From aade7008a2e6566d32584082b255a90eae8f37bf Mon Sep 17 00:00:00 2001 From: William Conti Date: Tue, 15 Jul 2025 11:27:52 -0400 Subject: [PATCH 13/13] fix lint error --- packages/datadog-instrumentations/src/mongoose.js | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js index ea9220e01db..41ee2f9fd87 100644 --- a/packages/datadog-instrumentations/src/mongoose.js +++ b/packages/datadog-instrumentations/src/mongoose.js @@ -9,9 +9,10 @@ const finishCh = channel('datadog:mongoose:model:filter:finish') // this channel is for wrapping the callback of exec methods and handling store context, it doesn't have any subscribers const callbackCh = channel('datadog:mongoose:model:exec:callback') -function wrapAddQueue (addQueue, ctx) { +function wrapAddQueue (addQueue) { return function (name, args, options) { if (typeof name === 'function') { + const ctx = {} arguments[0] = callbackCh.runStores(ctx, addQueue, this, ...arguments) } return addQueue.apply(this, arguments) @@ -54,7 +55,7 @@ addHook({ } if (mongoose.Collection) { - shimmer.wrap(mongoose.Collection.prototype, 'addQueue', addQueue => wrapAddQueue(addQueue, ctx)) + shimmer.wrap(mongoose.Collection.prototype, 'addQueue', wrapAddQueue) } return mongoose