fix the local span cannot ref other trace segment

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/TracingContext.java

@@ -16,7 +16,6 @@
  *
  */
 
-
 package org.apache.skywalking.apm.agent.core.context;
 
 import java.util.LinkedList;
@@ -50,6 +49,7 @@ import org.apache.skywalking.apm.agent.core.sampling.SamplingService;
  * ContextCarrier} or {@link ContextSnapshot}.
  *
  * @author wusheng
+ * @author zhang xin
  */
 public class TracingContext implements AbstractTracerContext {
     /**
@@ -159,7 +159,7 @@ public class TracingContext implements AbstractTracerContext {
         this.segment.relatedGlobalTraces(carrier.getDistributedTraceId());
         AbstractSpan span = this.activeSpan();
         if (span instanceof EntrySpan) {
-            ((EntrySpan)span).ref(ref);
+            span.ref(ref);
         }
     }
 
@@ -213,7 +213,9 @@ public class TracingContext implements AbstractTracerContext {
      */
     @Override
     public void continued(ContextSnapshot snapshot) {
-        this.segment.ref(new TraceSegmentRef(snapshot));
+        TraceSegmentRef segmentRef = new TraceSegmentRef(snapshot);
+        this.segment.ref(segmentRef);
+        this.activeSpan().ref(segmentRef);
         this.segment.relatedGlobalTraces(snapshot.getDistributedTraceId());
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/AbstractSpan.java

@@ -16,7 +16,6 @@
  *
  */
 
-
 package org.apache.skywalking.apm.agent.core.context.trace;
 
 import java.util.Map;
@@ -110,4 +109,11 @@ public interface AbstractSpan {
     String getOperationName();
 
     AbstractSpan setOperationId(int operationId);
+
+    /**
+     * Reference other trace segment.
+     *
+     * @param ref segment ref
+     */
+    void ref(TraceSegmentRef ref);
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/AbstractTracingSpan.java

@@ -64,6 +64,13 @@ public abstract class AbstractTracingSpan implements AbstractSpan {
      */
     protected List<LogDataEntity> logs;
 
+    /**
+     * The refs of parent trace segments, except the primary one. For most RPC call, {@link #refs} contains only one
+     * element, but if this segment is a start span of batch process, the segment faces multi parents, at this moment,
+     * we use this {@link #refs} to link them.
+     */
+    protected List<TraceSegmentRef> refs;
+
     protected AbstractTracingSpan(int spanId, int parentSpanId, String operationName) {
         this.operationName = operationName;
         this.operationId = DictionaryUtil.nullValue();
@@ -276,4 +283,13 @@ public abstract class AbstractTracingSpan implements AbstractSpan {
 
         return spanBuilder;
     }
+
+    @Override public void ref(TraceSegmentRef ref) {
+        if (refs == null) {
+            refs = new LinkedList<TraceSegmentRef>();
+        }
+        if (!refs.contains(ref)) {
+            refs.add(ref);
+        }
+    }
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/EntrySpan.java

@@ -38,12 +38,6 @@ import org.apache.skywalking.apm.network.trace.component.Component;
  * @author wusheng
  */
 public class EntrySpan extends StackBasedTracingSpan {
-    /**
-     * The refs of parent trace segments, except the primary one. For most RPC call, {@link #refs} contains only one
-     * element, but if this segment is a start span of batch process, the segment faces multi parents, at this moment,
-     * we use this {@link #refs} to link them.
-     */
-    private List<TraceSegmentRef> refs;
 
     private int currentMaxDepth;
 
@@ -146,15 +140,6 @@ public class EntrySpan extends StackBasedTracingSpan {
         return builder;
     }
 
-    public void ref(TraceSegmentRef ref) {
-        if (refs == null) {
-            refs = new LinkedList<TraceSegmentRef>();
-        }
-        if (!refs.contains(ref)) {
-            refs.add(ref);
-        }
-    }
-
     private void clearWhenRestart() {
         this.componentId = DictionaryUtil.nullValue();
         this.componentName = null;

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/trace/NoopSpan.java

@@ -99,4 +99,7 @@ public class NoopSpan implements AbstractSpan {
     @Override public AbstractSpan setOperationId(int operationId) {
         return this;
     }
+
+    @Override public void ref(TraceSegmentRef ref) {
+    }
 }