package com.a.eye.skywalking.trace;

import com.a.eye.skywalking.messages.ISerializable;
import com.a.eye.skywalking.trace.proto.SegmentMessage;
import com.a.eye.skywalking.trace.proto.SegmentRefMessage;
import com.a.eye.skywalking.trace.proto.SpanMessage;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

/**
 * {@link TraceSegment} is a segment or fragment of the distributed trace.
 * {@see https://github.com/opentracing/specification/blob/master/specification.md#the-opentracing-data-model}
 * A {@link
 * TraceSegment} means the segment, which exists in current {@link Thread}. And the distributed trace is formed by multi
 * {@link TraceSegment}s, because the distributed trace crosses multi-processes, multi-threads.
 *
 * Created by wusheng on 2017/2/17.
 */
public class TraceSegment implements ISerializable<SegmentMessage> {
    /**
     * The id of this trace segment.
     * Every segment has its unique-global-id.
     */
    private String traceSegmentId;

    /**
     * The start time of this trace segment.
     */
    private long startTime;

    /**
     * The end time of this trace segment.
     */
    private long endTime;

    /**
     * The primary ref of the parent trace segment.
     * Use {@link TraceSegmentRef}, we can link this trace segment to the primary parent segment.
     */
    private TraceSegmentRef primaryRef;

    /**
     * The refs of other parent trace segments, except the primary one.
     * For most RPC call, {@link #refs} stay in null,
     * 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;

    /**
     * The spans belong to this trace segment.
     * They all have finished.
     * All active spans are hold and controlled by "skywalking-api" module.
     */
    private List<Span> spans;

    /**
     * The <code>applicationCode</code> represents a name of current application/JVM and indicates which is business
     * role in the cluster.
     *
     * e.g. account_app, billing_app
     */
    private String applicationCode;

    /**
     * Create a trace segment, by given segmentId.
     * This segmentId is generated by TraceSegmentRef, AKA, from tracer/agent module.
     *
     * @param segmentId {@link #traceSegmentId}
     */
    public TraceSegment(String segmentId, String applicationCode) {
        this.traceSegmentId = segmentId;
        this.applicationCode = applicationCode;
        this.startTime = System.currentTimeMillis();
        this.spans = new LinkedList<Span>();
    }

    /**
     * Create a trace segment, by given {@link SegmentMessage}
     *
     * @param message from another {@link TraceSegment#serialize()}
     */
    public TraceSegment(SegmentMessage message) {
        deserialize(message);
    }

    /**
     * Establish the link between this segment and its parents.
     * When {@param primaryOnly} is true;
     * The first time, you {@link #ref(TraceSegmentRef)} to parent, it is affirmed as {@link #primaryRef}.
     * And others are affirmed as {@link #refs}.
     *
     * @param refSegment {@link TraceSegmentRef}
     * @param primaryOnly if true, set {@param refSegment} to {@link #primaryRef} only.
     */
    public void ref(TraceSegmentRef refSegment, boolean primaryOnly) {
        if(primaryOnly){
            if (primaryRef == null) {
                primaryRef = refSegment;
            }
        }else {
            if (primaryRef == null) {
                primaryRef = refSegment;
            } else {
                if (refs == null) {
                    refs = new LinkedList<TraceSegmentRef>();
                }
                refs.add(refSegment);
            }
        }
    }

    /**
     * Set to {@link #primaryRef} only,
     * based on {@link #ref(TraceSegmentRef, boolean)}
     *
     * @param refSegment {@link TraceSegmentRef}
     */
    public void ref(TraceSegmentRef refSegment){
        ref(refSegment, true);
    }

    /**
     * After {@link Span} is finished, as be controller by "skywalking-api" module,
     * notify the {@link TraceSegment} to archive it.
     *
     * @param finishedSpan
     */
    public void archive(Span finishedSpan) {
        spans.add(finishedSpan);
    }

    /**
     * Finish this {@link TraceSegment}.
     *
     * return this, for chaining
     */
    public TraceSegment finish() {
        this.endTime = System.currentTimeMillis();
        return this;
    }

    public String getTraceSegmentId() {
        return traceSegmentId;
    }

    public long getStartTime() {
        return startTime;
    }

    public long getEndTime() {
        return endTime;
    }

    public TraceSegmentRef getPrimaryRef() {
        return primaryRef;
    }

    public List<TraceSegmentRef> getRefs() {
        return Collections.unmodifiableList(refs);
    }

    public List<Span> getSpans() {
        return Collections.unmodifiableList(spans);
    }

    public String getApplicationCode() {
        return applicationCode;
    }

    @Override
    public String toString() {
        return "TraceSegment{" +
            "traceSegmentId='" + traceSegmentId + '\'' +
            ", endTime=" + endTime +
            ", primaryRef=" + primaryRef +
            ", spans.size=" + spans.size() +
            '}';
    }

    @Override
    public SegmentMessage serialize() {
        SegmentMessage.Builder segmentBuilder = SegmentMessage.newBuilder();
        segmentBuilder.setTraceSegmentId(traceSegmentId);
        segmentBuilder.setStartTime(startTime);
        segmentBuilder.setEndTime(endTime);
        segmentBuilder.setApplicationCode(applicationCode);
        segmentBuilder.setPrimaryRef(primaryRef.serialize());
        for (TraceSegmentRef ref : refs) {
            segmentBuilder.addRefs(ref.serialize());
        }
        for (Span span : spans) {
            segmentBuilder.addSpans(span.serialize());
        }
        return segmentBuilder.build();
    }

    @Override
    public void deserialize(SegmentMessage message) {
        traceSegmentId = message.getTraceSegmentId();
        startTime = message.getStartTime();
        endTime = message.getEndTime();
        applicationCode = message.getApplicationCode();
        (primaryRef = new TraceSegmentRef()).deserialize(message.getPrimaryRef());
        List<SegmentRefMessage> refsList = message.getRefsList();
        if (refsList != null) {
            this.refs = new LinkedList<TraceSegmentRef>();
            for (SegmentRefMessage refMessage : refsList) {
                TraceSegmentRef ref = new TraceSegmentRef();
                ref.deserialize(refMessage);
                refs.add(ref);
            }
        }

        List<SpanMessage> spansList = message.getSpansList();
        if (spansList != null) {
            this.spans = new LinkedList<Span>();
            for (SpanMessage spanMessage : spansList) {
                spans.add(new Span(spanMessage));
            }
        }
    }
}

@@ -4853,24 +4853,85 @@
 getRefs() %7B%0A
+        if(refs == null)%7B%0A            return null;%0A        %7D%0A
         retu

Fix NullPointException