dd-trace 6.10.0 → 6.12.0

package/packages/datadog-instrumentations/src/mariadb.js

@@ -1,7 +1,19 @@
 'use strict'
 
+const { performance } = require('node:perf_hooks')
+
 const shimmer = require('../../datadog-shimmer')
 const { channel, addHook } = require('./helpers/instrument')
+const {
+  acquireWait,
+  clearPoolWaitTime,
+  isPoolQueryAcquire,
+  runOutsidePoolQueryAcquire,
+  runPoolAcquireError,
+  runWithPoolWait,
+  takePoolWaitTime,
+  wrapPoolQueryMethod,
+} = require('./helpers/pool-acquire')
 
 const commandAddCh = channel('apm:mariadb:command:add')
 const connectionStartCh = channel('apm:mariadb:connection:start')
@@ -10,11 +22,37 @@ const startCh = channel('apm:mariadb:query:start')
 const finishCh = channel('apm:mariadb:query:finish')
 const errorCh = channel('apm:mariadb:query:error')
 const skipCh = channel('apm:mariadb:pool:skip')
+const acquireStartCh = channel('apm:mariadb:pool:acquire:start')
+const acquireFinishCh = channel('apm:mariadb:pool:acquire:finish')
+const poolAcquireChannels = {
+  connectionFinishCh,
+  acquireStartCh,
+  acquireFinishCh,
+}
+const wrappedPools = new WeakSet()
+
+/**
+ * @typedef {object} PoolAcquireTiming
+ * @property {number} [poolAcquireStart]
+ * @property {number} [poolAcquireStartedAt]
+ * @property {number} [poolWaitTime]
+ */
+
+/** @typedef {{ length: number, [index: number]: unknown } & Iterable<unknown>} ArgumentsLike */
+
+/** @type {PoolAcquireTiming|undefined} */
+let currentPoolAcquireTiming
+let explicitPoolAcquire = false
 
 function wrapCommandStart (start, ctx) {
   return shimmer.wrapFunction(start, start => function (...args) {
     if (!startCh.hasSubscribers) return start.apply(this, args)
 
+    const poolWaitTime = takePoolWaitTime(args[0])
+    if (poolWaitTime !== undefined) {
+      ctx.poolWaitTime = poolWaitTime
+    }
+
     const { reject, resolve } = this
     shimmer.wrap(this, 'resolve', function wrapResolve () {
       return function (...args) {
@@ -133,43 +171,262 @@ function wrapPoolMethod (createConnection) {
   }
 }
 
+/**
+ * @param {Function} getConnection
+ * @returns {Function}
+ */
+function wrapExplicitPoolGetConnection (getConnection) {
+  return function wrappedExplicitPoolGetConnection () {
+    if (!connectionStartCh.hasSubscribers) {
+      return getConnection.apply(this, arguments)
+    }
+
+    const previous = explicitPoolAcquire
+    explicitPoolAcquire = true
+    try {
+      return runOutsidePoolQueryAcquire(getConnection, this, arguments)
+    } finally {
+      explicitPoolAcquire = previous
+    }
+  }
+}
+
+/**
+ * @param {Function} method
+ * @returns {Function}
+ */
+function wrapPoolQuery (method) {
+  return wrapPoolQueryMethod(method, connectionStartCh)
+}
+
+/**
+ * @param {Function} Pool
+ * @returns {boolean}
+ */
+function claimPoolWrap (Pool) {
+  if (wrappedPools.has(Pool)) return false
+  wrappedPools.add(Pool)
+  return true
+}
+
+/**
+ * @param {Function} Pool
+ * @returns {Function}
+ */
+function wrapPublicPool (Pool) {
+  if (!claimPoolWrap(Pool)) return Pool
+
+  shimmer.wrap(Pool.prototype, 'getConnection', wrapExplicitPoolGetConnection)
+  shimmer.wrap(Pool.prototype, 'query', wrapPoolQuery)
+  shimmer.wrap(Pool.prototype, 'execute', wrapPoolQuery)
+  return Pool
+}
+
+/**
+ * @param {Function} Pool
+ * @returns {Function}
+ */
+function wrapInternalPool (Pool) {
+  if (!claimPoolWrap(Pool)) return Pool
+
+  // The idle callback must be replaced before validation and observed before the method returns,
+  // which cannot be expressed by Orchestrion's subscriber lifecycle.
+  shimmer.wrap(Pool.prototype, '_acquireIdleConnection', wrapPoolAcquireIdleMethod)
+  shimmer.wrap(Pool.prototype, 'getConnection', wrapPoolGetConnectionMethod)
+  shimmer.wrap(Pool.prototype, '_createPoolConnection', wrapPoolMethod)
+  return Pool
+}
+
+/**
+ * @param {Record<string, unknown>} acquireCtx
+ * @param {Function} callback
+ * @param {unknown} thisArg
+ * @param {ArgumentsLike} args
+ * @returns {unknown}
+ */
+function finishExplicitPoolAcquire (acquireCtx, callback, thisArg, args) {
+  acquireFinishCh.publish(acquireCtx)
+  return callback.apply(thisArg, args)
+}
+
+/**
+ * @param {PoolAcquireTiming} timing
+ */
+function startPoolAcquireTiming (timing) {
+  const start = performance.now()
+  timing.poolAcquireStart = start
+  timing.poolAcquireStartedAt ??= start
+}
+
+/**
+ * @param {PoolAcquireTiming} timing
+ * @returns {number}
+ */
+function finishPoolAcquireTiming (timing) {
+  const poolWaitTime = (timing.poolWaitTime ?? 0) + acquireWait(timing.poolAcquireStart)
+  timing.poolAcquireStart = undefined
+  timing.poolWaitTime = poolWaitTime
+  return poolWaitTime
+}
+
+/**
+ * @param {Function} acquireIdleConnection
+ * @returns {Function}
+ */
+function wrapPoolAcquireIdleMethod (acquireIdleConnection) {
+  return function wrappedAcquireIdleConnection (callback) {
+    const timing = currentPoolAcquireTiming
+    if (timing === undefined) {
+      return acquireIdleConnection.apply(this, arguments)
+    }
+    currentPoolAcquireTiming = undefined
+
+    let completed = false
+    let synchronous = true
+
+    /**
+     * @param {unknown} error
+     * @returns {unknown}
+     */
+    arguments[0] = function wrappedCallback (error) {
+      completed = true
+      if (error) {
+        if (synchronous) {
+          startPoolAcquireTiming(timing)
+        }
+      } else if (synchronous) {
+        timing.poolWaitTime = 0
+      } else {
+        finishPoolAcquireTiming(timing)
+      }
+      return callback.apply(this, arguments)
+    }
+
+    const result = acquireIdleConnection.apply(this, arguments)
+    synchronous = false
+    if (!completed) {
+      startPoolAcquireTiming(timing)
+    }
+    return result
+  }
+}
+
+/**
+ * @param {Function} getConnection
+ * @returns {Function}
+ */
 function wrapPoolGetConnectionMethod (getConnection) {
-  return function wrappedGetConnection (...args) {
-    const cb = args.at(-1)
-    if (typeof cb !== 'function') return getConnection.apply(this, args)
+  const poolWaitTransfer = { deferred: false }
+
+  return function wrappedGetConnection (cmdParam, callback) {
+    if (!connectionStartCh.hasSubscribers) {
+      return getConnection.apply(this, arguments)
+    }
 
     const ctx = {}
+    const poolQueryAcquire = isPoolQueryAcquire()
+    const acquireCtx = !poolQueryAcquire && explicitPoolAcquire && acquireStartCh.hasSubscribers
+      ? { conf: this.opts.connOptions }
+      : undefined
+    const pool = this
+
+    if (acquireCtx !== undefined) {
+      acquireStartCh.publish(acquireCtx)
+    }
 
-    args[args.length - 1] = function (...args) {
-      return connectionFinishCh.runStores(ctx, cb, this, ...args)
+    /**
+     * @param {unknown} error
+     * @param {{ streamOut: object }|undefined} connection
+     * @returns {unknown}
+     */
+    arguments[1] = function wrappedCallback (error, connection) {
+      const poolWaitTime = poolQueryAcquire || acquireCtx !== undefined
+        ? finishPoolAcquireTiming(ctx)
+        : undefined
+
+      if (poolQueryAcquire) {
+        if (error) {
+          return runPoolAcquireError(
+            ctx.poolAcquireStartedAt,
+            error,
+            { conf: pool.opts.connOptions },
+            poolAcquireChannels,
+            ctx,
+            callback,
+            this,
+            arguments,
+            poolWaitTime
+          )
+        }
+
+        return runWithPoolWait(
+          poolWaitTransfer,
+          connection.streamOut,
+          poolWaitTime,
+          connectionFinishCh,
+          ctx,
+          callback,
+          this,
+          arguments
+        )
+      }
+
+      if (acquireCtx !== undefined) {
+        if (!error && connection !== undefined) {
+          clearPoolWaitTime(connection.streamOut)
+        }
+        acquireCtx.error = error
+        acquireCtx.poolWaitTime = poolWaitTime
+        return connectionFinishCh.runStores(
+          ctx,
+          finishExplicitPoolAcquire,
+          undefined,
+          acquireCtx,
+          callback,
+          this,
+          arguments
+        )
+      }
+      return connectionFinishCh.runStores(ctx, callback, this, ...arguments)
     }
 
     connectionStartCh.publish(ctx)
 
-    return getConnection.apply(this, args)
+    if (!poolQueryAcquire && acquireCtx === undefined) {
+      return getConnection.apply(pool, arguments)
+    }
+
+    const previousTiming = currentPoolAcquireTiming
+    const previousExplicitPoolAcquire = explicitPoolAcquire
+    currentPoolAcquireTiming = ctx
+    if (previousExplicitPoolAcquire) explicitPoolAcquire = false
+    try {
+      return poolQueryAcquire
+        ? runOutsidePoolQueryAcquire(getConnection, pool, arguments)
+        : getConnection.apply(pool, arguments)
+    } finally {
+      currentPoolAcquireTiming = previousTiming
+      if (previousExplicitPoolAcquire) explicitPoolAcquire = true
+    }
   }
 }
 
 const name = 'mariadb'
 
-addHook({ name, file: 'lib/cmd/query.js', versions: ['>=3'] }, (Query) => {
-  return wrapCommand(Query)
-})
+addHook({ name, file: 'lib/cmd/query.js', versions: ['>=3'], patchDefault: true }, wrapCommand)
+addHook({ name, file: 'lib/cmd/execute.js', versions: ['>=3'], patchDefault: true }, wrapCommand)
 
-addHook({ name, file: 'lib/cmd/execute.js', versions: ['>=3'] }, (Execute) => {
-  return wrapCommand(Execute)
-})
+// mariadb 3.4.1 switched its internal getConnection from promises to callbacks. That exposes synchronous
+// recent-idle reuse without a clock read while still letting us time validation and queueing.
+// The same release renamed _createConnection to _createPoolConnection.
+addHook({ name, file: 'lib/pool.js', versions: ['>=3.4.1'], patchDefault: true }, wrapInternalPool)
 
-// mariadb 3.4.1 refactored the pool: getConnection switched from promises to
-// callbacks and _createConnection was renamed to _createPoolConnection.
-addHook({ name, file: 'lib/pool.js', versions: ['>=3.4.1'] }, (Pool) => {
-  shimmer.wrap(Pool.prototype, 'getConnection', wrapPoolGetConnectionMethod)
-  shimmer.wrap(Pool.prototype, '_createPoolConnection', wrapPoolMethod)
-
-  return Pool
-})
+// Orchestrion completion follows the returned Promise, but operation classification must end as
+// soon as the facade calls its private internal pool, so these methods need a synchronous bracket.
+addHook({ name, file: 'lib/pool-promise.js', versions: ['>=3.4.1'], patchDefault: true }, wrapPublicPool)
+addHook({ name, file: 'lib/pool-callback.js', versions: ['>=3.4.1'], patchDefault: true }, wrapPublicPool)
 
-addHook({ name, file: 'lib/pool.js', versions: ['>=3 <3.4.1'] }, (Pool) => {
+addHook({ name, file: 'lib/pool.js', versions: ['>=3 <3.4.1'], patchDefault: true }, (Pool) => {
   shimmer.wrap(Pool.prototype, '_createConnection', wrapPoolMethod)
 
   return Pool