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

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() {

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

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() {

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
   }

lib/instrumentation/context/async-context-local-storage.js

@@ -0,0 +1,25 @@
+/**
+ * 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()
+    }
+}
+
+module.exports = AsyncContextLocalStorage

lib/instrumentation/context/async-hooks-local-storage.js

@@ -0,0 +1,26 @@
+/**
+ * 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()
+    }
+}
+
+module.exports = AsyncHooksLocalStorage

lib/instrumentation/context/async-service-type.js

@@ -0,0 +1,11 @@
+/**
+ * Pinpoint Node.js Agent
+ * Copyright 2023-present NAVER Corp.
+ * Apache License v2.0
+ */
+
+'use strict'
+
+const ServiceType = require('../../context/service-type')
+
+module.exports = new ServiceType(100, 'ASYNC')

lib/instrumentation/context/local-storage.js

@@ -0,0 +1,34 @@
+/**
+ * Pinpoint Node.js Agent
+ * Copyright 2023-present NAVER Corp.
+ * Apache License v2.0
+ */
+
+'use strict'
+
+const semver = require('semver')
+const AsyncContextLocalStorage = require('./async-context-local-storage')
+const AsyncHooksLocalStorage = require('./async-hooks-local-storage')
+
+class LocalStorage {
+    constructor() {
+        this.storage = this.createLocalStorage()
+    }
+
+    createLocalStorage() {
+        if (!semver.satisfies(process.version, '>=16.4.0')) {
+            return new AsyncHooksLocalStorage()
+        }
+        return new AsyncContextLocalStorage()
+    }
+
+    run(store, callback) {
+        return this.storage.run(store, callback)
+    }
+
+    getStore() {
+        return this.storage.getStore()
+    }
+}
+
+module.exports = new LocalStorage()

lib/instrumentation/context/trace-builder.js

@@ -0,0 +1,61 @@
+/**
+ * Pinpoint Node.js Agent
+ * Copyright 2023-present NAVER Corp.
+ * Apache License v2.0
+ */
+
+'use strict'
+
+const Span = require('../../context/span')
+const SpanRecorder = require('../../context/span-recorder')
+const SpanChunk = require('../../context/span-chunk')
+const BufferedStorage = require('../../context/buffered-storage')
+const AsyncTrace = require('../../context/async-trace')
+const DisableAsyncTrace = require('../../context/disable-async-trace')
+const serviceType = require('./async-service-type')
+const defaultPredefinedMethodDescriptorRegistry = require('../../constant/default-predefined-method-descriptor-registry')
+
+class TraceBuilder {
+    constructor(traceId, agentInfo, dataSender, sampling) {
+        this.traceId = traceId
+        this.agentInfo = agentInfo
+        this.dataSender = dataSender
+        this.sampling = sampling
+
+        this.span = new Span(traceId, agentInfo)
+        this.spanRecorder = new SpanRecorder(this.span)
+
+        this.callStack = []
+        this.sequence = 0
+
+        const createSpanChunk = SpanChunk.getFactoryMethod(traceId, agentInfo)
+        this.storage = new BufferedStorage(dataSender, createSpanChunk)
+    }
+
+    static valueOfTrace(trace) {
+        const value = new TraceBuilder(trace.traceId, trace.agentInfo, trace.dataSender, trace.sampling)
+        value.span = trace.span
+        value.spanRecorder = trace.spanRecorder
+        value.callStack = trace.callStack
+        value.sequence = trace.sequence
+        value.storage = trace.storage
+        return value
+    }
+
+    buildAsyncTrace(asyncId) {
+        if (!asyncId) {
+            return new DisableAsyncTrace()
+        }
+
+        const value = new AsyncTrace(this.span, asyncId, this.traceId, this.agentInfo, this.dataSender, this.sampling)
+        const spanEventRecorder = value.traceAsyncBegin()
+        spanEventRecorder.recordServiceType(serviceType)
+        spanEventRecorder.recordApiId(defaultPredefinedMethodDescriptorRegistry.asyncInvocationDescriptor.getApiId())
+        spanEventRecorder.spanEvent.endPoint = null
+        spanEventRecorder.spanEvent.destinationId = null
+        value.storage.storeSpanEvent(spanEventRecorder.spanEvent)
+        return value
+    }
+}
+
+module.exports = TraceBuilder