Merge pull request #480 from ieven/master

invokechain support custom async
uavorg · Mar 29, 2019 · 1ed0a81 · 1ed0a81
2 parents 17cc36a + 315d31f
commit 1ed0a81
Showing 1 changed file with 19 additions and 1 deletion.
diff --git a/...nitorframework.apm/src/main/java/com/creditease/uav/apm/invokechain/span/SpanFactory.java b/...nitorframework.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();
     }
 
-}
+}