Marks contexts extracted by AWSPropagation and cuts overhead

This marks contexts extracted by AWSPropagation so that calls to `rootField` are reliable even when AWS isn't the current propagation factory. This also dramatically reduces extraction overhead by avoiding allocation of constants, and early returning when required fields are absent. Before: ``` Benchmark Mode Cnt Score Error Units PropagationBenchmarks.extract_aws thrpt 15 4.145 ± 0.292 ops/us PropagationBenchmarks.extract_aws_nothing thrpt 15 132.670 ± 1.465 ops/us PropagationBenchmarks.extract_b3 thrpt 15 10.697 ± 0.394 ops/us PropagationBenchmarks.extract_b3_nothing thrpt 15 170.315 ± 5.142 ops/us ``` After: (higher numbers are better) ``` Benchmark Mode Cnt Score Error Units PropagationBenchmarks.extract_aws thrpt 15 4.256 ± 0.245 ops/us PropagationBenchmarks.extract_aws_nothing thrpt 15 214.205 ± 4.829 ops/us PropagationBenchmarks.extract_b3 thrpt 15 10.769 ± 0.316 ops/us PropagationBenchmarks.extract_b3_nothing thrpt 15 268.972 ± 12.388 ops/us ```
daxwang · Oct 16, 2017 · 1256c82 · 1256c82
1 parent c9567c7
commit 1256c82
Show file tree

Hide file tree

Showing 12 changed files with 228 additions and 98 deletions.
diff --git a/brave/src/main/java/brave/propagation/B3Propagation.java b/brave/src/main/java/brave/propagation/B3Propagation.java
@@ -1,29 +1,30 @@
 package brave.propagation;
 
-import brave.internal.HexCodec;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 
 import static brave.internal.HexCodec.lowerHexToUnsignedLong;
+import static brave.internal.HexCodec.toLowerHex;
 
 /**
  * Implements <a href="https://github.com/openzipkin/b3-propagation">B3 Propagation</a>
  */
 public final class B3Propagation<K> implements Propagation<K> {
-  public static final class Factory extends Propagation.Factory {
+
+  public static final Propagation.Factory FACTORY = new Propagation.Factory() {
     @Override public <K> Propagation<K> create(KeyFactory<K> keyFactory) {
-      return B3Propagation.create(keyFactory);
+      return new B3Propagation<>(keyFactory);
     }
 
     @Override public boolean supportsJoin() {
       return true;
     }
-  }
 
-  public static <K> B3Propagation<K> create(KeyFactory<K> keyFactory) {
-    return new B3Propagation<>(keyFactory);
-  }
+    @Override public String toString() {
+      return "B3PropagationFactory";
+    }
+  };
 
   /**
    * 128 or 64-bit trace ID lower-hex encoded into 32 or 16 characters (required)
@@ -84,10 +85,9 @@ static final class B3Injector<C, K> implements TraceContext.Injector<C> {
 
     @Override public void inject(TraceContext traceContext, C carrier) {
       setter.put(carrier, propagation.traceIdKey, traceContext.traceIdString());
-      setter.put(carrier, propagation.spanIdKey, HexCodec.toLowerHex(traceContext.spanId()));
+      setter.put(carrier, propagation.spanIdKey, toLowerHex(traceContext.spanId()));
       if (traceContext.parentId() != null) {
-        setter.put(carrier, propagation.parentSpanIdKey,
-            HexCodec.toLowerHex(traceContext.parentId()));
+        setter.put(carrier, propagation.parentSpanIdKey, toLowerHex(traceContext.parentId()));
       }
       if (traceContext.debug()) {
         setter.put(carrier, propagation.debugKey, "1");
@@ -114,31 +114,34 @@ static final class B3Extractor<C, K> implements TraceContext.Extractor<C> {
     @Override public TraceContextOrSamplingFlags extract(C carrier) {
       if (carrier == null) throw new NullPointerException("carrier == null");
 
-      String sampledString = getter.get(carrier, propagation.sampledKey);
+      String traceId = getter.get(carrier, propagation.traceIdKey);
+      String sampled = getter.get(carrier, propagation.sampledKey);
+      String debug = getter.get(carrier, propagation.debugKey);
+      if (traceId == null && sampled == null && debug == null) {
+        return TraceContextOrSamplingFlags.EMPTY;
+      }
+
       // Official sampled value is 1, though some old instrumentation send true
-      Boolean sampled = sampledString != null
-          ? sampledString.equals("1") || sampledString.equalsIgnoreCase("true")
+      Boolean sampledV = sampled != null
+          ? sampled.equals("1") || sampled.equalsIgnoreCase("true")
           : null;
-      boolean debug = "1".equals(getter.get(carrier, propagation.debugKey));
+      boolean debugV = "1".equals(debug);
 
-      String traceIdString = getter.get(carrier, propagation.traceIdKey);
-      String spanIdString = getter.get(carrier, propagation.spanIdKey);
-      if (traceIdString == null || spanIdString == null) { // return early if there's no trace ID
+      String spanId = getter.get(carrier, propagation.spanIdKey);
+      if (spanId == null) { // return early if there's no span ID
         return TraceContextOrSamplingFlags.create(
-            new SamplingFlags.Builder().sampled(sampled).debug(debug).build()
+            debugV ? SamplingFlags.DEBUG : SamplingFlags.Builder.build(sampledV)
         );
       }
 
-      TraceContext.Builder result = TraceContext.newBuilder().sampled(sampled).debug(debug);
+      TraceContext.Builder result = TraceContext.newBuilder().sampled(sampledV).debug(debugV);
       result.traceIdHigh(
-          traceIdString.length() == 32 ? lowerHexToUnsignedLong(traceIdString, 0) : 0
+          traceId.length() == 32 ? lowerHexToUnsignedLong(traceId, 0) : 0
       );
-      result.traceId(lowerHexToUnsignedLong(traceIdString));
-      result.spanId(lowerHexToUnsignedLong(spanIdString));
+      result.traceId(lowerHexToUnsignedLong(traceId));
+      result.spanId(lowerHexToUnsignedLong(spanId));
       String parentSpanIdString = getter.get(carrier, propagation.parentSpanIdKey);
-      if (parentSpanIdString != null) {
-        result.parentId(lowerHexToUnsignedLong(parentSpanIdString));
-      }
+      if (parentSpanIdString != null) result.parentId(lowerHexToUnsignedLong(parentSpanIdString));
       return TraceContextOrSamplingFlags.create(result.build());
     }
   }

diff --git a/brave/src/main/java/brave/propagation/Propagation.java b/brave/src/main/java/brave/propagation/Propagation.java
@@ -18,10 +18,10 @@
  * @param <K> Usually, but not always a String
  */
 public interface Propagation<K> {
-  Propagation<String> B3_STRING = Propagation.Factory.B3.create(Propagation.KeyFactory.STRING);
+  Propagation<String> B3_STRING = B3Propagation.FACTORY.create(Propagation.KeyFactory.STRING);
 
   abstract class Factory {
-    public static final Factory B3 = new B3Propagation.Factory();
+    public static final Factory B3 = B3Propagation.FACTORY;
 
     /**
      * Does the propagation implementation support sharing client and server span IDs. For example,

diff --git a/brave/src/main/java/brave/propagation/SamplingFlags.java b/brave/src/main/java/brave/propagation/SamplingFlags.java
@@ -39,14 +39,16 @@ public Builder debug(boolean debug) {
       return this;
     }
 
-    public SamplingFlags build() {
-      if (debug) {
-        return DEBUG;
-      } else if (sampled != null) {
-        return sampled ? SAMPLED : NOT_SAMPLED;
-      }
+    /** Allows you to create flags from a boolean value without allocating a builder instance */
+    public static SamplingFlags build(@Nullable Boolean sampled) {
+      if (sampled != null) return sampled ? SAMPLED : NOT_SAMPLED;
       return EMPTY;
     }
+
+    public SamplingFlags build() {
+      if (debug) return DEBUG;
+      return build(sampled);
+    }
   }
 
   static final class SamplingFlagsImpl extends SamplingFlags {

diff --git a/brave/src/main/java/brave/propagation/TraceContext.java b/brave/src/main/java/brave/propagation/TraceContext.java
@@ -166,9 +166,8 @@ public static abstract class Builder {
     public abstract TraceContext autoBuild();
 
     public final TraceContext build() {
-      if (extra().isEmpty()) return autoBuild();
       // make sure the extra data is immutable and unmodifiable
-      return extra(Collections.unmodifiableList(new ArrayList<>(extra()))).autoBuild();
+      return extra(ensureImmutable(extra())).autoBuild();
     }
 
     @Nullable abstract Boolean sampled();
@@ -181,4 +180,11 @@ public final TraceContext build() {
 
   TraceContext() { // no external implementations
   }
+
+  static List<Object> ensureImmutable(List<Object> extra) {
+    if (extra == Collections.EMPTY_LIST) return extra;
+    // Faster to make a copy than check the type to see if it is already a singleton list
+    if (extra.size() == 1) return Collections.singletonList(extra.get(0));
+    return Collections.unmodifiableList(new ArrayList<>(extra));
+  }
 }
diff --git a/brave/src/main/java/brave/propagation/TraceContextOrSamplingFlags.java b/brave/src/main/java/brave/propagation/TraceContextOrSamplingFlags.java
@@ -6,6 +6,8 @@
 import java.util.Collections;
 import java.util.List;
 
+import static brave.propagation.TraceContext.ensureImmutable;
+
 /**
  * Union type that contains only one of trace context, trace ID context or sampling flags. This type
  * is designed for use with {@link Tracer#nextSpan(TraceContextOrSamplingFlags)}.
@@ -30,6 +32,12 @@
  */
 //@Immutable
 public final class TraceContextOrSamplingFlags {
+  public static final TraceContextOrSamplingFlags EMPTY =
+      TraceContextOrSamplingFlags.create(SamplingFlags.EMPTY);
+
+  public static Builder newBuilder(){
+    return new Builder();
+  }
 
   /** Returns {@link SamplingFlags#sampled()}, regardless of subtype. */
   @Nullable public Boolean sampled() {
@@ -43,7 +51,7 @@ public TraceContextOrSamplingFlags sampled(@Nullable Boolean sampled) {
             .sampled(sampled).build(), extra);
       case 2:
         return new TraceContextOrSamplingFlags(type, ((TraceIdContext) value).toBuilder()
-            .sampled(sampled).build(), extra());
+            .sampled(sampled).build(), extra);
       case 3:
         return new TraceContextOrSamplingFlags(type, new SamplingFlags.Builder()
             .sampled(sampled).debug(value.debug()).build(), extra);
@@ -116,16 +124,53 @@ public static TraceContextOrSamplingFlags create(TraceContext.Builder builder) {
   final List<Object> extra;
 
   TraceContextOrSamplingFlags(int type, SamplingFlags value, List<Object> extra) {
+    if (value == null) throw new NullPointerException("value == null");
+    if (extra == null) throw new NullPointerException("extra == null");
     this.type = type;
     this.value = value;
-    this.extra = extra.isEmpty() ? extra : Collections.unmodifiableList(new ArrayList<>(extra));
+    this.extra = ensureImmutable(extra);
   }
 
   public static final class Builder {
     int type;
     SamplingFlags value;
     List<Object> extra = Collections.emptyList();
 
+    /** @see TraceContextOrSamplingFlags#context() */
+    public final Builder context(TraceContext context) {
+      if (context == null) throw new NullPointerException("context == null");
+      type = 1;
+      value = context;
+      return this;
+    }
+
+    /** @see TraceContextOrSamplingFlags#traceIdContext() */
+    public final Builder traceIdContext(TraceIdContext traceIdContext) {
+      if (traceIdContext == null) throw new NullPointerException("traceIdContext == null");
+      type = 2;
+      value = traceIdContext;
+      return this;
+    }
+
+    /** @see TraceContextOrSamplingFlags#samplingFlags() */
+    public final Builder samplingFlags(SamplingFlags samplingFlags) {
+      if (samplingFlags == null) throw new NullPointerException("samplingFlags == null");
+      type = 3;
+      value = samplingFlags;
+      return this;
+    }
+
+    /**
+     * Shares the input with the builder, replacing any current data in the builder.
+     *
+     * @see TraceContextOrSamplingFlags#extra()
+     */
+    public final Builder extra(List<Object> extra) {
+      if (extra == null) throw new NullPointerException("extra == null");
+      this.extra = extra; // sharing a copy in case it is immutable
+      return this;
+    }
+
     /** @see TraceContextOrSamplingFlags#extra() */
     public final Builder addExtra(Object extra) {
       if (extra == null) throw new NullPointerException("extra == null");
@@ -136,6 +181,7 @@ public final Builder addExtra(Object extra) {
       return this;
     }
 
+    /** Returns an immutable result from the values currently in the builder */
     public final TraceContextOrSamplingFlags build() {
       if (!extra.isEmpty() && type == 1) { // move extra to the trace context
         TraceContext context = (TraceContext) value;

diff --git a/brave/src/test/java/brave/TracerTest.java b/brave/src/test/java/brave/TracerTest.java
@@ -146,7 +146,7 @@ public class TracerTest {
     tracer = Tracing.newBuilder()
         .propagationFactory(new Propagation.Factory(){
           @Override public <K> Propagation<K> create(Propagation.KeyFactory<K> keyFactory) {
-            return B3Propagation.create(keyFactory);
+            return B3Propagation.FACTORY.create(keyFactory);
           }
         })
         .spanReporter(spans::add).build().tracer();

diff --git a/brave/src/test/java/brave/propagation/TraceContextTest.java b/brave/src/test/java/brave/propagation/TraceContextTest.java
@@ -1,5 +1,8 @@
 package brave.propagation;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import org.junit.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -36,4 +39,23 @@ public void testToString() {
     assertThat(context.toString())
         .isEqualTo("000000000000014d00000000000001bc/0000000000000003");
   }
+
+  @Test(expected = UnsupportedOperationException.class)
+  public void ensureImmutable_returnsImmutableEmptyList() {
+    TraceContext.ensureImmutable(new ArrayList<>()).add("foo");
+  }
+
+  @Test public void ensureImmutable_convertsToSingletonList() {
+    List<Object> list = new ArrayList<>();
+    list.add("foo");
+    TraceContext.ensureImmutable(list);
+    assertThat(TraceContext.ensureImmutable(list).getClass().getSimpleName())
+        .isEqualTo("SingletonList");
+  }
+
+  @Test public void ensureImmutable_returnsEmptyList() {
+    List<Object> list = Collections.emptyList();
+    assertThat(TraceContext.ensureImmutable(list))
+        .isSameAs(list);
+  }
 }
diff --git a/instrumentation/benchmarks/src/main/java/brave/EndToEndBenchmarks.java b/instrumentation/benchmarks/src/main/java/brave/EndToEndBenchmarks.java
@@ -82,7 +82,7 @@ public Traced128() {
   public static class TracedAWS extends ForwardingTracingFilter {
     public TracedAWS() {
       super(Tracing.newBuilder()
-          .propagationFactory(new AWSPropagation.Factory())
+          .propagationFactory(AWSPropagation.FACTORY)
           .spanReporter(Reporter.NOOP)
           .build());
     }

diff --git a/instrumentation/benchmarks/src/main/java/brave/internal/PropagationBenchmarks.java b/instrumentation/benchmarks/src/main/java/brave/internal/PropagationBenchmarks.java
@@ -19,6 +19,7 @@
 import brave.propagation.TraceContext.Injector;
 import brave.propagation.TraceContextOrSamplingFlags;
 import brave.propagation.aws.AWSPropagation;
+import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -47,11 +48,11 @@ public class PropagationBenchmarks {
   static final Injector<Map<String, String>> b3Injector =
       Propagation.B3_STRING.injector(Map::put);
   static final Injector<Map<String, String>> awsInjector =
-      new AWSPropagation.Factory().create(Propagation.KeyFactory.STRING).injector(Map::put);
+      AWSPropagation.FACTORY.create(Propagation.KeyFactory.STRING).injector(Map::put);
   static final Extractor<Map<String, String>> b3Extractor =
       Propagation.B3_STRING.extractor(Map::get);
   static final Extractor<Map<String, String>> awsExtractor =
-      new AWSPropagation.Factory().create(Propagation.KeyFactory.STRING).extractor(Map::get);
+      AWSPropagation.FACTORY.create(Propagation.KeyFactory.STRING).extractor(Map::get);
 
   static final TraceContext context = TraceContext.newBuilder()
       .traceIdHigh(HexCodec.lowerHexToUnsignedLong("67891233abcdef01"))
@@ -67,6 +68,8 @@ public class PropagationBenchmarks {
     }
   };
 
+  static final Map<String, String> nothingIncoming = Collections.emptyMap();
+
   Map<String, String> carrier = new LinkedHashMap<>();
 
   @Benchmark public void inject_b3() {
@@ -85,10 +88,19 @@ public class PropagationBenchmarks {
     return awsExtractor.extract(incoming);
   }
 
+  @Benchmark public TraceContextOrSamplingFlags extract_b3_nothing() {
+    return b3Extractor.extract(nothingIncoming);
+  }
+
+  @Benchmark public TraceContextOrSamplingFlags extract_aws_nothing() {
+    return awsExtractor.extract(nothingIncoming);
+  }
+
   // Convenience main entry-point
   public static void main(String[] args) throws RunnerException {
+    new PropagationBenchmarks().extract_aws();
     Options opt = new OptionsBuilder()
-        .include(".*" + PropagationBenchmarks.class.getSimpleName() + ".extract_aws")
+        .include(".*" + PropagationBenchmarks.class.getSimpleName())
         .build();
 
     new Runner(opt).run();

diff --git a/propagation/aws/README.md b/propagation/aws/README.md
@@ -2,11 +2,11 @@
 This changes brave to use "x-amzn-trace-id" as opposed to "x-b3" prefixed headers to propagate trace
 context across processes.
 
-To enable this, configure `brave.Tracing` with `AWSPropagation.Factory` like so:
+To enable this, configure `brave.Tracing` with `AWSPropagation.FACTORY` like so:
 
 ```java
 tracing = Tracing.newBuilder()
-    .propagationFactory(new AWSPropagation.Factory())
+    .propagationFactory(AWSPropagation.FACTORY)
     ...
     .build();
 ```
@@ -23,10 +23,10 @@ tracing = Tracing.newBuilder()
 ## Utilities
 There are a couple added utilities for parsing and generating an AWS trace ID string:
 
-* `AWSPropagation.traceIdString` - used to generate a formatted trace ID for correlation purposes.
+* `AWSPropagation.rootField` - used to generate a formatted root field ID for correlation purposes.
 * `AWSPropagation.extract` - extracts a trace context from a string such as an environment variable.
 
-Ex to extract the trace ID from the built-in AWS Lambda variable
+Ex. to extract a trace context from the built-in AWS Lambda variable
 ```java
 extracted = AWSPropagation.extract(System.getenv("_X_AMZN_TRACE_ID"));
 ```