[pinpoint-apm#140] Support AsyncLocalStorage [email protected]

* Rename refactoring MethodDescriptor2 to MethodDescriptor * AsyncLocalStorage above [email protected] The AsyncHook has been deprecated. Node Document recommend AsyncLocalStorage.
feelform · Nov 21, 2023 · 6ddb781 · 6ddb781
1 parent 8383b14
commit 6ddb781
Show file tree

Hide file tree

Showing 32 changed files with 1,257 additions and 1,254 deletions.
diff --git a/lib/context/async-trace.js b/lib/context/async-trace.js
@@ -10,22 +10,15 @@ const SpanEventRecorder = require('./span-event-recorder')
 const SpanEvent = require('./span-event')
 const SpanChunk = require('./span-chunk')
 const BufferedStorage = require('./buffered-storage')
+const Trace = require('./trace')
 
-class AsyncTrace {
+class AsyncTrace extends Trace {
   constructor(span, asyncId, traceId, agentInfo, dataSender, sampling) {
+    super(traceId, agentInfo, dataSender, sampling)
     this.span = span
-    this.traceId = traceId
-    this.agentInfo = agentInfo
     this.asyncId = asyncId
-
-    this.spanEventRecorder = null
-
     const createAsyncSpanChunk = SpanChunk.getAsyncFactoryMethod(traceId, agentInfo, asyncId)
     this.storage = new BufferedStorage(dataSender, createAsyncSpanChunk)
-
-    this.callStack = []
-    this.sequence = 0
-    this.sampling = sampling
   }
 
   traceAsyncBegin() {

diff --git a/lib/context/context-manager.js b/lib/context/context-manager.js
@@ -41,13 +41,13 @@ const setObject = (value) => {
   traceObjectMap.set(asyncId, value)
 }
 
-const getAllObject = () => {
-  return traceObjectMap
+const disable = () => {
+  traceObjectMap.clear()
 }
 
 module.exports = {
   getObject,
   setObject,
-  getAllObject,
   start,
+  disable
 }
diff --git a/lib/context/disable-async-trace.js b/lib/context/disable-async-trace.js
@@ -9,20 +9,19 @@
 const DisableSpanEventRecorder = require('./disable-span-event-recorder')
 
 class DisableAsyncTrace {
-    traceAsyncBegin () {    
-        return new DisableSpanEventRecorder()
-      }
+  traceAsyncBegin () {    
+    return new DisableSpanEventRecorder()
+  }
 
-      traceAsyncEnd (spanEventRecorder) {
-
-      }
+  traceAsyncEnd () {
+  }
 
-      canSampled () {
-        return false
-      }
+  canSampled () {
+    return false
+  }
 
-      close () {
-      }
+  close () {
+  }
 }
 
 module.exports = DisableAsyncTrace
diff --git a/lib/context/disable-trace.js b/lib/context/disable-trace.js
@@ -6,7 +6,6 @@
 
 'use strict'
 
-const SpanRecorder = require('./span-recorder')
 const Span = require('./span')
 const DisableSpanEventRecorder = require('./disable-span-event-recorder')
 const DisableAsyncTrace = require('./disable-async-trace')
@@ -24,22 +23,9 @@ class DisableTrace {
         return new DisableSpanEventRecorder()
     }
 
-    traceBlockEnd(spanEventRecorder) {}
+    traceBlockEnd() {}
 
-    completeSpanEvent(spanEvent) {
-    }
-
-    newAsyncTrace(spanEventRecorder) {
-        if (spanEventRecorder) {
-            const asyncId = spanEventRecorder.recordNextAsyncId()
-            return new DisableAsyncTrace()
-        }
-    }
-
-    newAsyncTraceWithId(asyncId) {
-        if (asyncId) {
-            return new DisableAsyncTrace()
-        }
+    completeSpanEvent() {
     }
 
     canSampled() {

diff --git a/lib/context/method-descriptor-builder.js b/lib/context/method-descriptor-builder.js
@@ -6,7 +6,7 @@
 
 'use strict'
 
-const MethodDescriptor2 = require('../instrumentation/method-descriptor2')
+const MethodDescriptor = require('../instrumentation/method-descriptor')
 const MethodType = require('../constant/method-type')
 
 // https://v8.dev/docs/stack-trace-api#appendix%3A-stack-trace-format
@@ -191,7 +191,7 @@ class MethodDescriptorBuilder {
     }
 
     build() {
-        const value = new MethodDescriptor2(this.functionName)
+        const value = new MethodDescriptor(this.functionName)
         if (this.type !== 'Function' && typeof this.methodName === 'undefined') {
             value.methodName = this.functionName
         } else {

diff --git a/lib/context/method-descriptor.js b/lib/context/method-descriptor.js
diff --git a/lib/context/trace-context.js b/lib/context/trace-context.js
@@ -6,7 +6,6 @@
 
 'use strict'
 
-const contextManager = require('./context-manager')
 const Trace = require('./trace')
 const TransactionId = require('./transaction-id')
 const TraceId = require('./trace-id')
@@ -15,6 +14,7 @@ const activeTrace = require('../metric/active-trace')
 const log = require('../utils/logger')
 const sampler = require('../sampler/sampler')
 const DisableTrace = require('./disable-trace')
+const localStorage = require('../instrumentation/context/local-storage')
 
 class TraceContext {
   constructor() {
@@ -35,8 +35,6 @@ class TraceContext {
       throw new Error('Fail to initialize pinpoint context')
     }
 
-    contextManager.start()
-
     const instance = new TraceContext()
     instance.agentInfo = {
       agentId: options.agentId,
@@ -79,13 +77,11 @@ class TraceContext {
 
     if (false == sampling) {
       const disableTrace = new DisableTrace(traceId, this.agentInfo, requestData)
-      this.setCurrentTraceObject(disableTrace)
       return disableTrace
     }
 
     try {
       const trace = new Trace(traceId, this.agentInfo, this.dataSender, sampling, requestData)
-      this.setCurrentTraceObject(trace)
       activeTrace.register(trace)
       return trace
     } catch (e) {
@@ -107,11 +103,7 @@ class TraceContext {
   }
 
   currentTraceObject() {
-    return contextManager.getObject()
-  }
-
-  setCurrentTraceObject(traceObject) {
-    contextManager.setObject(traceObject)
+    return localStorage.getStore()
   }
 
   makeTrace(requestData) {
@@ -125,19 +117,12 @@ class TraceContext {
           requestData.spanId,
           requestData.parentSpanId,
           requestData.flag)
-        const disableTrace = new DisableTrace(traceId, this.agentInfo, requestData)
-        this.setCurrentTraceObject(disableTrace)
-        return disableTrace
+        return new DisableTrace(traceId, this.agentInfo, requestData)
       }
     } else {
       return this.newTraceObject(this.isSampling())
     }
   }
-
-  // only test
-  getAllTraceObject() {
-    return contextManager.getAllObject()
-  }
 }
 
 module.exports = TraceContext
diff --git a/lib/context/trace.js b/lib/context/trace.js
@@ -10,13 +10,9 @@ const SpanRecorder = require('./span-recorder')
 const SpanEventRecorder = require('./span-event-recorder')
 const Span = require('./span')
 const SpanEvent = require('./span-event')
-const AsyncTrace = require('./async-trace')
 const BufferedStorage = require('./buffered-storage')
 const SpanChunk = require('./span-chunk')
 const log = require('../utils/logger')
-const defaultPredefinedMethodDescriptorRegistry = require('../constant/default-predefined-method-descriptor-registry')
-const DisableAsyncTrace = require('./disable-async-trace')
-const ServiceType = require('../context/service-type')
 const calledTraceBlockEnd = Symbol('called-traceBlockEnd')
 
 class Trace {
@@ -89,29 +85,6 @@ class Trace {
     }
   }
 
-  newAsyncTrace(spanEventRecorder) {
-    if (!spanEventRecorder || typeof spanEventRecorder.recordNextAsyncId !== 'function') {
-      return new DisableAsyncTrace()
-    }
-    const asyncId = spanEventRecorder.recordNextAsyncId()
-    return this.newAsyncTraceWithId(asyncId)
-  }
-
-  newAsyncTraceWithId(asyncId) {
-    if (!asyncId) {
-      return new DisableAsyncTrace()
-    }
-
-    const asyncTrace = new AsyncTrace(this.span, asyncId, this.traceId, this.agentInfo, this.dataSender, this.sampling)
-    const spanEventRecorder = asyncTrace.traceAsyncBegin()
-    spanEventRecorder.recordServiceType(ServiceType.async)
-    spanEventRecorder.recordApiId(defaultPredefinedMethodDescriptorRegistry.asyncInvocationMethodDescriptor.getApiId())
-    spanEventRecorder.spanEvent.endPoint = null
-    spanEventRecorder.spanEvent.destinationId = null
-    asyncTrace.storage.storeSpanEvent(spanEventRecorder.spanEvent)
-    return asyncTrace
-  }
-
   canSampled() {
     return this.sampling
   }

diff --git a/lib/instrumentation/call-stack.js b/lib/instrumentation/call-stack.js
@@ -85,7 +85,7 @@ const captureLocationFileNameNamedGroup = (stack) => {
 }
 
 const locationFileNameLineNumber = (stack) => {
-    return stack.match(/ \((?<location>.+\/)(?<fileName>[^:/]+):(?<lineNumber>[0-9]+):(?<columnNumber>[0-9]+)\)$/)
+    return stack.match(/ \(?(?<location>[^\s]+\/)(?<fileName>[^:/]+):(?<lineNumber>[0-9]+):(?<columnNumber>[0-9]+)\)?$/)
 }
 
 const fileNameOfStack = (stack) => {

diff --git a/lib/instrumentation/context/async-context-local-storage.js b/lib/instrumentation/context/async-context-local-storage.js
@@ -0,0 +1,29 @@
+/**
+ * Pinpoint Node.js Agent
+ * Copyright 2023-present NAVER Corp.
+ * Apache License v2.0
+ */
+
+'use strict'
+
+const { AsyncLocalStorage } = require('node:async_hooks')
+
+class AsyncContextLocalStorage {
+    constructor() {
+        this.storage = new AsyncLocalStorage()
+    }
+
+    run(store, callback) {
+        this.storage.run(store, callback)
+    }
+
+    getStore() {
+        return this.storage.getStore()
+    }
+
+    disable() {
+        this.storage.disable()
+    }
+}
+
+module.exports = AsyncContextLocalStorage
diff --git a/lib/instrumentation/context/async-hooks-local-storage.js b/lib/instrumentation/context/async-hooks-local-storage.js
@@ -0,0 +1,30 @@
+/**
+ * Pinpoint Node.js Agent
+ * Copyright 2023-present NAVER Corp.
+ * Apache License v2.0
+ */
+
+'use strict'
+
+const contextManager = require('../../context/context-manager')
+
+class AsyncHooksLocalStorage {
+    constructor() {
+        contextManager.start()
+    }
+
+    run(store, callback) {
+        contextManager.setObject(store)
+        callback()
+    }
+
+    getStore() {
+        return contextManager.getObject()
+    }
+
+    disable() {
+        contextManager.disable()
+    }
+}
+
+module.exports = AsyncHooksLocalStorage