Add baggage to B3 codec

jaeger-core/README.md

@@ -40,7 +40,8 @@ The `config` objects lazily builds and configures Jaeger Tracer. Multiple calls
 Jaeger tracer can also work in the environment where B3 propagation is used. This is mostly related 
 to systems instrumented with Zipkin. Once you register `B3TextMapCodec`, Jaeger can join traces 
 started by other Zipkin instrumented applications. This includes reading headers 
-like "X-B3-TraceId".
+like "X-B3-TraceId". Jaeger B3 implementation automatically propagates baggage and by default it
+uses `baggage-` prefix.
 
 Example configuration:
 ```java

jaeger-core/src/main/java/io/jaegertracing/propagation/B3TextMapCodec.java

@@ -15,7 +15,10 @@
 package io.jaegertracing.propagation;
 
 import io.jaegertracing.SpanContext;
+import io.jaegertracing.Tracer;
+import io.jaegertracing.baggage.BaggageRestrictionManager;
 import io.opentracing.propagation.TextMap;
+import java.util.HashMap;
 import java.util.Map;
 
 /**
@@ -34,17 +37,34 @@
  *
  * <p>
  * See <a href="http://zipkin.io/pages/instrumenting.html">Instrumenting a Library</a>
+ *
+ * Note that this codec automatically propagates baggage
+ * (with {@value io.jaegertracing.propagation.B3TextMapCodec#BAGGAGE_NAME} prefix).
+ * Baggage whitelisting can be configured in {@link BaggageRestrictionManager} and then
+ * passed to {@link Tracer.Builder#baggageRestrictionManager}
  */
 public class B3TextMapCodec implements Codec<TextMap> {
   protected static final String TRACE_ID_NAME = "X-B3-TraceId";
   protected static final String SPAN_ID_NAME = "X-B3-SpanId";
   protected static final String PARENT_SPAN_ID_NAME = "X-B3-ParentSpanId";
   protected static final String SAMPLED_NAME = "X-B3-Sampled";
   protected static final String FLAGS_NAME = "X-B3-Flags";
+  protected static final String BAGGAGE_NAME = "baggage-";
   // NOTE: uber's flags aren't the same as B3/Finagle ones
   protected static final byte SAMPLED_FLAG = 1;
   protected static final byte DEBUG_FLAG = 2;
 
+  private static final PrefixedKeys keys = new PrefixedKeys();
+  private final String baggageName;
+
+  public B3TextMapCodec() {
+    this(new Builder());
+  }
+
+  private B3TextMapCodec(Builder builder) {
+    this.baggageName = builder.baggagePrefix;
+  }
+
   @Override
   public void inject(SpanContext spanContext, TextMap carrier) {
     carrier.put(TRACE_ID_NAME, HexCodec.toLowerHex(spanContext.getTraceId()));
@@ -56,6 +76,9 @@ public void inject(SpanContext spanContext, TextMap carrier) {
     if (spanContext.isDebug()) {
       carrier.put(FLAGS_NAME, "1");
     }
+    for (Map.Entry<String, String> entry : spanContext.baggageItems()) {
+      carrier.put(keys.prefixedKey(entry.getKey(), baggageName), entry.getValue());
+    }
   }
 
   @Override
@@ -64,6 +87,7 @@ public SpanContext extract(TextMap carrier) {
     Long spanId = null;
     Long parentId = 0L; // Conventionally, parent id == 0 means the root span
     byte flags = 0;
+    Map<String, String> baggage = null;
     for (Map.Entry<String, String> entry : carrier) {
       if (entry.getKey().equalsIgnoreCase(SAMPLED_NAME)) {
         String value = entry.getValue();
@@ -80,12 +104,37 @@ public SpanContext extract(TextMap carrier) {
         if (entry.getValue().equals("1")) {
           flags |= DEBUG_FLAG;
         }
+      } else if (entry.getKey().startsWith(BAGGAGE_NAME)) {
+        if (baggage == null) {
+          baggage = new HashMap<String, String>();
+        }
+        baggage.put(keys.unprefixedKey(entry.getKey(), BAGGAGE_NAME), entry.getValue());
       }
     }
 
     if (null != traceId && null != parentId && null != spanId) {
-      return new SpanContext(traceId, spanId, parentId, flags);
+      SpanContext spanContext = new SpanContext(traceId, spanId, parentId, flags);
+      if (baggage != null) {
+        spanContext = spanContext.withBaggage(baggage);
+      }
+      return spanContext;
     }
     return null;
   }
+
+  public static class Builder {
+    private String baggagePrefix = BAGGAGE_NAME;
+
+    /**
+     * Specify baggage prefix. The default is {@value B3TextMapCodec#BAGGAGE_NAME}
+     */
+    public Builder withBaggagePrefix(String baggagePrefix) {
+      this.baggagePrefix = baggagePrefix;
+      return this;
+    }
+
+    public B3TextMapCodec build() {
+      return new B3TextMapCodec(this);
+    }
+  }
 }

jaeger-core/src/test/java/io/jaegertracing/propagation/B3TextMapCodecTest.java

@@ -15,14 +15,19 @@
 package io.jaegertracing.propagation;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 import io.jaegertracing.SpanContext;
 import io.opentracing.propagation.TextMap;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
 import org.junit.Test;
 
 /**
@@ -48,6 +53,8 @@ public void downgrades128BitTraceIdToLower64Bits() throws Exception {
     textMap.put(B3TextMapCodec.PARENT_SPAN_ID_NAME, "0");
     textMap.put(B3TextMapCodec.SAMPLED_NAME, "1");
     textMap.put(B3TextMapCodec.FLAGS_NAME, "1");
+    textMap.put(B3TextMapCodec.BAGGAGE_NAME + "foo", "bar");
+    textMap.put("random-foo", "bar");
 
     SpanContext context = b3Codec.extract(textMap);
 
@@ -56,6 +63,41 @@ public void downgrades128BitTraceIdToLower64Bits() throws Exception {
     assertEquals(HexCodec.lowerHexToUnsignedLong(lower64Bits).longValue(), context.getSpanId());
     assertEquals(0, context.getParentId());
     assertEquals(B3TextMapCodec.SAMPLED_FLAG | B3TextMapCodec.DEBUG_FLAG, context.getFlags());
+    assertEquals(1, ((Set)context.baggageItems()).size());
+    assertEquals("bar", context.getBaggageItem("foo"));
+  }
+
+  @Test
+  public void testDefault() {
+    B3TextMapCodec b3Codec = new B3TextMapCodec.Builder()
+        .build();
+
+    DelegatingTextMap entries = new DelegatingTextMap();
+    SpanContext spanContext = new SpanContext(1, 2, 3, (byte)0)
+        .withBaggageItem("foo", "bar");
+
+    b3Codec.inject(spanContext, entries);
+    assertEquals(5, entries.delegate.size());
+    assertNotNull(entries.delegate.get(B3TextMapCodec.TRACE_ID_NAME));
+    assertNotNull(entries.delegate.get(B3TextMapCodec.SPAN_ID_NAME));
+    assertNotNull(entries.delegate.get(B3TextMapCodec.PARENT_SPAN_ID_NAME));
+    assertNotNull(entries.delegate.get(B3TextMapCodec.SAMPLED_NAME));
+    assertEquals("bar", entries.delegate.get("baggage-foo"));
+  }
+
+  @Test
+  public void testChangeBaggagePrefix() {
+    B3TextMapCodec b3Codec = new B3TextMapCodec.Builder()
+        .withBaggagePrefix("foo")
+        .build();
+
+    DelegatingTextMap entries = new DelegatingTextMap();
+    SpanContext spanContext = new SpanContext(1, 2, 3, (byte)0)
+        .withBaggageItem("foo", "bar");
+
+    b3Codec.inject(spanContext, entries);
+    assertEquals(5, entries.delegate.size());
+    assertEquals("bar", entries.delegate.get("foofoo"));
   }
 
   @Test