Resolves state conflict between client and local spans

Before, local spans used the same thread local as client ones. This
meant a client span in the same trace would overwite a local span.

Ex. [server -> local -> client] would truncate to [server -> client]

This change adds state to ensure local spans aren't overwritten by
client spans.

This change does not introduce a state stack, eventhough local spans
should be nestable. The reason is that Brave's model isn't currently
designed for that and a nested model would significantly impact the api.

brave-core/src/main/java/com/github/kristofa/brave/Brave.java

@@ -27,14 +27,14 @@ public class Brave {
      * Builds Brave api objects with following defaults if not overridden:
      * <p>
      * <ul>
-     * <li>ThreadLocalServerAndClientSpanState which binds trace/span state to current thread.</li>
+     * <li>ThreadLocalServerClientAndLocalSpanState which binds trace/span state to current thread.</li>
      * <li>FixedSampleRateTraceFilter which traces every request.</li>
      * <li>LoggingSpanCollector</li>
      * </ul>
      */
     public static class Builder {
 
-        private final ServerAndClientSpanState state;
+        private final ServerClientAndLocalSpanState state;
         private final List<TraceFilter> traceFilters = new ArrayList<>();
         private SpanCollector spanCollector = new LoggingSpanCollector();
         private Random random = new Random();
@@ -64,7 +64,7 @@ public Builder() {
         public Builder(String serviceName) {
             try {
                 int ip = toInt(getLocalHostLANAddress());
-                state = new ThreadLocalServerAndClientSpanState(ip, 0, serviceName);
+                state = new ThreadLocalServerClientAndLocalSpanState(ip, 0, serviceName);
             } catch (UnknownHostException e) {
                 throw new IllegalStateException("Unable to get Inet address", e);
             }
@@ -79,14 +79,14 @@ public Builder(String serviceName) {
          * @param serviceName Name of service. Is only relevant when we do server side tracing.
          */
         public Builder(int ip, int port, String serviceName) {
-            state = new ThreadLocalServerAndClientSpanState(ip, port, serviceName);
+            state = new ThreadLocalServerClientAndLocalSpanState(ip, port, serviceName);
             traceFilters.add(new FixedSampleRateTraceFilter(1));
         }
 
         /**
          * Use for control of how tracing state propagates across threads.
          */
-        public Builder(ServerAndClientSpanState state) {
+        public Builder(ServerClientAndLocalSpanState state) {
             this.state = Util.checkNotNull(state, "state must be specified.");
             traceFilters.add(new FixedSampleRateTraceFilter(1));
         }
@@ -209,7 +209,12 @@ private Brave(Builder builder) {
                 .state(builder.state)
                 .traceFilters(builder.traceFilters).build();
 
-        localTracer = LocalTracer.create(clientTracer, builder.spanCollector, builder.state);
+        localTracer = LocalTracer.builder()
+                .randomGenerator(builder.random)
+                .spanCollector(builder.spanCollector)
+                .spanAndEndpoint(SpanAndEndpoint.LocalSpanAndEndpoint.create(builder.state))
+                .traceFilters(builder.traceFilters).build();
+
         serverRequestInterceptor = new ServerRequestInterceptor(serverTracer);
         serverResponseInterceptor = new ServerResponseInterceptor(serverTracer);
         clientRequestInterceptor = new ClientRequestInterceptor(clientTracer);

brave-core/src/main/java/com/github/kristofa/brave/ClientSpanState.java

@@ -5,8 +5,15 @@
 import com.twitter.zipkin.gen.Span;
 
 /**
- * Maintains client span state.
- * 
+ * Maintains state for a single client span.
+ *
+ * <p/>Client spans can be at the following locations in the span tree.
+ * <ul>
+ *     <li>The root-span of a trace originated by Brave</li>
+ *     <li>A child of a server span originated by Brave</li>
+ *     <li>A child of a local span originated by Brave</li>
+ * </ul>
+ *
  * @author kristof
  */
 public interface ClientSpanState extends CommonSpanState {

brave-core/src/main/java/com/github/kristofa/brave/ClientTracer.java

@@ -42,7 +42,7 @@ public static Builder builder() {
     @AutoValue.Builder
     public abstract static class Builder {
 
-        public Builder state(ServerAndClientSpanState state) {
+        public Builder state(ServerClientAndLocalSpanState state) {
             return spanAndEndpoint(ClientSpanAndEndpoint.create(state));
         }
 
@@ -148,13 +148,19 @@ public void setCurrentClientServiceName(String serviceName) {
 
     private SpanId getNewSpanId() {
 
-        Span currentServerSpan = spanAndEndpoint().state().getCurrentServerSpan().getSpan();
+        Span parentSpan = spanAndEndpoint().state().getCurrentLocalSpan();
+        if (parentSpan == null) {
+            ServerSpan serverSpan = spanAndEndpoint().state().getCurrentServerSpan();
+            if (serverSpan != null) {
+                parentSpan = serverSpan.getSpan();
+            }
+        }
         long newSpanId = randomGenerator().nextLong();
-        if (currentServerSpan == null) {
+        if (parentSpan == null) {
             return SpanId.create(newSpanId, newSpanId, null);
         }
 
-        return SpanId.create(currentServerSpan.getTrace_id(), newSpanId, currentServerSpan.getId());
+        return SpanId.create(parentSpan.getTrace_id(), newSpanId, parentSpan.getId());
     }
 
     ClientTracer() {

brave-core/src/main/java/com/github/kristofa/brave/LocalSpanState.java

@@ -0,0 +1,35 @@
+package com.github.kristofa.brave;
+
+import com.github.kristofa.brave.internal.Nullable;
+import com.twitter.zipkin.gen.Span;
+
+/**
+ * Maintains state for a single local span. This means nesting is not supported.
+ *
+ * <p/>Local spans can be at the following locations in the span tree.
+ * <ul>
+ *     <li>The root-span of a trace originated by Brave</li>
+ *     <li>A child of a server span originated by Brave</li>
+ * </ul>
+ */
+public interface LocalSpanState extends CommonSpanState {
+
+    /**
+     * Gets the Span for the local request that was started as part of current request.
+     * <p/>
+     * Should be thread-aware to support multiple parallel requests.
+     * 
+     * @return Local request span for current thread.
+     */
+    @Nullable
+    Span getCurrentLocalSpan();
+
+    /**
+     * Sets current local span.
+     * <p/>
+     * Should be thread-aware to support multiple parallel requests.
+     * 
+     * @param span Local span.
+     */
+    void setCurrentLocalSpan(Span span);
+}

brave-core/src/main/java/com/github/kristofa/brave/LocalTracer.java

@@ -1,9 +1,17 @@
 package com.github.kristofa.brave;
 
+import com.github.kristofa.brave.SpanAndEndpoint.LocalSpanAndEndpoint;
+import com.github.kristofa.brave.internal.Util;
 import com.google.auto.value.AutoValue;
+import com.twitter.zipkin.gen.AnnotationType;
+import com.twitter.zipkin.gen.BinaryAnnotation;
 import com.twitter.zipkin.gen.Span;
 import com.twitter.zipkin.gen.zipkinCoreConstants;
 
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Random;
+
 import static com.twitter.zipkin.gen.zipkinCoreConstants.LOCAL_COMPONENT;
 
 /**
@@ -18,7 +26,7 @@
  *
  * <p/>Here's an example of allocating precise duration for a local span:
  * <pre>
- * tracer.startSpan("codec", "encode");
+ * tracer.startNewSpan("codec", "encode");
  * try {
  *   return codec.encode(input);
  * } finally {
@@ -29,14 +37,38 @@
  * @see zipkinCoreConstants#LOCAL_COMPONENT
  */
 @AutoValue
-public abstract class LocalTracer {
+public abstract class LocalTracer extends AnnotationSubmitter {
+
+    static Builder builder() {
+        return new AutoValue_LocalTracer.Builder();
+    }
+
+    // visible for testing
+    static Builder builder(LocalTracer source) {
+        return new AutoValue_LocalTracer.Builder(source);
+    }
+
+    @Override
+    abstract LocalSpanAndEndpoint spanAndEndpoint();
+
+    abstract Random randomGenerator();
 
-    abstract ClientTracer clientTracer();
     abstract SpanCollector spanCollector();
-    abstract ServerAndClientSpanState state();
 
-    static LocalTracer create(ClientTracer clientTracer, SpanCollector spanCollector, ServerAndClientSpanState state) {
-        return new AutoValue_LocalTracer(clientTracer, spanCollector, state);
+    abstract List<TraceFilter> traceFilters();
+
+    @AutoValue.Builder
+    abstract static class Builder {
+
+        abstract Builder spanAndEndpoint(LocalSpanAndEndpoint spanAndEndpoint);
+
+        abstract Builder randomGenerator(Random randomGenerator);
+
+        abstract Builder spanCollector(SpanCollector spanCollector);
+
+        abstract Builder traceFilters(List<TraceFilter> traceFilters);
+
+        abstract LocalTracer build();
     }
 
     /**
@@ -47,16 +79,23 @@ static LocalTracer create(ClientTracer clientTracer, SpanCollector spanCollector
      * @return metadata about the new span or null if one wasn't started due to sampling policy.
      * @see zipkinCoreConstants#LOCAL_COMPONENT
      */
-    public SpanId startSpan(String component, String operation) {
-        SpanId spanId = clientTracer().startNewSpan(operation);
+    public SpanId startNewSpan(String component, String operation) {
+        SpanId spanId = startNewSpan(component, operation, currentTimeMicroseconds());
         if (spanId == null) return null;
-        clientTracer().submitBinaryAnnotation(LOCAL_COMPONENT, component);
-        state().getCurrentClientSpan()
-                .setTimestamp(clientTracer().currentTimeMicroseconds())
-                .startTick = System.nanoTime(); // embezzle start tick into an internal field.
+        spanAndEndpoint().span().startTick = System.nanoTime(); // embezzle start tick into an internal field.
         return spanId;
     }
 
+    private SpanId getNewSpanId() {
+        Span currentServerSpan = spanAndEndpoint().state().getCurrentServerSpan().getSpan();
+        long newSpanId = randomGenerator().nextLong();
+        if (currentServerSpan == null) {
+            return SpanId.create(newSpanId, newSpanId, null);
+        }
+
+        return SpanId.create(currentServerSpan.getTrace_id(), newSpanId, currentServerSpan.getId());
+    }
+
     /**
      * Request a new local span, which started at the given timestamp.
      *
@@ -66,42 +105,40 @@ public SpanId startSpan(String component, String operation) {
      * @return metadata about the new span or null if one wasn't started due to sampling policy.
      * @see zipkinCoreConstants#LOCAL_COMPONENT
      */
-    public SpanId startSpan(String component, String operation, long timestamp) {
-        SpanId spanId = clientTracer().startNewSpan(operation);
-        if (spanId == null) return null;
-        clientTracer().submitBinaryAnnotation(LOCAL_COMPONENT, component);
-        state().getCurrentClientSpan().setTimestamp(timestamp);
-        return spanId;
-    }
+    public SpanId startNewSpan(String component, String operation, long timestamp) {
 
-    /**
-     * Associates an event that explains latency with the current system time.
-     *
-     * @param value A short tag indicating the event, like "ApplicationReady"
-     */
-    public void submitAnnotation(String value) {
-        clientTracer().submitAnnotation(value);
-    }
+        Boolean sample = spanAndEndpoint().state().sample();
+        if (Boolean.FALSE.equals(sample)) {
+            spanAndEndpoint().state().setCurrentLocalSpan(null);
+            return null;
+        }
 
-    /**
-     * Associates an event that explains latency with a timestamp.
-     *
-     * @param value     A short tag indicating the event, like "ApplicationReady"
-     * @param timestamp microseconds from epoch
-     */
-    public void submitAnnotation(String value, long timestamp) {
-        clientTracer().submitAnnotation(value, timestamp);
-    }
+        SpanId newSpanId = getNewSpanId();
+        if (sample == null) {
+            // No sample indication is present.
+            for (TraceFilter traceFilter : traceFilters()) {
+                if (!traceFilter.trace(newSpanId.getSpanId(), operation)) {
+                    spanAndEndpoint().state().setCurrentLocalSpan(null);
+                    return null;
+                }
+            }
+        }
 
-    /**
-     * Binary annotations are tags applied to a Span to give it context. For
-     * example, a key "your_app.version" would let you lookup spans by version.
-     *
-     * @param key   Name used to lookup spans, such as "your_app.version"
-     * @param value String value, should not be <code>null</code>.
-     */
-    public void submitBinaryAnnotation(String key, String value) {
-        clientTracer().submitBinaryAnnotation(key, value);
+        Span newSpan = new Span();
+        newSpan.setId(newSpanId.getSpanId());
+        newSpan.setTrace_id(newSpanId.getTraceId());
+        if (newSpanId.getParentSpanId() != null) {
+            newSpan.setParent_id(newSpanId.getParentSpanId());
+        }
+        newSpan.setName(operation);
+        newSpan.setTimestamp(timestamp);
+        newSpan.addToBinary_annotations(new BinaryAnnotation()
+                .setKey(LOCAL_COMPONENT)
+                .setValue(ByteBuffer.wrap(component.getBytes(Util.UTF_8)))
+                .setAnnotation_type(AnnotationType.STRING)
+                .setHost(spanAndEndpoint().endpoint()));
+        spanAndEndpoint().state().setCurrentLocalSpan(newSpan);
+        return newSpanId;
     }
 
     /**
@@ -110,15 +147,15 @@ public void submitBinaryAnnotation(String key, String value) {
     public void finishSpan() {
         long endTick = System.nanoTime();
 
-        Span span = state().getCurrentClientSpan();
+        Span span = spanAndEndpoint().span();
         if (span == null) return;
 
         Long startTick = span.startTick;
         final long duration;
         if (startTick != null) {
             duration = (endTick - startTick) / 1000;
         } else {
-            duration = clientTracer().currentTimeMicroseconds() - span.getTimestamp();
+            duration = currentTimeMicroseconds() - span.getTimestamp();
         }
         finishSpan(duration);
     }
@@ -127,16 +164,15 @@ public void finishSpan() {
      * Completes the span, which took {@code duration} microseconds.
      */
     public void finishSpan(long duration) {
-        Span span = state().getCurrentClientSpan();
+        Span span = spanAndEndpoint().span();
         if (span == null) return;
 
         synchronized (span) {
             span.setDuration(duration);
             spanCollector().collect(span);
         }
 
-        state().setCurrentClientSpan(null);
-        state().setCurrentClientServiceName(null);
+        spanAndEndpoint().state().setCurrentLocalSpan(null);
     }
 
     LocalTracer() {

brave-core/src/main/java/com/github/kristofa/brave/ServerAndClientSpanState.java → brave-core/src/main/java/com/github/kristofa/brave/ServerClientAndLocalSpanState.java

@@ -5,6 +5,6 @@
  * 
  * @author kristof
  */
-public interface ServerAndClientSpanState extends ServerSpanState, ClientSpanState {
+public interface ServerClientAndLocalSpanState extends ServerSpanState, ClientSpanState, LocalSpanState {
 
 }

brave-core/src/main/java/com/github/kristofa/brave/ServerSpanState.java

@@ -5,8 +5,14 @@
 import com.twitter.zipkin.gen.Endpoint;
 
 /**
- * Maintains server span state.
- * 
+ * Maintains state for a single server span.
+ *
+ * <p/>Server spans can be at the following locations in the span tree.
+ * <ul>
+ *     <li>The root-span of a trace originated by Brave</li>
+ *     <li>A child of a span propagated to Brave</li>
+ * </ul>
+ *
  * @author kristof
  */
 public interface ServerSpanState extends CommonSpanState {