[pinpoint-apm#140] Support nested async trace

* Remove no needs an argument TraceContext

lib/agent.js

@@ -26,13 +26,10 @@ class Agent {
 
     log.warn('[Pinpoint Agent] Configuration', this.config)
 
-    if (!this.config.enable || this.config.enable.toString() !== 'true') {
-      global.__PINPOINT_ENABLED__ = false
+    if (!this.config || !this.config.enable || this.config.enable.toString() !== 'true') {
       log.warn('[Pinpoint Agent][' + this.config.agentId + '] Disabled')
       return
     }
-
-    global.__PINPOINT_ENABLED__ = true
     log.warn('[Pinpoint Agent][' + this.config.agentId + '] Init Started')
 
     const agentId = this.config.agentId

lib/context/context-manager.js

@@ -7,7 +7,6 @@
 'use strict'
 
 const async_hooks = require('async_hooks')
-
 const traceObjectMap = new Map()
 
 const init = (asyncId, type, triggerAsyncId, resource) => {
@@ -24,9 +23,7 @@ const destroy = (asyncId) => {
 }
 
 const start = () => {
-  if (global.__PINPOINT_ENABLED__) {
-    async_hooks.createHook({init, destroy}).enable()
-  }
+  async_hooks.createHook({init, destroy}).enable()
 }
 
 const getObject = () => {

lib/instrumentation/context/local-storage.js

@@ -9,13 +9,18 @@
 const semver = require('semver')
 const AsyncContextLocalStorage = require('./async-context-local-storage')
 const AsyncHooksLocalStorage = require('./async-hooks-local-storage')
-
+const { getConfig } = require('../../config')
 class LocalStorage {
     constructor() {
         this.storage = this.createLocalStorage()
     }
 
     createLocalStorage() {
+        const config = getConfig()
+        if (config && !config.enable) {
+            return new EmptyLocalStorage()
+        }
+
         if (!semver.satisfies(process.version, '>=16.4.0')) {
             return new AsyncHooksLocalStorage()
         }
@@ -35,4 +40,15 @@ class LocalStorage {
     }
 }
 
+class EmptyLocalStorage {
+    run() {
+    }
+
+    getStore() {
+    }
+
+    disable() {
+    }
+}
+
 module.exports = new LocalStorage()

lib/instrumentation/instrument-method-context.js

@@ -7,22 +7,13 @@
 'use strict'
 
 class InstrumentMethodContext {
-    constructor(methodDescriptorBuilder, traceContext) {
+    constructor(methodDescriptorBuilder) {
         this.methodDescriptorBuilder = methodDescriptorBuilder
-        this.traceContext = traceContext
     }
 
     getMethodDescriptorBuilder() {
         return this.methodDescriptorBuilder
     }
-
-    getTraceContext() {
-        return this.traceContext
-    }
-
-    currentTraceObject() {
-        return this.traceContext.currentTraceObject()
-    }
 }
 
 module.exports = InstrumentMethodContext

lib/instrumentation/instrument-method.js

@@ -13,17 +13,17 @@ const { callSite } = require('./call-stack')
 const InstrumentMethodContext = require('./instrument-method-context')
 const config = require('../config')
 const TraceBuilder = require('./context/trace-builder')
+const localStorage = require('./context/local-storage')
 
 // ref: SpanEventSimpleAroundInterceptorForPlugin.java
 class InstrumentMethod {
-    constructor(target, method, traceContext) {
+    constructor(target, method) {
         this.target = target
         this.method = method
-        this.traceContext = traceContext
     }
 
-    static make(target, method, traceContext) {
-        return new InstrumentMethod(target, method, traceContext)
+    static make(target, method) {
+        return new InstrumentMethod(target, method)
     }
 
     addScopedInterceptor(interceptor) {
@@ -35,7 +35,6 @@ class InstrumentMethod {
             return
         }
 
-        const traceContext = this.traceContext
         const target = this.target
         shimmer.wrap(this.target, this.method, function (original) {
             return function () {
@@ -46,13 +45,13 @@ class InstrumentMethod {
                 } else {
                     builder = interceptor.methodDescriptorBuilder
                 }
-                const context = new InstrumentMethodContext(builder, traceContext)
-
+
+                const context = new InstrumentMethodContext(builder)
                 if (typeof interceptor.prepareBeforeTrace === 'function') {
                     interceptor.prepareBeforeTrace(target, arguments, this, context)
                 }
 
-                const trace = traceContext.currentTraceObject()
+                const trace = localStorage.getStore()
                 let recorder
                 if (trace) {
                     recorder = trace.traceBlockBegin()

lib/instrumentation/module/express.js

@@ -23,8 +23,8 @@ module.exports = function (agent, version, express) {
     log.error('Invalid Version error')
     return express
   }
-  InstrumentMethod.make(express.Router, 'use', agent.traceContext).addScopedInterceptor(new ExpressUseInterceptor(agent.traceContext))
-  InstrumentMethod.make(express.Router, 'route', agent.traceContext).addScopedInterceptor(new ExpressRouteInterceptor(agent.traceContext))
+  InstrumentMethod.make(express.Router, 'use').addScopedInterceptor(new ExpressUseInterceptor())
+  InstrumentMethod.make(express.Router, 'route').addScopedInterceptor(new ExpressRouteInterceptor())
 
   return express
 }

lib/instrumentation/module/express/express-route-interceptor.js

@@ -10,8 +10,7 @@ const MethodDescriptorBuilder = require('../../../context/method-descriptor-buil
 const patchLayer = require('./patch-layer')
 
 class ExpressRouteInterceptor {
-    constructor(traceContext) {
-        this.traceContext = traceContext
+    constructor() {
         this.methodDescriptorBuilder = MethodDescriptorBuilder.makeRuntimeDetectionMethodDescriptorBuilder()
             .setClassName('Router')
             .setFileNameIndex(3)

lib/instrumentation/module/express/patch-layer.js

@@ -11,13 +11,14 @@ const shimmer = require('@pinpoint-apm/shimmer')
 const serviceType = require('./express-service-type')
 const util = require('util')
 const errorReportedSymbol = Symbol('PinpointErrorReported')
+const localStorage = require('../../../instrumentation/context/local-storage')
 
 function patchLayer(layer, context) {
     shimmer.wrap(layer, 'handle', function (original) {
         let handle
         if (original.length !== 4) {
             handle = function(req, res, next) {
-                const trace = context.currentTraceObject()
+                const trace = localStorage.getStore()
                 if (!trace) {
                     return original.apply(this, arguments)
                 }
@@ -45,7 +46,7 @@ function patchLayer(layer, context) {
             }
         } else {
             handle = function (err, req, res, next) {
-                const trace = context.currentTraceObject()
+                const trace = localStorage.getStore()
                 let spanEventRecorder = null
                 if (shouldReport(err, trace)) {
                     spanEventRecorder = trace.traceBlockBegin()

lib/instrumentation/module/koa-router.js

@@ -17,7 +17,7 @@ module.exports = function (agent, version, router) {
     return router
   }
 
-  InstrumentMethod.make(router.prototype, 'register', agent.traceContext).addScopedInterceptor(new KoaRegisterInterceptor())
+  InstrumentMethod.make(router.prototype, 'register').addScopedInterceptor(new KoaRegisterInterceptor())
 
   return router
 }

lib/instrumentation/module/koa/koa-register-interceptor.js

@@ -9,6 +9,7 @@
 const apiMetaService = require('../../../context/api-meta-service')
 const MethodDescriptorBuilder = require('../../../context/method-descriptor-builder')
 const serviceType = require('./koa-service-type')
+const localStorage = require('../../../instrumentation/context/local-storage')
 
 class KoaRegisterInterceptor {
     constructor() {
@@ -45,7 +46,7 @@ class KoaRegisterInterceptor {
                 builder = builder.makeCloneOf(method.toLowerCase())
             }
             const methodDescriptor = apiMetaService.cacheApiWithBuilder(builder)
-            const trace = context.currentTraceObject()
+            const trace = localStorage.getStore()
             let spanEventRecorder = null
             let result
             try {

lib/instrumentation/module/mysql.js

@@ -21,13 +21,13 @@ module.exports = function (agent, version, mysql) {
   }
 
   log.debug('shimming mysql.createPool')
-  InstrumentMethod.make(mysql, 'createPool', agent.traceContext).addScopedInterceptor(new MySQLCreatePoolInterceptor(agent.traceContext))
+  InstrumentMethod.make(mysql, 'createPool').addScopedInterceptor(new MySQLCreatePoolInterceptor())
 
   log.debug('shimming mysql.createPoolCluster')
-  InstrumentMethod.make(mysql, 'createPoolCluster', agent.traceContext).addScopedInterceptor(new MySQLCreatePoolClusterInterceptor(agent.traceContext))
+  InstrumentMethod.make(mysql, 'createPoolCluster').addScopedInterceptor(new MySQLCreatePoolClusterInterceptor())
 
   log.debug('shimming mysql.createConnection')
-  InstrumentMethod.make(mysql, 'createConnection', agent.traceContext).addScopedInterceptor(new MySQLCreateConnectionInterceptor(agent.traceContext))
+  InstrumentMethod.make(mysql, 'createConnection').addScopedInterceptor(new MySQLCreateConnectionInterceptor())
 
   return mysql
 }

lib/instrumentation/module/mysql/mysql-cluster-of-interceptor.js

@@ -11,8 +11,7 @@ const InstrumentMethod = require('../../instrument-method')
 const MySQLGetConnectionInterceptor = require('./mysql-get-connection-interceptor')
 
 class MySQLClusterOfInterceptor {
-    constructor(traceContext) {
-        this.traceContext = traceContext
+    constructor() {
         this.methodDescriptorBuilder = new MethodDescriptorBuilder('of')
     }
 
@@ -24,7 +23,7 @@ class MySQLClusterOfInterceptor {
         if (!poolNamespace) {
             return
         }
-        InstrumentMethod.make(poolNamespace, 'getConnection', this.traceContext).addScopedInterceptor(new MySQLGetConnectionInterceptor(this.traceContext))
+        InstrumentMethod.make(poolNamespace, 'getConnection').addScopedInterceptor(new MySQLGetConnectionInterceptor())
     }
 }
 

lib/instrumentation/module/mysql/mysql-create-connection-interceptor.js

@@ -15,10 +15,9 @@ const databaseInfoSymbol = require('./mysql-database-information-symbol')
 const MySQLDatabaseInformationContext = require('./mysql-database-information-context')
 
 class MySQLCreateConnectionInterceptor {
-    constructor(traceContext) {
-        this.traceContext = traceContext
+    constructor() {
         this.methodDescriptorBuilder = new MethodDescriptorBuilder('createConnection')
-        this.methodDescriptorBuilder.setParameterVariableNames('connectionUri')
+                                        .setParameterVariableNames('connectionUri')
     }
 
     doInBeforeTrace(recorder) {
@@ -28,8 +27,8 @@ class MySQLCreateConnectionInterceptor {
     doInAfterTrace(recorder, target, args, returned) {
         if (returned) {
             returned[databaseInfoSymbol] = MySQLDatabaseInformationContext.recordAndReturnDatabaseInfo(recorder, target, args)
-            InstrumentMethod.make(returned, 'connect', this.traceContext).addScopedInterceptor(new MySQLConnectionInterceptor())
-            InstrumentMethod.make(returned, 'query', this.traceContext).addScopedInterceptor(new MySQLStatementExecuteQueryInterceptor('Connection'))
+            InstrumentMethod.make(returned, 'connect').addScopedInterceptor(new MySQLConnectionInterceptor())
+            InstrumentMethod.make(returned, 'query').addScopedInterceptor(new MySQLStatementExecuteQueryInterceptor('Connection'))
         }
     }
 }

lib/instrumentation/module/mysql/mysql-create-pool-cluster-interceptor.js

@@ -11,8 +11,7 @@ const InstrumentMethod = require('../../instrument-method')
 const MySQLClusterOfInterceptor = require('./mysql-cluster-of-interceptor')
 
 class MySQLCreatePoolClusterInterceptor {
-    constructor(traceContext) {
-        this.traceContext = traceContext
+    constructor() {
         this.methodDescriptorBuilder = new MethodDescriptorBuilder('createPoolCluster')
     }
 
@@ -24,7 +23,7 @@ class MySQLCreatePoolClusterInterceptor {
         if (!poolCluster) {
             return
         }
-        InstrumentMethod.make(poolCluster, 'of', this.traceContext).addScopedInterceptor(new MySQLClusterOfInterceptor(this.traceContext))
+        InstrumentMethod.make(poolCluster, 'of').addScopedInterceptor(new MySQLClusterOfInterceptor())
     }
 }
 

lib/instrumentation/module/mysql/mysql-create-pool-interceptor.js

@@ -14,8 +14,7 @@ const MySQLGetConnectionInterceptor = require('./mysql-get-connection-intercepto
 const MySQLDatabaseInformationContext = require('./mysql-database-information-context')
 
 class MySQLCreatePoolInterceptor {
-    constructor(traceContext) {
-        this.traceContext = traceContext
+    constructor() {
         this.methodDescriptorBuilder = new MethodDescriptorBuilder('createPool')
     }
 
@@ -25,7 +24,7 @@ class MySQLCreatePoolInterceptor {
 
     prepareAfterTrace(target, args, pool) {
         if (pool) {
-            InstrumentMethod.make(pool, 'getConnection', this.traceContext).addScopedInterceptor(new MySQLGetConnectionInterceptor(this.traceContext))
+            InstrumentMethod.make(pool, 'getConnection').addScopedInterceptor(new MySQLGetConnectionInterceptor())
         }
     }
 

lib/instrumentation/module/mysql/mysql-get-connection-interceptor.js

@@ -13,8 +13,7 @@ const InstrumentMethod = require('../../instrument-method')
 const MySQLStatementExecuteQueryInterceptor = require('./mysql-statement-execute-query-interceptor')
 
 class MySQLGetConnectionInterceptor {
-    constructor(traceContext) {
-        this.traceContext = traceContext
+    constructor() {
         this.methodDescriptorBuilder = new MethodDescriptorBuilder('getConnection')
             .setClassName('Pool')
         this.serviceType = serviceType
@@ -31,7 +30,7 @@ class MySQLGetConnectionInterceptor {
     prepareBeforeAsyncTrace(target, args) {
         if (args.length > 1 && args[1]) {
             const connection = args[1]
-            InstrumentMethod.make(connection, 'query', this.traceContext).addScopedInterceptor(new MySQLStatementExecuteQueryInterceptor('PoolConnection'))
+            InstrumentMethod.make(connection, 'query').addScopedInterceptor(new MySQLStatementExecuteQueryInterceptor('PoolConnection'))
         }
     }
 

lib/instrumentation/module/mysql2.js

@@ -20,13 +20,13 @@ module.exports = function (agent, version, mysql) {
   }
 
   log.debug('shimming mysql.createPool')
-  InstrumentMethod.make(mysql, 'createPool', agent.traceContext).addScopedInterceptor(new MySQLCreatePoolInterceptor(agent.traceContext))
+  InstrumentMethod.make(mysql, 'createPool').addScopedInterceptor(new MySQLCreatePoolInterceptor())
 
   log.debug('shimming mysql.createPoolCluster')
-  InstrumentMethod.make(mysql, 'createPoolCluster', agent.traceContext).addScopedInterceptor(new MySQLCreatePoolClusterInterceptor(agent.traceContext))
+  InstrumentMethod.make(mysql, 'createPoolCluster').addScopedInterceptor(new MySQLCreatePoolClusterInterceptor())
 
   log.debug('shimming mysql.createConnection')
-  InstrumentMethod.make(mysql, 'createConnection', agent.traceContext).addScopedInterceptor(new MySQLCreateConnectionInterceptor(agent.traceContext))
+  InstrumentMethod.make(mysql, 'createConnection').addScopedInterceptor(new MySQLCreateConnectionInterceptor())
 
   return mysql
 }