Replacing active Object from Span to TraceContext on OT span finish (#…

…459)
elastic · Jan 29, 2019 · f254571 · f254571
1 parent 24424d6
commit f254571
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/...ng-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/ApmSpanInstrumentation.java b/...ng-plugin/src/main/java/co/elastic/apm/agent/opentracing/impl/ApmSpanInstrumentation.java
@@ -22,6 +22,7 @@
 import co.elastic.apm.agent.bci.VisibleForAdvice;
 import co.elastic.apm.agent.impl.transaction.AbstractSpan;
 import co.elastic.apm.agent.impl.transaction.Span;
+import co.elastic.apm.agent.impl.transaction.TraceContext;
 import co.elastic.apm.agent.impl.transaction.Transaction;
 import co.elastic.apm.agent.web.ResultUtil;
 import net.bytebuddy.asm.Advice;
@@ -66,15 +67,16 @@ public FinishInstrumentation() {
 
         @Advice.OnMethodEnter(suppress = Throwable.class)
         private static void finishInternal(@Advice.FieldValue(value = "dispatcher", readOnly = false, typing = Assigner.Typing.DYNAMIC) @Nullable AbstractSpan<?> span,
-                                           @Advice.Argument(0) long finishMicros) {
+                                           @Advice.Argument(0) long finishMicros,
+                                           @Advice.Argument(value = 1, optional = true) @Nullable Object traceContext) {
             if (span != null) {
-                doFinishInternal(span, finishMicros);
+                doFinishInternal(span, finishMicros, traceContext);
                 span = null;
             }
         }
 
         @VisibleForAdvice
-        public static void doFinishInternal(AbstractSpan<?> span, long finishMicros) {
+        public static void doFinishInternal(AbstractSpan<?> span, long finishMicros, @Nullable Object traceContext) {
             if (span.getType() == null) {
                 if (span instanceof Transaction) {
                     Transaction transaction = (Transaction) span;
@@ -95,6 +97,12 @@ public static void doFinishInternal(AbstractSpan<?> span, long finishMicros) {
             } else {
                 span.end();
             }
+
+            // If the finished span is the active span, replace with the corresponding TraceContext
+            if (tracer != null && traceContext != null && span == tracer.getActive() && traceContext instanceof TraceContext) {
+                tracer.deactivate(span);
+                tracer.activate((TraceContext) traceContext);
+            }
         }
     }
 

diff --git a/apm-opentracing/src/main/java/co/elastic/apm/opentracing/ApmSpan.java b/apm-opentracing/src/main/java/co/elastic/apm/opentracing/ApmSpan.java
@@ -91,13 +91,13 @@ public void finish(long finishMicros) {
             if (traceContext != null) {
                 // prevents race conditions with ScopeManager#activate(Span)
                 synchronized (traceContext) {
-                    finishInternal(finishMicros);
+                    finishInternal(finishMicros, traceContext);
                 }
             }
         }
     }
 
-    private void finishInternal(long finishMicros) {
+    private void finishInternal(long finishMicros, Object traceContext) {
         // implementation injected at runtime by co.elastic.apm.agent.opentracing.impl.ApmSpanInstrumentation.FinishInstrumentation.finishInternal
     }