From 315d31f9c306377922086f19cd33f8451b149f6e Mon Sep 17 00:00:00 2001
From: chongli8 <chongli8@creditease.cn>
Date: Wed, 27 Mar 2019 17:35:30 +0800
Subject: [PATCH] invokechain support custom async
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

调用链支持异步，really！！

	支持用户自定义的异步接口。如应用接收到请求之后，用户启动一个线程池处理当前请求，而不等待处理结果立即返回情况。这种情况下调用链也可连起来。
---
 .../uav/apm/invokechain/span/SpanFactory.java | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/com.creditease.uav.monitorframework.apm/src/main/java/com/creditease/uav/apm/invokechain/span/SpanFactory.java b/com.creditease.uav.monitorframework.apm/src/main/java/com/creditease/uav/apm/invokechain/span/SpanFactory.java
index 4bd32292..e3980874 100644
--- a/com.creditease.uav.monitorframework.apm/src/main/java/com/creditease/uav/apm/invokechain/span/SpanFactory.java
+++ b/com.creditease.uav.monitorframework.apm/src/main/java/com/creditease/uav/apm/invokechain/span/SpanFactory.java
@@ -49,6 +49,7 @@ private static class TraceIDGen {
         private volatile AtomicLong count = new AtomicLong(0);
 
         public TraceIDGen() {
+
             ip = NetworkHelper.getLocalIP();
         }
 
@@ -103,6 +104,8 @@ protected Map<String, Span> initialValue() {
         }
     };
 
+    public static final TransmittableThreadLocal<Span> mainThreadLocalTracker = new TransmittableThreadLocal<Span>();
+
     /**
      * 通过一个父span来，创建一个span
      *
@@ -215,6 +218,10 @@ public Map<String, Span> getSpansFromContext() {
      */
     public Span getSpanFromContext(String key) {
 
+        if (key.equals("main")) {
+            return mainThreadLocalTracker.get();
+        }
+
         Map<String, Span> spanMap = threadLocalTracker.get();
 
         return spanMap.get(key);
@@ -228,6 +235,12 @@ public Span getSpanFromContext(String key) {
      */
     public Span getRemoveSpanFromContext(String key) {
 
+        if (key.equals("main")) {
+            Span span = mainThreadLocalTracker.get();
+            mainThreadLocalTracker.remove();
+            return span;
+        }
+
         Map<String, Span> spanMap = threadLocalTracker.get();
 
         return spanMap.remove(key);
@@ -240,6 +253,10 @@ public Span getRemoveSpanFromContext(String key) {
      */
     public void setSpanToContext(String key, Span span) {
 
+        if (key.equals("main")) {
+            mainThreadLocalTracker.set(span);
+            return;
+        }
         Map<String, Span> spanMap = threadLocalTracker.get();
 
         spanMap.put(key, span);
@@ -250,7 +267,8 @@ public void setSpanToContext(String key, Span span) {
      */
     public void removeCurrentThreadValue() {
 
+        mainThreadLocalTracker.removeCurrent();
         threadLocalTracker.removeCurrent();
     }
 
-}
+}
\ No newline at end of file