Set incoming baggage from parent on Span (#8364)

Signed-off-by: Tim Quinn <[email protected]>

tracing/opentelemetry/src/main/java/io/helidon/tracing/opentelemetry/OpenTelemetrySpanBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, 2023 Oracle and/or its affiliates.
+ * Copyright (c) 2022, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,13 +20,15 @@
 import io.helidon.tracing.Span;
 import io.helidon.tracing.SpanContext;
 
+import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.api.trace.SpanBuilder;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 
 class OpenTelemetrySpanBuilder implements Span.Builder<OpenTelemetrySpanBuilder> {
     private final SpanBuilder spanBuilder;
     private boolean parentSet;
+    private Baggage parentBaggage;
 
     OpenTelemetrySpanBuilder(SpanBuilder spanBuilder) {
         this.spanBuilder = spanBuilder;
@@ -41,6 +43,9 @@ public Span build() {
     public OpenTelemetrySpanBuilder parent(SpanContext spanContext) {
         this.parentSet = true;
         spanContext.asParent(this);
+        if (spanContext instanceof OpenTelemetrySpanContext otsc) {
+            parentBaggage = Baggage.fromContext(otsc.openTelemetry());
+        }
         return this;
     }
 
@@ -87,7 +92,11 @@ public Span start(Instant instant) {
         }
         spanBuilder.setStartTimestamp(instant);
         io.opentelemetry.api.trace.Span span = spanBuilder.startSpan();
-        return new OpenTelemetrySpan(span);
+        Span result = new OpenTelemetrySpan(span);
+        if (parentBaggage != null) {
+            parentBaggage.forEach((key, baggageEntry) -> result.baggage(key, baggageEntry.getValue()));
+        }
+        return result;
     }
 
     // used to set open telemetry context as parent, to be equivalent in function to

tracing/opentelemetry/src/test/java/io/helidon/tracing/opentelemetry/MapHeaderProvider.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 2024 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.helidon.tracing.opentelemetry;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import io.helidon.tracing.HeaderProvider;
+
+record MapHeaderProvider(Map<String, List<String>> values) implements HeaderProvider {
+
+    @Override
+    public Iterable<String> keys() {
+        return values.keySet();
+    }
+
+    @Override
+    public Optional<String> get(String key) {
+        if (!values.containsKey(key)) {
+            return Optional.empty();
+        }
+        List<String> matches = values.get(key);
+        return matches.isEmpty() ? Optional.empty() : Optional.of(matches.get(0));
+    }
+
+    @Override
+    public Iterable<String> getAll(String key) {
+        return values.get(key);
+    }
+
+    @Override
+    public boolean contains(String key) {
+        return values.containsKey(key);
+    }
+}

tracing/opentelemetry/src/test/java/io/helidon/tracing/opentelemetry/TestSpanAndBaggage.java

@@ -15,11 +15,15 @@
  */
 package io.helidon.tracing.opentelemetry;
 
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 import io.helidon.common.testing.junit5.OptionalMatcher;
+import io.helidon.tracing.HeaderProvider;
 import io.helidon.tracing.Scope;
 import io.helidon.tracing.Span;
+import io.helidon.tracing.SpanContext;
 import io.helidon.tracing.Tracer;
 
 import org.junit.jupiter.api.AfterAll;
@@ -109,4 +113,17 @@ void testActiveSpanScopeWithBaggage() {
                    currentSpanAfterTryResourcesBlock.get().context().spanId(),
                    containsString("00000000"));
     }
+
+
+    @Test
+    void testIncomingBaggage() {
+        Tracer tracer = Tracer.global();
+        HeaderProvider inboundHeaders = new MapHeaderProvider(Map.of("baggage", List.of("bag1=val1,bag2=val2")));
+        Optional<SpanContext> spanContextOpt = tracer.extract(inboundHeaders);
+        assertThat("Span context from inbound headers", spanContextOpt, OptionalMatcher.optionalPresent());
+        Span span = tracer.spanBuilder("inbound").parent(spanContextOpt.get()).start();
+        span.end();
+        assertThat("Inbound baggage bag1", span.baggage("bag1"), OptionalMatcher.optionalValue(is("val1")));
+        assertThat("Inbound baggage bag1", span.baggage("bag2"), OptionalMatcher.optionalValue(is("val2")));
+    }
 }

tracing/opentracing/src/main/java/io/helidon/tracing/opentracing/OpenTracingSpanBuilder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022 Oracle and/or its affiliates.
+ * Copyright (c) 2022, 2024 Oracle and/or its affiliates.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 package io.helidon.tracing.opentracing;
 
 import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import io.helidon.tracing.Span;
@@ -27,6 +29,7 @@
 class OpenTracingSpanBuilder implements Span.Builder<OpenTracingSpanBuilder> {
     private final Tracer.SpanBuilder delegate;
     private final Tracer tracer;
+    private Map<String, String> baggage;
 
     OpenTracingSpanBuilder(Tracer tracer, Tracer.SpanBuilder delegate) {
         this.tracer = tracer;
@@ -42,6 +45,12 @@ public Span build() {
     public OpenTracingSpanBuilder parent(SpanContext spanContext) {
         if (spanContext instanceof OpenTracingContext otc) {
             delegate.asChildOf(otc.openTracing());
+            if (baggage == null) {
+                baggage = new HashMap<>();
+            } else {
+                baggage.clear();
+            }
+            otc.openTracing().baggageItems().forEach(entry -> baggage.put(entry.getKey(), entry.getValue()));
         }
         return this;
     }
@@ -83,6 +92,10 @@ public OpenTracingSpanBuilder tag(String key, Number value) {
     @Override
     public Span start(Instant instant) {
         long micro = TimeUnit.MILLISECONDS.toMicros(instant.toEpochMilli());
-        return new OpenTracingSpan(tracer, delegate.withStartTimestamp(micro).start());
+        Span result = new OpenTracingSpan(tracer, delegate.withStartTimestamp(micro).start());
+        if (baggage != null) {
+            baggage.forEach(result::baggage);
+        }
+        return result;
     }
 }