fix mongoose instru not supporting deferred callback

DataDog · Feb 9, 2024 · 012a485 · 012a485
1 parent 12f4bbc
commit 012a485
Showing 1 changed file with 21 additions and 10 deletions.
diff --git a/packages/datadog-instrumentations/src/mongoose.js b/packages/datadog-instrumentations/src/mongoose.js
@@ -79,20 +79,25 @@ addHook({
         })
 
         let callbackWrapped = false
-        const lastArgumentIndex = arguments.length - 1
 
-        if (typeof arguments[lastArgumentIndex] === 'function') {
-          // is a callback, wrap it to execute finish()
-          shimmer.wrap(arguments, lastArgumentIndex, originalCb => {
-            return function () {
-              finish()
+        const wrapCallbackIfExist = (args) => {
+          const lastArgumentIndex = args.length - 1
 
-              return originalCb.apply(this, arguments)
-            }
-          })
+          if (typeof args[lastArgumentIndex] === 'function') {
+            // is a callback, wrap it to execute finish()
+            shimmer.wrap(args, lastArgumentIndex, originalCb => {
+              return function () {
+                finish()
+
+                return originalCb.apply(this, arguments)
+              }
+            })
 
-          callbackWrapped = true
+            callbackWrapped = true
+          }
         }
+
+        wrapCallbackIfExist(arguments)
 
         return asyncResource.runInAsyncScope(() => {
           startCh.publish({
@@ -106,8 +111,14 @@ addHook({
           if (!callbackWrapped) {
             shimmer.wrap(res, 'exec', originalExec => {
               return function wrappedExec () {
+                wrapCallbackIfExist(arguments)
+
                 const execResult = originalExec.apply(this, arguments)
 
+                if (callbackWrapped || typeof execResult?.then !== 'function') {
+                  return execResult
+                }
+
                 // wrap them method, wrap resolve and reject methods
                 shimmer.wrap(execResult, 'then', originalThen => {
                   return function wrappedThen () {