OtTracerPropagator: inject and extract baggage (open-telemetry#2594)

Sergei Malafeev · web-flow · commit 24cfb4acbfb0 · 2021-02-01T14:07:12.000+01:00
diff --git a/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/OtTracerPropagator.java b/extensions/trace-propagators/src/main/java/io/opentelemetry/extension/trace/propagation/OtTracerPropagator.java
@@ -7,6 +7,8 @@
 
 import static io.opentelemetry.extension.trace.propagation.Common.MAX_TRACE_ID_LENGTH;
 
+import io.opentelemetry.api.baggage.Baggage;
+import io.opentelemetry.api.baggage.BaggageBuilder;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceId;
@@ -31,6 +33,7 @@ public final class OtTracerPropagator implements TextMapPropagator {
   static final String TRACE_ID_HEADER = "ot-tracer-traceid";
   static final String SPAN_ID_HEADER = "ot-tracer-spanid";
   static final String SAMPLED_HEADER = "ot-tracer-sampled";
+  static final String PREFIX_BAGGAGE_HEADER = "ot-baggage-";
   private static final Collection<String> FIELDS =
       Collections.unmodifiableList(Arrays.asList(TRACE_ID_HEADER, SPAN_ID_HEADER, SAMPLED_HEADER));
 
@@ -66,6 +69,16 @@ public <C> void inject(Context context, C carrier, Setter<C> setter) {
         spanContext.getTraceIdAsHexString().substring(TraceId.getHexLength() / 2));
     setter.set(carrier, SPAN_ID_HEADER, spanContext.getSpanIdAsHexString());
     setter.set(carrier, SAMPLED_HEADER, String.valueOf(spanContext.isSampled()));
+
+    // Baggage is only injected if there is a valid SpanContext
+    Baggage baggage = Baggage.fromContext(context);
+    if (!baggage.isEmpty()) {
+      baggage.forEach(
+          (key, value, metadata) -> {
+            // Metadata is not supported by OpenTracing
+            setter.set(carrier, PREFIX_BAGGAGE_HEADER + key, value);
+          });
+    }
   }
 
   @Override
@@ -85,7 +98,29 @@ public <C> Context extract(Context context, @Nullable C carrier, Getter<C> gette
     if (!spanContext.isValid()) {
       return context;
     }
-    return context.with(Span.wrap(spanContext));
+
+    Context extractedContext = context.with(Span.wrap(spanContext));
+
+    // Baggage is only extracted if there is a valid SpanContext
+    if (carrier != null) {
+      BaggageBuilder baggageBuilder = Baggage.builder();
+      for (String key : getter.keys(carrier)) {
+        if (!key.startsWith(PREFIX_BAGGAGE_HEADER)) {
+          continue;
+        }
+        String value = getter.get(carrier, key);
+        if (value == null) {
+          continue;
+        }
+        baggageBuilder.put(key.replace(OtTracerPropagator.PREFIX_BAGGAGE_HEADER, ""), value);
+      }
+      Baggage baggage = baggageBuilder.build();
+      if (!baggage.isEmpty()) {
+        extractedContext = extractedContext.with(baggage);
+      }
+    }
+
+    return extractedContext;
   }
 
   static SpanContext buildSpanContext(String traceId, String spanId, String sampled) {
diff --git a/extensions/trace-propagators/src/test/java/io/opentelemetry/extension/trace/propagation/OtTracerPropagatorTest.java b/extensions/trace-propagators/src/test/java/io/opentelemetry/extension/trace/propagation/OtTracerPropagatorTest.java
@@ -7,6 +7,7 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.SpanId;
@@ -114,6 +115,37 @@ void inject_NotSampledContext() {
     assertThat(carrier).containsEntry(OtTracerPropagator.SAMPLED_HEADER, "false");
   }
 
+  @Test
+  void inject_Baggage() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    Baggage baggage = Baggage.builder().put("foo", "bar").put("key", "value").build();
+    propagator.inject(
+        withSpanContext(
+            SpanContext.create(TRACE_ID, SPAN_ID, SAMPLED_TRACE_OPTIONS, TRACE_STATE_DEFAULT),
+            Context.current().with(baggage)),
+        carrier,
+        setter);
+    assertThat(carrier).containsEntry(OtTracerPropagator.PREFIX_BAGGAGE_HEADER + "foo", "bar");
+    assertThat(carrier).containsEntry(OtTracerPropagator.PREFIX_BAGGAGE_HEADER + "key", "value");
+  }
+
+  @Test
+  void inject_Baggage_InvalidContext() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    Baggage baggage = Baggage.builder().put("foo", "bar").put("key", "value").build();
+    propagator.inject(
+        withSpanContext(
+            SpanContext.create(
+                TraceId.getInvalid(),
+                SpanId.getInvalid(),
+                SAMPLED_TRACE_OPTIONS,
+                TRACE_STATE_DEFAULT),
+            Context.current().with(baggage)),
+        carrier,
+        setter);
+    assertThat(carrier).isEmpty();
+  }
+
   @Test
   void extract_Nothing() {
     // Context remains untouched.
@@ -246,4 +278,33 @@ void extract_emptyCarrier() {
     assertThat(getSpanContext(propagator.extract(Context.current(), emptyHeaders, getter)))
         .isEqualTo(SpanContext.getInvalid());
   }
+
+  @Test
+  void extract_Baggage() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    carrier.put(OtTracerPropagator.TRACE_ID_HEADER, TRACE_ID);
+    carrier.put(OtTracerPropagator.SPAN_ID_HEADER, SPAN_ID);
+    carrier.put(OtTracerPropagator.SAMPLED_HEADER, Common.TRUE_INT);
+    carrier.put(OtTracerPropagator.PREFIX_BAGGAGE_HEADER + "foo", "bar");
+    carrier.put(OtTracerPropagator.PREFIX_BAGGAGE_HEADER + "key", "value");
+
+    Context context = propagator.extract(Context.current(), carrier, getter);
+
+    Baggage expectedBaggage = Baggage.builder().put("foo", "bar").put("key", "value").build();
+    assertThat(Baggage.fromContext(context)).isEqualTo(expectedBaggage);
+  }
+
+  @Test
+  void extract_Baggage_InvalidContext() {
+    Map<String, String> carrier = new LinkedHashMap<>();
+    carrier.put(OtTracerPropagator.TRACE_ID_HEADER, TraceId.getInvalid());
+    carrier.put(OtTracerPropagator.SPAN_ID_HEADER, SPAN_ID);
+    carrier.put(OtTracerPropagator.SAMPLED_HEADER, Common.TRUE_INT);
+    carrier.put(OtTracerPropagator.PREFIX_BAGGAGE_HEADER + "foo", "bar");
+    carrier.put(OtTracerPropagator.PREFIX_BAGGAGE_HEADER + "key", "value");
+
+    Context context = propagator.extract(Context.current(), carrier, getter);
+
+    assertThat(Baggage.fromContext(context).isEmpty()).isTrue();
+  }
 }
