[#249] Fix InstrumentArrowFunction local frame stack variable capturi…

…ng bug [#249] Fix InstrumentArrowFunction local frame stack variable capturing bug [#249] Fix InstrumentArrowFunction local frame stack variable capturing bug [#249] Fix InstrumentArrowFunction local frame stack variable capturing bug [#249] Fix InstrumentArrowFunction local frame stack variable capturing bug [#249] Fix InstrumentArrowFunction local frame stack variable capturing bug
pinpoint-apm · Dec 20, 2024 · 9adf4d7 · 9adf4d7
1 parent 6f0e263
commit 9adf4d7
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 121 deletions.
diff --git a/lib/instrumentation/callback-interceptor-runner.js b/lib/instrumentation/callback-interceptor-runner.js
diff --git a/lib/instrumentation/instrument-arrow-function.js b/lib/instrumentation/instrument-arrow-function.js
@@ -38,7 +38,7 @@ class InstrumentArrowFunction {
         })
     }
 
-    addCallbackInterceptor(interceptor, callerSpanEventRecorder) {
+    addChildTraceInterceptor(callerSpanEventRecorder) {
         if (!this.target) {
             return
         }
@@ -50,14 +50,28 @@ class InstrumentArrowFunction {
         if (!callerSpanEventRecorder || typeof callerSpanEventRecorder.getNextAsyncId !== 'function') {
             return
         }
-        const asyncId = callerSpanEventRecorder.getNextAsyncId()
 
         const target = this.target
         const method = this.method
-        const trace = localStorage.getStore()
         const traceContext = this.traceContext
+        const trace = localStorage.getStore()
+        const asyncId = callerSpanEventRecorder.getNextAsyncId()
         shimmer.wrap(target, method, function (original) {
-            return interceptor.makeHookCallbackFunction(trace, asyncId, new HookFunctionArguments(target, method, original), traceContext)
+            return function () {
+                if (!trace) {
+                    return original.apply(this, arguments)
+                }
+                if (!asyncId) {
+                    return original.apply(this, arguments)
+                }
+
+                const childTraceBuilder = traceContext.continueAsyncContextTraceObject(trace.getTraceRoot(), asyncId.nextLocalAsyncId2())
+                return localStorage.run(childTraceBuilder, () => {
+                    const result = original.apply(this, arguments)
+                    childTraceBuilder.close()
+                    return result
+                })
+            }
         })
     }
 }

diff --git a/lib/instrumentation/instrument-method.js b/lib/instrumentation/instrument-method.js
@@ -73,20 +73,21 @@ class InstrumentMethod {
                     if (recorder && arguments.length > 0 && typeof interceptor.callbackIndexOf === 'function' && typeof interceptor.callbackIndexOf(arguments) === 'number') {
                         const asyncId = recorder.getNextAsyncId()
                         const callbackIndex = interceptor.callbackIndexOf(arguments)
-                        const callback = arguments[callbackIndex]
-                        arguments[callbackIndex] = function () {
-                            try {
-                                interceptor.prepareBeforeAsyncTrace?.(this, arguments)
-                            } catch (error) {
-                                log.error(`arguments: ${error}`)
+                        shimmer.wrap(arguments, callbackIndex, function (original) {
+                            return function () {
+                                try {
+                                    interceptor.prepareBeforeAsyncTrace?.(this, arguments)
+                                } catch (error) {
+                                    log.error(`arguments: ${error}`)
+                                }
+                                const childTraceBuilder = traceContext.continueAsyncContextTraceObject(trace.getTraceRoot(), asyncId.nextLocalAsyncId2())
+                                return localStorage.run(childTraceBuilder, () => {
+                                    const returnedValue = original.apply(this, arguments)
+                                    childTraceBuilder?.close()
+                                    return returnedValue
+                                })
                             }
-                            const childTraceBuilder = traceContext.continueAsyncContextTraceObject(trace.getTraceRoot(), asyncId.nextLocalAsyncId2())
-                            return localStorage.run(childTraceBuilder, () => {
-                                const returnedValue = callback.apply(this, arguments)
-                                childTraceBuilder?.close()
-                                return returnedValue
-                            })
-                        }
+                        })
                     }
                 } catch (error) {
                     // TODO: INTERNAL ERROR logger
@@ -107,44 +108,6 @@ class InstrumentMethod {
             }
         })
     }
-
-    addCallbackInterceptor(interceptor, callerSpanEventRecorder) {
-        if (!this.target) {
-            return
-        }
-
-        if (!this.method || typeof this.method !== 'string' || !this.target[this.method] || typeof this.target[this.method] !== 'function') {
-            return
-        }
-
-        if (!callerSpanEventRecorder || typeof callerSpanEventRecorder.getNextAsyncId !== 'function') {
-            return
-        }
-
-        const traceContext = this.traceContext
-        shimmer.wrap(this.target, this.method, function (original) {
-            return function () {
-                const trace = localStorage.getStore()
-                if (!trace) {
-                    return original.apply(this, arguments)
-                }
-                const asyncId = callerSpanEventRecorder.getNextAsyncId()
-
-                try {
-                    interceptor.prepareBeforeAsyncTrace?.(this, arguments)
-                } catch (error) {
-                    log.error(`addCallbackInterceptor: ${error}`)
-                }
-
-                const childTraceBuilder = traceContext.continueAsyncContextTraceObject(trace.getTraceRoot(), asyncId.nextLocalAsyncId2())
-                return localStorage.run(childTraceBuilder, () => {
-                    const returnedValue = original.apply(this, arguments)
-                    childTraceBuilder?.close()
-                    return returnedValue
-                })
-            }
-        })
-    }
 }
 
 module.exports = InstrumentMethod
diff --git a/lib/instrumentation/module/callback-interceptor.js b/lib/instrumentation/module/callback-interceptor.js
diff --git a/lib/instrumentation/module/ioredis/ioredis-send-command-interceptor.js b/lib/instrumentation/module/ioredis/ioredis-send-command-interceptor.js
@@ -10,7 +10,6 @@ const MethodDescriptorBuilder = require('../../../context/method-descriptor-buil
 const serviceType = require('../redis/redis-service-type')
 const { addressStringOf } = require('../../../utils/convert-utils')
 const InstrumentArrowFunction = require('../../instrument-arrow-function')
-const CallbackInterceptor = require('../callback-interceptor')
 const endTraceSymbol = Symbol('ioredis-end-trace')
 
 class IoredisSendCommandInterceptor {
@@ -39,7 +38,7 @@ class IoredisSendCommandInterceptor {
 
         const command = args[0]
         command[endTraceSymbol] = function() {}
-        InstrumentArrowFunction.make(command, endTraceSymbol, this.traceContext).addCallbackInterceptor(new CallbackInterceptor(), recorder)
+        InstrumentArrowFunction.make(command, endTraceSymbol, this.traceContext).addChildTraceInterceptor(recorder)
         command.promise.finally(command[endTraceSymbol])
     }
 }

diff --git a/lib/instrumentation/module/redis/redis-internal-send-command-interceptor.js b/lib/instrumentation/module/redis/redis-internal-send-command-interceptor.js
@@ -10,7 +10,6 @@ const serviceType = require('./redis-service-type')
 const MethodDescriptorBuilder = require('../../../context/method-descriptor-builder')
 const { addressStringOf } = require('../../../utils/convert-utils')
 const InstrumentArrowFunction = require('../../instrument-arrow-function')
-const CallbackInterceptor = require('../callback-interceptor')
 
 class RedisInternalSendCommandInterceptor {
     constructor(traceContext) {
@@ -40,7 +39,7 @@ class RedisInternalSendCommandInterceptor {
         }
 
         const command = args[0]
-        InstrumentArrowFunction.make(command, 'callback', this.traceContext).addCallbackInterceptor(new CallbackInterceptor(), recorder)
+        InstrumentArrowFunction.make(command, 'callback', this.traceContext).addChildTraceInterceptor(recorder)
     }
 }