[#231] Support ChildTrace

Author: feelform

lib/context/async-id.js

@@ -15,7 +15,7 @@ class AsyncId {
 
   // DefaultAsyncIdGenerator.java: nextAsyncId()
   static make() {
-    return new AsyncId(AsyncId.asyncIdGenerator.next, 0)
+    return new AsyncId(AsyncId.asyncIdGenerator.getAndIncrement(), 0)
   }
 
   constructor(asyncId, sequence) {

lib/context/remote-trace-root-builder.js

@@ -22,25 +22,39 @@ class RemoteTraceRoot {
     isSampled() {
         return true
     }
+
+    getTraceId() {
+        return this.traceId
+    }
 }
 
 class RemoteTraceRootBuilder {
-    constructor(agentInfo) {
+    constructor(agentInfo, transactionId) {
         this.agentInfo = agentInfo
+        this.transactionId = transactionId
+    }
+
+    make(transactionId) {
+        return new RemoteTraceRootBuilder(this.agentInfo, transactionId)
     }
 
     setTraceId(traceId) {
         this.traceId = traceId
         return this
     }
 
-    build(transactionId) {
-        if (this.traceId) {
-            return new RemoteTraceRoot(this.traceId, new TraceRootBuilder(this.agentInfo.agentId).build(transactionId))
+    build() {
+        const agentId = this.agentInfo.getAgentId()
+        if (this.isNewTraceRoot()) {
+            const traceId = new TraceIdBuilder(this.agentInfo, this.transactionId).build()
+            return new RemoteTraceRoot(traceId, new TraceRootBuilder(agentId, this.transactionId).build())
         }
 
-        const traceId = new TraceIdBuilder(this.agentInfo, transactionId).build()
-        return new RemoteTraceRoot(traceId, new TraceRootBuilder(this.agentInfo.agentId).build(transactionId))
+        return new RemoteTraceRoot(this.traceId, new TraceRootBuilder(agentId, this.transactionId).build())
+    }
+
+    isNewTraceRoot() {
+        return !this.traceId
     }
 }
 

lib/context/span-builder.js

@@ -64,6 +64,10 @@ class SpanBuilder {
         return this
     }
 
+    getTraceRoot() {
+        return this.traceRoot
+    }
+
     build() {
         const span = new Span(this.traceRoot)
         this.startTime = this.startTime || Date.now()

lib/context/span-id.js

@@ -8,19 +8,29 @@
 
 // SpanId.java in Java agent
 class SpanId {
-  static nullSpanId = -1
-  
-  constructor () {
+  constructor() {
     this.MAX_NUM = Number.MAX_SAFE_INTEGER
   }
 
-  get next () {
+  nullSpanId() {
+    return '-1'
+  }
+
+  get next() {
     return Math.floor(Math.random() * this.MAX_NUM)
   }
 
   newSpanId() {
     return this.next.toString()
   }
+
+  nextSpanId(spanId, parentSpanId) {
+    let newId = this.newSpanId()
+    while (newId === spanId || newId === parentSpanId) {
+      newId = this.newSpanId()
+    }
+    return newId
+  }
 }
 
 module.exports = new SpanId()

lib/context/trace-context.js

@@ -20,6 +20,8 @@ const SpanBuilder = require('./span-builder')
 const SpanChunkBuilder = require('./span-chunk-builder')
 const SpanRepository = require('./trace/span-repository')
 const Trace2 = require('./trace/trace2')
+const AsyncSpanChunkBuilder = require('./trace/async-span-chunk-builder')
+const ChildTraceBuilder = require('./trace/child-trace-builder')
 
 class TraceContext {
   constructor(agentInfo, dataSender, config) {
@@ -33,6 +35,10 @@ class TraceContext {
     this.traceSampler = new TraceSampler(agentInfo, config)
   }
 
+  getAgentInfo() {
+    return this.agentInfo
+  }
+
   continueTraceObject(requestData) {
     const traceId = new TraceId(
       requestData.transactionId,
@@ -125,14 +131,16 @@ class TraceContext {
     const spanBuilder = new SpanBuilder(traceRoot)
     const spanChunkBuilder = new SpanChunkBuilder(traceRoot)
     const repository = new SpanRepository(spanChunkBuilder, this.dataSender)
-    return new Trace2(spanBuilder, repository, this.agentInfo.getServiceType())
+    return new Trace2(spanBuilder, repository)
   }
 
   // DefaultAsyncContext.java: newAsyncContextTrace
   // DefaultBaseTraceFactory.java: continueAsyncContextTraceObject
-  continueAsyncContextTraceObject(traceRoot, asyncId) {
-    const localAsyncId = asyncId.nextLocalAsyncId2()
-    
+  // AsyncContextSpanEventEndPointApiAwareInterceptor.java : before
+  continueAsyncContextTraceObject(traceRoot, localAsyncId) {
+    const spanChunkBuilder = new AsyncSpanChunkBuilder(traceRoot, localAsyncId)
+    const repository = new SpanRepository(spanChunkBuilder, this.dataSender)
+    return new ChildTraceBuilder(traceRoot, repository, localAsyncId)
   }
 }
 

lib/context/trace-root-builder.js

@@ -23,12 +23,19 @@ class TraceRoot {
 }
 
 class TraceRootBuilder {
-    constructor(agentId) {
+    constructor(agentId, localTransactionId) {
         this.agentId = agentId
+        this.localTransactionId = localTransactionId
+        this.traceStartTime = Date.now()
     }
 
-    build(transactionId) {
-        return new TraceRoot(this.agentId, Date.now(), transactionId)
+    make(localTransactionId) {
+        return new TraceRootBuilder(this.agentId, localTransactionId)
+    }
+
+    // DefaultTraceRootFactory.java: newDisableTraceRoot
+    build() {
+        return new TraceRoot(this.agentId, this.traceStartTime, this.localTransactionId)
     }
 }
 

lib/context/trace/async-span-chunk-builder.js

@@ -19,6 +19,10 @@ class AsyncSpanChunk extends SpanChunk {
         return this.localAsyncId
     }
 
+    isAsyncSpanChunk() {
+        return true
+    }
+
     toString() {
         return `AsyncSpanChunk(traceRoot=${this.traceRoot}, spanEventList=${this.spanEventList}, localAsyncId=${this.localAsyncId})`
     }
@@ -32,7 +36,7 @@ class AsyncSpanChunkBuilder {
     }
 
     build(spanEventList) {
-        return new AsyncSpanChunk(this.traceRoot, this.asyncId, spanEventList)
+        return new AsyncSpanChunk(this.traceRoot, spanEventList, this.asyncId)
     }
 
     toString() {

lib/context/trace/call-stack.js

@@ -0,0 +1,49 @@
+/**
+ * Pinpoint Node.js Agent
+ * Copyright 2020-present NAVER Corp.
+ * Apache License v2.0
+ */
+
+'use strict'
+
+const SpanEventBuilder = require('./span-event-builder')
+const SpanEventRecorder = require('./span-event-recorder2')
+
+/**
+ * DefaultCallStack.java in Java agent
+ */
+class CallStack {
+    constructor() {
+        this.stack = []
+        this.sequence = 0
+        this.depth = 0
+    }
+
+    makeSpanEventRecorder(stackId) {
+        const recorder = new SpanEventRecorder(SpanEventBuilder.make(stackId))
+        this.push(recorder.getSpanEventBuilder())
+        return recorder
+    }
+
+    push(spanEventBuilder) {
+        if (spanEventBuilder.needsSequence()) {
+            spanEventBuilder.setSequence(this.sequence++)
+        }
+
+        if (spanEventBuilder.needsDepth()) {
+            spanEventBuilder.setDepth(this.stack.length + 1)
+        }
+
+        this.stack.push(spanEventBuilder)
+    }
+
+    // pop in java agent
+    pop() {
+        if (this.stack.length < 1) {
+            return SpanEventBuilder.nullObject()
+        }
+        return this.stack.pop()
+    }
+}
+
+module.exports = CallStack

lib/context/trace/child-trace-builder.js

@@ -6,8 +6,11 @@
 
 'use strict'
 
-const SpanEventRecorder = require('./span-event-recorder')
+const SpanEventRecorder = require('./span-event-recorder2')
 const Trace = require('./trace2')
+const TraceRootSpanRecorder = require('./trace-root-span-recorder')
+const StackId = require('./stack-id')
+const CallStack = require('./call-stack')
 
 class ChildTrace extends Trace {
     constructor(spanBuilder, repository, localAsyncId) {
@@ -17,14 +20,48 @@ class ChildTrace extends Trace {
 }
 
 class ChildTraceBuilder {
-    constructor(traceRoot, localAsyncId) {
+    constructor(traceRoot, repository, localAsyncId) {
         this.traceRoot = traceRoot
+        this.repository = repository
+        this.spanRecorder = new TraceRootSpanRecorder(traceRoot)
         this.localAsyncId = localAsyncId
 
-        this.callStack = []
+        this.callStack = new CallStack()
         this.closed = false
 
-        this.spanEventRecorder = SpanEventRecorder.nullObject()
+        this.traceBlockBegin(StackId.asyncBeginStackId)
+    }
+
+    traceBlockBegin(stackId = StackId.default) {
+        if (this.closed) {
+            return SpanEventRecorder.nullObject()
+        }
+
+        return this.callStack.makeSpanEventRecorder(stackId)
+    }
+
+    traceBlockEnd() {
+        if (this.closed) {
+            return
+        }
+
+        const spanEventBuilder = this.callStack.pop()
+        spanEventBuilder.markAfterTime()
+        this.repository.storeSpanEvent(spanEventBuilder)
+    }
+
+    getTraceRoot() {
+        return this.traceRoot
+    }
+
+    getTraceId() {
+        return this.traceRoot.getTraceId()
+    }
+
+    close() {
+        this.traceBlockEnd(StackId.asyncBeginStackId)
+        this.closed = true
+        this.repository.flush()
     }
 }