Skip to content

Commit 875e585

Browse files
BridgeARBridgeAR
authored
fix: wrap mysql2 addCommand only once (#7459)
* fix: wrap mysql2 addCommand only once Prepare calls would currently cause the method to be rewrapped frequently. This is now only done a single time. In addition this adds types to the instrumentation and adds a few more safe guards for type safety. * perf: improve database plugin performance The database plugin has a few intermediate objects that are not necessary as well as calling methods when they do not have to be called. This is not a huge difference but it is not needed. * chore: make private methods actually private * chore: add types to tracing and outbound plugins * test: add regression test for single addCommand wrapping Fixes: #7074
1 parent e130d50 commit 875e585

5 files changed

Lines changed: 286 additions & 114 deletions

File tree

‎packages/datadog-instrumentations/src/mysql2.js‎

Lines changed: 131 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,22 @@ const shimmer = require('../../datadog-shimmer')
66
const satisfies = require('../../../vendor/dist/semifies')
77
const { channel, addHook } = require('./helpers/instrument')
88

9+
/** @type {WeakMap<object, Function>} */
10+
const wrappedOnResult = new WeakMap()
11+
12+
/**
13+
* @param {unknown} sql
14+
* @returns {string|undefined}
15+
*/
16+
function resolveSqlString (sql) {
17+
return typeof sql === 'string' ? sql : /** @type {{ sql?: string }} */ (sql)?.sql
18+
}
19+
20+
/**
21+
* @param {Function} Connection
22+
* @param {string} version
23+
* @returns {Function}
24+
*/
925
function wrapConnection (Connection, version) {
1026
const startCh = channel('apm:mysql2:query:start')
1127
const finishCh = channel('apm:mysql2:query:finish')
@@ -19,28 +35,39 @@ function wrapConnection (Connection, version) {
1935
shimmer.wrap(Connection.prototype, 'addCommand', addCommand => function (cmd) {
2036
if (!startCh.hasSubscribers) return addCommand.apply(this, arguments)
2137

22-
const name = cmd && cmd.constructor && cmd.constructor.name
23-
const isCommand = typeof cmd.execute === 'function'
24-
const isQuery = isCommand && (name === 'Execute' || name === 'Query')
38+
const command = /** @type {{ execute?: Function, constructor?: { name?: string } }} */ (cmd)
39+
if (typeof command.execute !== 'function') return addCommand.apply(this, arguments)
40+
41+
const name = command.constructor?.name
42+
const isQuery = name === 'Execute' || name === 'Query'
2543
const ctx = {}
2644

27-
// TODO: consider supporting all commands and not just queries
28-
cmd.execute = isQuery
29-
? wrapExecute(cmd, cmd.execute, ctx, this.config)
30-
: bindExecute(cmd.execute, ctx)
45+
if (isQuery) {
46+
command.execute = wrapExecute(command, command.execute, ctx, this.config)
47+
48+
return commandAddCh.runStores(ctx, addCommand, this, ...arguments)
49+
}
50+
51+
wrapCommandOnResult(command, ctx)
52+
53+
command.execute = shimmer.wrapFunction(
54+
command.execute,
55+
execute => function executeWithTrace (_packet_, _connection_) {
56+
return commandStartCh.runStores(ctx, execute, this, ...arguments)
57+
}
58+
)
3159

3260
return commandAddCh.runStores(ctx, addCommand, this, ...arguments)
3361
})
3462

3563
shimmer.wrap(Connection.prototype, 'query', query => function (sql, values, cb) {
3664
if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments)
3765

38-
if (sql !== null && typeof sql === 'object') sql = sql.sql
39-
40-
if (!sql) return query.apply(this, arguments)
66+
const resolvedSql = resolveSqlString(sql)
67+
if (resolvedSql === undefined) return query.apply(this, arguments)
4168

4269
const abortController = new AbortController()
43-
startOuterQueryCh.publish({ sql, abortController })
70+
startOuterQueryCh.publish({ sql: resolvedSql, abortController })
4471

4572
if (abortController.signal.aborted) {
4673
const addCommand = this.addCommand
@@ -56,7 +83,7 @@ function wrapConnection (Connection, version) {
5683
cb = queryCommand.onResult
5784

5885
process.nextTick(() => {
59-
if (cb) {
86+
if (typeof cb === 'function') {
6087
cb(abortController.signal.reason)
6188
} else {
6289
queryCommand.emit('error', abortController.signal.reason)
@@ -76,12 +103,11 @@ function wrapConnection (Connection, version) {
76103
shimmer.wrap(Connection.prototype, 'execute', execute => function (sql, values, cb) {
77104
if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments)
78105

79-
if (sql !== null && typeof sql === 'object') sql = sql.sql
80-
81-
if (!sql) return execute.apply(this, arguments)
106+
const resolvedSql = resolveSqlString(sql)
107+
if (resolvedSql === undefined) return execute.apply(this, arguments)
82108

83109
const abortController = new AbortController()
84-
startOuterQueryCh.publish({ sql, abortController })
110+
startOuterQueryCh.publish({ sql: resolvedSql, abortController })
85111

86112
if (abortController.signal.aborted) {
87113
const addCommand = this.addCommand
@@ -94,7 +120,9 @@ function wrapConnection (Connection, version) {
94120
this.addCommand = addCommand
95121
}
96122

97-
result?.onResult(abortController.signal.reason)
123+
if (typeof result?.onResult === 'function') {
124+
result.onResult(abortController.signal.reason)
125+
}
98126

99127
return result
100128
}
@@ -104,33 +132,48 @@ function wrapConnection (Connection, version) {
104132

105133
return Connection
106134

107-
function bindExecute (execute, ctx) {
108-
return shimmer.wrapFunction(execute, execute => function executeWithTrace (packet, connection) {
109-
const onResult = this.onResult
135+
/**
136+
* @param {object} cmd
137+
* @param {object} ctx
138+
* @returns {void}
139+
*/
140+
function wrapCommandOnResult (cmd, ctx) {
141+
const onResult = cmd?.onResult
142+
if (typeof onResult !== 'function') return
110143

111-
if (onResult) {
112-
this.onResult = function () {
113-
return commandFinishCh.runStores(ctx, onResult, this, ...arguments)
114-
}
115-
}
144+
const cached = wrappedOnResult.get(cmd)
116145

117-
return commandStartCh.runStores(ctx, execute, this, ...arguments)
118-
})
146+
if (cached === onResult) return
147+
148+
const wrapped = function () {
149+
return commandFinishCh.runStores(ctx, onResult, this, ...arguments)
150+
}
151+
152+
wrappedOnResult.set(cmd, wrapped)
153+
cmd.onResult = wrapped
119154
}
120155

156+
/**
157+
* @param {object} cmd
158+
* @param {Function} execute
159+
* @param {object} ctx
160+
* @param {object} config
161+
* @returns {Function}
162+
*/
121163
function wrapExecute (cmd, execute, ctx, config) {
122164
return shimmer.wrapFunction(execute, execute => function executeWithTrace (packet, connection) {
123-
ctx.sql = cmd.statement ? cmd.statement.query : cmd.sql
165+
const command = /** @type {{ statement?: { query?: unknown }, sql?: unknown }} */ (cmd)
166+
ctx.sql = command.statement ? command.statement.query : command.sql
124167
ctx.conf = config
125168

126169
return startCh.runStores(ctx, () => {
127-
if (cmd.statement) {
128-
cmd.statement.query = ctx.sql
170+
if (command.statement) {
171+
command.statement.query = ctx.sql
129172
} else {
130-
cmd.sql = ctx.sql
173+
command.sql = ctx.sql
131174
}
132175

133-
if (this.onResult) {
176+
if (typeof this.onResult === 'function') {
134177
const onResult = this.onResult
135178

136179
this.onResult = shimmer.wrapFunction(onResult, onResult => function (error) {
@@ -141,11 +184,14 @@ function wrapConnection (Connection, version) {
141184
finishCh.runStores(ctx, onResult, this, ...arguments)
142185
})
143186
} else {
144-
this.once(errorMonitor, error => {
145-
ctx.error = error
146-
errorCh.publish(ctx)
147-
})
148-
this.once('end', () => finishCh.publish(ctx))
187+
const command = /** @type {{ once?: Function }} */ (this)
188+
if (typeof command.once === 'function') {
189+
command.once(errorMonitor, error => {
190+
ctx.error = error
191+
errorCh.publish(ctx)
192+
})
193+
command.once('end', () => finishCh.publish(ctx))
194+
}
149195
}
150196

151197
this.execute = execute
@@ -157,22 +203,26 @@ function wrapConnection (Connection, version) {
157203
errorCh.publish(ctx)
158204
}
159205
})
160-
}, cmd)
206+
})
161207
}
162208
}
209+
/**
210+
* @param {Function} Pool
211+
* @param {string} version
212+
* @returns {Function}
213+
*/
163214
function wrapPool (Pool, version) {
164215
const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
165216
const shouldEmitEndAfterQueryAbort = satisfies(version, '>=1.3.3')
166217

167218
shimmer.wrap(Pool.prototype, 'query', query => function (sql, values, cb) {
168219
if (!startOuterQueryCh.hasSubscribers) return query.apply(this, arguments)
169220

170-
if (sql !== null && typeof sql === 'object') sql = sql.sql
171-
172-
if (!sql) return query.apply(this, arguments)
221+
const resolvedSql = resolveSqlString(sql)
222+
if (resolvedSql === undefined) return query.apply(this, arguments)
173223

174224
const abortController = new AbortController()
175-
startOuterQueryCh.publish({ sql, abortController })
225+
startOuterQueryCh.publish({ sql: resolvedSql, abortController })
176226

177227
if (abortController.signal.aborted) {
178228
const getConnection = this.getConnection
@@ -206,21 +256,22 @@ function wrapPool (Pool, version) {
206256
shimmer.wrap(Pool.prototype, 'execute', execute => function (sql, values, cb) {
207257
if (!startOuterQueryCh.hasSubscribers) return execute.apply(this, arguments)
208258

209-
if (sql !== null && typeof sql === 'object') sql = sql.sql
210-
211-
if (!sql) return execute.apply(this, arguments)
259+
const resolvedSql = resolveSqlString(sql)
260+
if (resolvedSql === undefined) return execute.apply(this, arguments)
212261

213262
const abortController = new AbortController()
214-
startOuterQueryCh.publish({ sql, abortController })
263+
startOuterQueryCh.publish({ sql: resolvedSql, abortController })
215264

216265
if (abortController.signal.aborted) {
217266
if (typeof values === 'function') {
218267
cb = values
219268
}
220269

221-
process.nextTick(() => {
222-
cb(abortController.signal.reason)
223-
})
270+
if (typeof cb === 'function') {
271+
process.nextTick(() => {
272+
/** @type {Function} */ (cb)(abortController.signal.reason)
273+
})
274+
}
224275
return
225276
}
226277

@@ -230,6 +281,10 @@ function wrapPool (Pool, version) {
230281
return Pool
231282
}
232283

284+
/**
285+
* @param {Function} PoolCluster
286+
* @returns {Function}
287+
*/
233288
function wrapPoolCluster (PoolCluster) {
234289
const startOuterQueryCh = channel('datadog:mysql2:outerquery:start')
235290
const wrappedPoolNamespaces = new WeakSet()
@@ -239,12 +294,11 @@ function wrapPoolCluster (PoolCluster) {
239294

240295
if (startOuterQueryCh.hasSubscribers && !wrappedPoolNamespaces.has(poolNamespace)) {
241296
shimmer.wrap(poolNamespace, 'query', query => function (sql, values, cb) {
242-
if (sql !== null && typeof sql === 'object') sql = sql.sql
243-
244-
if (!sql) return query.apply(this, arguments)
297+
const resolvedSql = resolveSqlString(sql)
298+
if (resolvedSql === undefined) return query.apply(this, arguments)
245299

246300
const abortController = new AbortController()
247-
startOuterQueryCh.publish({ sql, abortController })
301+
startOuterQueryCh.publish({ sql: resolvedSql, abortController })
248302

249303
if (abortController.signal.aborted) {
250304
const getConnection = this.getConnection
@@ -274,21 +328,22 @@ function wrapPoolCluster (PoolCluster) {
274328
})
275329

276330
shimmer.wrap(poolNamespace, 'execute', execute => function (sql, values, cb) {
277-
if (sql !== null && typeof sql === 'object') sql = sql.sql
278-
279-
if (!sql) return execute.apply(this, arguments)
331+
const resolvedSql = resolveSqlString(sql)
332+
if (resolvedSql === undefined) return execute.apply(this, arguments)
280333

281334
const abortController = new AbortController()
282-
startOuterQueryCh.publish({ sql, abortController })
335+
startOuterQueryCh.publish({ sql: resolvedSql, abortController })
283336

284337
if (abortController.signal.aborted) {
285338
if (typeof values === 'function') {
286339
cb = values
287340
}
288341

289-
process.nextTick(() => {
290-
cb(abortController.signal.reason)
291-
})
342+
if (typeof cb === 'function') {
343+
process.nextTick(() => {
344+
/** @type {Function} */ (cb)(abortController.signal.reason)
345+
})
346+
}
292347

293348
return
294349
}
@@ -305,9 +360,21 @@ function wrapPoolCluster (PoolCluster) {
305360
return PoolCluster
306361
}
307362

308-
addHook({ name: 'mysql2', file: 'lib/base/connection.js', versions: ['>=3.11.5'] }, wrapConnection)
309-
addHook({ name: 'mysql2', file: 'lib/connection.js', versions: ['1 - 3.11.4'] }, wrapConnection)
310-
addHook({ name: 'mysql2', file: 'lib/pool.js', versions: ['1 - 3.11.4'] }, wrapPool)
363+
addHook(
364+
{ name: 'mysql2', file: 'lib/base/connection.js', versions: ['>=3.11.5'] },
365+
/** @type {(moduleExports: unknown, version: string) => unknown} */ (wrapConnection)
366+
)
367+
addHook(
368+
{ name: 'mysql2', file: 'lib/connection.js', versions: ['1 - 3.11.4'] },
369+
/** @type {(moduleExports: unknown, version: string) => unknown} */ (wrapConnection)
370+
)
371+
addHook(
372+
{ name: 'mysql2', file: 'lib/pool.js', versions: ['1 - 3.11.4'] },
373+
/** @type {(moduleExports: unknown, version: string) => unknown} */ (wrapPool)
374+
)
311375

312376
// PoolNamespace.prototype.query does not exist in mysql2<2.3.0
313-
addHook({ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['2.3.0 - 3.11.4'] }, wrapPoolCluster)
377+
addHook(
378+
{ name: 'mysql2', file: 'lib/pool_cluster.js', versions: ['2.3.0 - 3.11.4'] },
379+
/** @type {(moduleExports: unknown, version: string) => unknown} */ (wrapPoolCluster)
380+
)

0 commit comments

Comments
 (0)