Support jaeger-baggage header for ad-hoc baggage (#525)

* Support `jaeger-baggage` header for ad-hoc baggage Signed-off-by: Yuri Shkuro <[email protected]> * prettify Signed-off-by: Yuri Shkuro <[email protected]>
jaegertracing · Aug 28, 2018 · 73c8e0e · 73c8e0e
1 parent ce175a4
commit 73c8e0e
Show file tree

Hide file tree

Showing 7 changed files with 112 additions and 20 deletions.
diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/Constants.java b/jaeger-core/src/main/java/io/jaegertracing/internal/Constants.java
@@ -35,6 +35,15 @@ public class Constants {
    */
   public static final String DEBUG_ID_HEADER_KEY = "jaeger-debug-id";
 
+  /**
+   * The name of HTTP header or a TextMap carrier key that can be used to pass
+   * additional baggage to the span, e.g. when executing an ad-hoc curl request:
+   * <pre>
+   * curl -H 'jaeger-baggage: k1=v1,k2=v2' http://...
+   * </pre>
+   */
+  public static final String BAGGAGE_HEADER_KEY = "jaeger-baggage";
+
   /**
    * The name of the tag used to report client version.
    */

diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/JaegerSpanContext.java b/jaeger-core/src/main/java/io/jaegertracing/internal/JaegerSpanContext.java
@@ -53,7 +53,7 @@ protected JaegerSpanContext(
       String debugId,
       JaegerObjectFactory objectFactory) {
     if (baggage == null) {
-      throw new NullPointerException();
+      baggage = Collections.<String, String>emptyMap();
     }
     this.traceId = traceId;
     this.spanId = spanId;

diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/JaegerTracer.java b/jaeger-core/src/main/java/io/jaegertracing/internal/JaegerTracer.java
@@ -410,9 +410,7 @@ public JaegerSpan start() {
       }
 
       String debugId = debugId();
-      if (references.isEmpty()) {
-        context = createNewContext(null);
-      } else if (debugId != null) {
+      if (references.isEmpty() || debugId != null) {
         context = createNewContext(debugId);
       } else {
         context = createChildContext();

diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/propagation/TextMapCodec.java b/jaeger-core/src/main/java/io/jaegertracing/internal/propagation/TextMapCodec.java
@@ -22,15 +22,18 @@
 import io.jaegertracing.internal.exceptions.MalformedTracerStateStringException;
 import io.jaegertracing.spi.Codec;
 import io.opentracing.propagation.TextMap;
+
 import java.io.UnsupportedEncodingException;
 import java.math.BigInteger;
 import java.net.URLDecoder;
 import java.net.URLEncoder;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.Locale;
 import java.util.Map;
 
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
 public class TextMapCodec implements Codec<TextMap> {
   /**
    * Key used to store serialized span context representation
@@ -125,24 +128,35 @@ public JaegerSpanContext extract(TextMap carrier) {
           baggage = new HashMap<String, String>();
         }
         baggage.put(keys.unprefixedKey(key, baggagePrefix), decodedValue(entry.getValue()));
+      } else if (key.equals(Constants.BAGGAGE_HEADER_KEY)) {
+        baggage = parseBaggageHeader(decodedValue(entry.getValue()), baggage);
       }
     }
-    if (context == null) {
-      if (debugId != null) {
-        return objectFactory.createSpanContext(
-            0,
-            0,
-            0,
-            (byte) 0,
-            Collections.<String, String>emptyMap(),
-            debugId);
-      }
-      return null;
-    }
-    if (baggage == null) {
+    if (debugId == null && baggage == null) {
       return context;
     }
-    return context.withBaggage(baggage);
+    return objectFactory.createSpanContext(
+      context == null ? 0 : context.getTraceId(),
+      context == null ? 0 : context.getSpanId(),
+      context == null ? 0 : context.getParentId(),
+      context == null ? (byte)0 : context.getFlags(),
+      baggage,
+      debugId);
+  }
+
+  private Map<String, String> parseBaggageHeader(String header, Map<String, String> baggage) {
+    for (String part : header.split("\\s*,\\s*")) {
+      String[] kv = part.split("\\s*=\\s*");
+      if (kv.length == 2) {
+        if (baggage == null) {
+          baggage = new HashMap<String, String>();
+        }
+        baggage.put(kv[0], kv[1]);
+      } else {
+        log.debug("malformed token in {} header: {}", Constants.BAGGAGE_HEADER_KEY, part);
+      }
+    }
+    return baggage;
   }
 
   @Override

diff --git a/jaeger-core/src/test/java/io/jaegertracing/internal/JaegerTracerTest.java b/jaeger-core/src/test/java/io/jaegertracing/internal/JaegerTracerTest.java
@@ -37,8 +37,13 @@
 import io.opentracing.Span;
 import io.opentracing.propagation.Format;
 import io.opentracing.propagation.TextMap;
+import io.opentracing.propagation.TextMapExtractAdapter;
+import io.opentracing.propagation.TextMapInjectAdapter;
 import io.opentracing.tag.Tags;
 import java.io.Closeable;
+import java.util.HashMap;
+import java.util.Map;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mockito;
@@ -193,4 +198,29 @@ public void testCustomSpanOnSpanManager() {
     // check
     assertEquals(activeSpan, tracer.activeSpan());
   }
+
+  @Test
+  public void testStartTraceWithAdhocBaggage() {
+    traceWithAdhocBaggage(new HashMap<String, String>());
+  }
+
+  @Test
+  public void testJoinTraceWithAdhocBaggage() {
+    Span span = tracer.buildSpan("test").start();
+    Map<String, String> headers = new HashMap<String, String>();
+    tracer.inject(span.context(), Format.Builtin.HTTP_HEADERS, new TextMapInjectAdapter(headers));
+    assertEquals(1, headers.size());
+
+    traceWithAdhocBaggage(headers);
+  }
+
+  private void traceWithAdhocBaggage(Map<String, String> headers) {
+    headers.put("jaeger-baggage", "k1=v1, k2 = v2");
+
+    JaegerSpanContext parent = tracer.extract(Format.Builtin.HTTP_HEADERS, new TextMapExtractAdapter(headers));
+    Span span = tracer.buildSpan("test").asChildOf(parent).start();
+
+    assertEquals("must have baggage", "v1", span.getBaggageItem("k1"));
+    assertEquals("must have baggage", "v2", span.getBaggageItem("k2"));
+  }
 }
diff --git a/jaeger-core/src/test/java/io/jaegertracing/internal/PropagationTest.java b/jaeger-core/src/test/java/io/jaegertracing/internal/PropagationTest.java
@@ -16,6 +16,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 
@@ -33,6 +34,8 @@
 import io.opentracing.propagation.TextMap;
 import io.opentracing.propagation.TextMapExtractAdapter;
 import java.util.Collections;
+import java.util.Map;
+
 import org.junit.Test;
 
 public class PropagationTest {
@@ -42,9 +45,11 @@ public void testDebugCorrelationId() {
             .withReporter(new InMemoryReporter())
             .withSampler(new ConstSampler(true))
             .build();
-    TextMap carrier = new TextMapExtractAdapter(Collections.singletonMap(Constants.DEBUG_ID_HEADER_KEY, "Coraline"));
+    Map<String, String> headers = Collections.singletonMap(Constants.DEBUG_ID_HEADER_KEY, "Coraline");
+    TextMap carrier = new TextMapExtractAdapter(headers);
 
     JaegerSpanContext jaegerSpanContext = tracer.extract(Format.Builtin.TEXT_MAP, carrier);
+    assertNotNull(jaegerSpanContext);
     assertTrue(jaegerSpanContext.isDebugIdContainerOnly());
     assertEquals("Coraline", jaegerSpanContext.getDebugId());
 

diff --git a/jaeger-core/src/test/java/io/jaegertracing/internal/propagation/TextMapCodecTest.java b/jaeger-core/src/test/java/io/jaegertracing/internal/propagation/TextMapCodecTest.java
@@ -21,6 +21,12 @@
 import io.jaegertracing.internal.JaegerSpanContext;
 import io.jaegertracing.internal.exceptions.EmptyTracerStateStringException;
 import io.jaegertracing.internal.exceptions.MalformedTracerStateStringException;
+import io.opentracing.propagation.TextMapExtractAdapter;
+import io.opentracing.propagation.TextMapInjectAdapter;
+
+import java.util.HashMap;
+import java.util.Map;
+
 import org.junit.Test;
 
 public class TextMapCodecTest {
@@ -83,4 +89,34 @@ public void testContextAsStringFormatsPositiveFields() {
     assertEquals(parentId, contextFromStr.getParentId());
     assertEquals(flags, contextFromStr.getFlags());
   }
+
+  /**
+   * Tests that the codec will include baggage from header "jaeger-baggage".
+   */
+  @Test
+  public void testAdhocBaggageWithTraceId() {
+    TextMapCodec codec = new TextMapCodec(false);
+    Map<String, String> headers = new HashMap<String, String>();
+    codec.inject(new JaegerSpanContext(42, 1, 0, (byte)1), new TextMapInjectAdapter(headers));
+    headers.put("jaeger-baggage", "k1=v1, k2 = v2");
+    JaegerSpanContext context = codec.extract(new TextMapExtractAdapter(headers));
+    assertEquals("must have trace ID", 42, context.getTraceId());
+    assertEquals("must have bagggae", "v1", context.getBaggageItem("k1"));
+    assertEquals("must have bagggae", "v2", context.getBaggageItem("k2"));
+  }
+
+  /**
+   * Tests that the codec will return non-null SpanContext even if the only header
+   * present is "jaeger-baggage".
+   */
+  @Test
+  public void testAdhocBaggageWithoutTraceId() {
+    Map<String, String> headers = new HashMap<String, String>();
+    headers.put("jaeger-baggage", "k1=v1, k2 = v2, k3=v3=d3");
+    TextMapCodec codec = new TextMapCodec(false);
+    JaegerSpanContext context = codec.extract(new TextMapExtractAdapter(headers));
+    assertEquals("v1", context.getBaggageItem("k1"));
+    assertEquals("v2", context.getBaggageItem("k2"));
+    assertEquals(null, context.getBaggageItem("k3"));
+  }
 }