Ensures that we always set baggage in scope when setting a value

fixes gh-181

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageInScope.java

@@ -50,12 +50,14 @@ public String get(TraceContext traceContext) {
     }
 
     @Override
+    @Deprecated
     public Baggage set(String value) {
         this.delegate.updateValue(value);
         return this;
     }
 
     @Override
+    @Deprecated
     public Baggage set(TraceContext traceContext, String value) {
         this.delegate.updateValue(BraveTraceContext.toBrave(traceContext), value);
         return this;
@@ -66,6 +68,18 @@ public BaggageInScope makeCurrent() {
         return this;
     }
 
+    @Override
+    public BaggageInScope makeCurrent(String value) {
+        this.delegate.updateValue(value);
+        return makeCurrent();
+    }
+
+    @Override
+    public BaggageInScope makeCurrent(TraceContext traceContext, String value) {
+        this.delegate.updateValue(BraveTraceContext.toBrave(traceContext), value);
+        return makeCurrent();
+    }
+
     @Override
     public void close() {
     }

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageManager.java

@@ -17,6 +17,7 @@
 
 import brave.baggage.BaggageField;
 import io.micrometer.tracing.Baggage;
+import io.micrometer.tracing.BaggageInScope;
 import io.micrometer.tracing.BaggageManager;
 import io.micrometer.tracing.TraceContext;
 
@@ -32,7 +33,7 @@
  */
 public class BraveBaggageManager implements Closeable, BaggageManager {
 
-    private static final Map<String, Baggage> CACHE = new ConcurrentHashMap<>();
+    private static final Map<String, BraveBaggageInScope> CACHE = new ConcurrentHashMap<>();
 
     @Override
     public Map<String, String> getAllBaggage() {
@@ -54,13 +55,29 @@ public Baggage getBaggage(TraceContext traceContext, String name) {
     }
 
     @Override
+    @Deprecated
     public Baggage createBaggage(String name) {
+        return baggage(name);
+    }
+
+    private BraveBaggageInScope baggage(String name) {
         return CACHE.computeIfAbsent(name, s -> new BraveBaggageInScope(BaggageField.create(s)));
     }
 
     @Override
+    @Deprecated
     public Baggage createBaggage(String name, String value) {
-        return createBaggage(name).set(value);
+        return baggage(name).set(value);
+    }
+
+    @Override
+    public BaggageInScope createBaggageInScope(String name, String value) {
+        return baggage(name).makeCurrent(value);
+    }
+
+    @Override
+    public BaggageInScope createBaggageInScope(TraceContext traceContext, String name, String value) {
+        return baggage(name).makeCurrent(traceContext, value);
     }
 
     @Override

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveTracer.java

@@ -122,15 +122,27 @@ public Baggage getBaggage(TraceContext traceContext, String name) {
     }
 
     @Override
+    @Deprecated
     public Baggage createBaggage(String name) {
         return this.braveBaggageManager.createBaggage(name);
     }
 
     @Override
+    @Deprecated
     public Baggage createBaggage(String name, String value) {
         return this.braveBaggageManager.createBaggage(name).set(value);
     }
 
+    @Override
+    public BaggageInScope createBaggageInScope(String name, String value) {
+        return this.braveBaggageManager.createBaggageInScope(name, value);
+    }
+
+    @Override
+    public BaggageInScope createBaggageInScope(TraceContext traceContext, String name, String value) {
+        return this.braveBaggageManager.createBaggageInScope(traceContext, name, value);
+    }
+
     @Override
     public CurrentTraceContext currentTraceContext() {
         return this.currentTraceContext;

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/BaseTests.java

@@ -15,10 +15,6 @@
  */
 package io.micrometer.tracing.brave;
 
-import java.util.Optional;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-
 import brave.Tracing;
 import brave.handler.MutableSpan;
 import brave.propagation.StrictCurrentTraceContext;
@@ -34,6 +30,10 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
 import static org.assertj.core.api.BDDAssertions.then;
 
 /**

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveFinishedSpanTests.java

@@ -19,6 +19,7 @@
 import brave.Tracing;
 import brave.handler.MutableSpan;
 import io.micrometer.tracing.exporter.FinishedSpan;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
 import java.time.Instant;
@@ -28,7 +29,14 @@
 
 class BraveFinishedSpanTests {
 
-    Tracer tracer = Tracing.newBuilder().build().tracer();
+    Tracing tracing = Tracing.newBuilder().build();
+
+    Tracer tracer = tracing.tracer();
+
+    @AfterEach
+    void cleanup() {
+        tracing.close();
+    }
 
     @Test
     void should_calculate_instant_from_brave_timestamps() {

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveHttpClientHandlerTests.java

@@ -18,13 +18,20 @@
 import brave.Tracing;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
 class BraveHttpClientHandlerTests {
 
+    Tracing tracing = Tracing.newBuilder().build();
+
+    @AfterEach
+    void cleanup() {
+        tracing.close();
+    }
+
     @Test
     void should_not_throw_exception_when_response_null() {
-        Tracing tracing = Tracing.newBuilder().build();
         brave.http.HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> delegate = HttpClientHandler
                 .create(HttpTracing.newBuilder(tracing).build());
         BraveHttpClientHandler handler = new BraveHttpClientHandler(delegate);

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BravePropagatorTests.java

@@ -20,21 +20,21 @@
 import brave.baggage.BaggagePropagation;
 import brave.baggage.BaggagePropagationConfig;
 import brave.propagation.B3Propagation;
-import brave.propagation.CurrentTraceContext;
 import brave.propagation.Propagation;
 import brave.propagation.StrictCurrentTraceContext;
 import io.micrometer.tracing.BaggageInScope;
 import io.micrometer.tracing.Span;
 import io.micrometer.tracing.Tracer;
 import org.assertj.core.api.BDDAssertions;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
 import java.util.HashMap;
 import java.util.Map;
 
 class BravePropagatorTests {
 
-    CurrentTraceContext currentTraceContext = new StrictCurrentTraceContext();
+    StrictCurrentTraceContext currentTraceContext = new StrictCurrentTraceContext();
 
     Tracing tracing = Tracing.newBuilder()
             .propagationFactory(micrometerTracingPropagationWithBaggage(b3PropagationFactory()))
@@ -47,6 +47,12 @@ class BravePropagatorTests {
 
     BravePropagator bravePropagator = new BravePropagator(tracing);
 
+    @AfterEach
+    void cleanup() {
+        tracing.close();
+        currentTraceContext.close();
+    }
+
     @Test
     void should_propagate_context_with_trace_and_baggage() {
         Map<String, String> carrier = new HashMap<>();
@@ -56,7 +62,7 @@ void should_propagate_context_with_trace_and_baggage() {
 
         Span span = extract.start();
 
-        BaggageInScope baggage = braveBaggageManager.getBaggage(span.context(), "foo").makeCurrent();
+        BaggageInScope baggage = tracer.getBaggage(span.context(), "foo").makeCurrent();
         try {
             BDDAssertions.then(baggage.get(span.context())).isEqualTo("bar");
         }
@@ -73,7 +79,7 @@ void should_propagate_context_with_baggage_only() {
 
         Span span = extract.start();
 
-        BaggageInScope baggage = braveBaggageManager.getBaggage(span.context(), "foo").makeCurrent();
+        BaggageInScope baggage = tracer.getBaggage(span.context(), "foo").makeCurrent();
         try {
             BDDAssertions.then(baggage.get(span.context())).isEqualTo("bar");
         }

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveSpanBuilderTests.java

@@ -18,15 +18,23 @@
 import brave.Tracer;
 import brave.Tracing;
 import io.micrometer.tracing.Span;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.BDDAssertions.then;
 
 class BraveSpanBuilderTests {
 
+    Tracing tracing = Tracing.newBuilder().build();
+
+    @AfterEach
+    void cleanup() {
+        tracing.close();
+    }
+
     @Test
     void should_set_child_span_when_using_builders() {
-        Tracer tracer = Tracing.newBuilder().build().tracer();
+        Tracer tracer = tracing.tracer();
         Span.Builder builder = new BraveSpanBuilder(tracer);
         Span parentSpan = BraveSpan.fromBrave(tracer.nextSpan());
 

micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveTracingApiTests.java

@@ -66,13 +66,17 @@ class BraveTracingApiTests {
     // [Brave component] Tracer is a component that handles the life-cycle of a span
     brave.Tracer braveTracer = this.tracing.tracer();
 
+    // [Micrometer Tracing component] Baggage manager
+    BraveBaggageManager braveBaggageManager = new BraveBaggageManager();
+
     // [Micrometer Tracing component] A Micrometer Tracing wrapper for Brave's Tracer
-    Tracer tracer = new BraveTracer(this.braveTracer, this.bridgeContext, new BraveBaggageManager());
+    Tracer tracer = new BraveTracer(this.braveTracer, this.bridgeContext, this.braveBaggageManager);
 
     @AfterEach
     void close() {
         this.tracing.close();
         this.braveCurrentTraceContext.close();
+        this.braveBaggageManager.close();
     }
 
     @Test
@@ -151,7 +155,7 @@ void should_start_a_span_with_explicit_parent() throws Exception {
     }
 
     @Test
-    void should_work_with_baggage() {
+    void should_work_with_baggage_with_legacy_api() {
         Span span = tracer.nextSpan().name("parent").start();
 
         // Assuming that there's a span in scope...
@@ -163,21 +167,20 @@ void should_work_with_baggage() {
             Baggage baggageForSpanInScopeTwo = tracer.createBaggage("from_span_in_scope 2").set("value 2");
 
             try (BaggageInScope baggage = baggageForSpanInScopeOne.makeCurrent()) {
-                then(baggageForSpanInScopeOne.get()).as("[In scope] Baggage 1").isEqualTo("value 1");
+                then(baggage.get()).as("[In scope] Baggage 1").isEqualTo("value 1");
                 then(tracer.getBaggage("from_span_in_scope 1").get()).as("[In scope] Baggage 1").isEqualTo("value 1");
             }
 
             try (BaggageInScope baggage = baggageForSpanInScopeTwo.makeCurrent()) {
-                then(baggageForSpanInScopeTwo.get()).as("[In scope] Baggage 2").isEqualTo("value 2");
+                then(baggage.get()).as("[In scope] Baggage 2").isEqualTo("value 2");
                 then(tracer.getBaggage("from_span_in_scope 2").get()).as("[In scope] Baggage 2").isEqualTo("value 2");
             }
         }
 
         // Assuming that you have a handle to the span
         Baggage baggageForExplicitSpan = tracer.createBaggage("from_span").set(span.context(), "value 3");
         try (BaggageInScope baggage = baggageForExplicitSpan.makeCurrent()) {
-            then(baggageForExplicitSpan.get(span.context())).as("[Span passed explicitly] Baggage 3")
-                    .isEqualTo("value 3");
+            then(baggage.get(span.context())).as("[Span passed explicitly] Baggage 3").isEqualTo("value 3");
             then(tracer.getBaggage("from_span").get(span.context())).as("[Span passed explicitly] Baggage 3")
                     .isEqualTo("value 3");
         }
@@ -188,6 +191,51 @@ void should_work_with_baggage() {
         // When there's no span in scope, there will never be any baggage - even if you
         // make it current
         try (BaggageInScope baggage = baggageFour.makeCurrent()) {
+            then(baggage.get()).as("[Out of span scope] Baggage 1").isNull();
+            then(tracer.getBaggage("from_span_in_scope 1").get()).as("[Out of span scope] Baggage 1").isNull();
+        }
+        then(tracer.getBaggage("from_span_in_scope 1").get()).as("[Out of scope] Baggage 1").isNull();
+        then(tracer.getBaggage("from_span_in_scope 2").get()).as("[Out of scope] Baggage 2").isNull();
+        then(tracer.getBaggage("from_span").get()).as("[Out of scope] Baggage 3").isNull();
+
+        // You will retrieve the baggage value ALWAYS when you pass the context explicitly
+        then(tracer.getBaggage("from_span").get(span.context())).as("[Out of scope - with context] Baggage 3")
+                .isEqualTo("value 3");
+    }
+
+    @Test
+    void should_work_with_baggage() {
+        Span span = tracer.nextSpan().name("parent").start();
+
+        // Assuming that there's a span in scope...
+        try (Tracer.SpanInScope ws = tracer.withSpan(span)) {
+
+            try (BaggageInScope baggageForSpanInScopeOne = tracer.createBaggageInScope("from_span_in_scope 1",
+                    "value 1")) {
+                then(baggageForSpanInScopeOne.get()).as("[In scope] Baggage 1").isEqualTo("value 1");
+                then(tracer.getBaggage("from_span_in_scope 1").get()).as("[In scope] Baggage 1").isEqualTo("value 1");
+            }
+
+            try (BaggageInScope baggageForSpanInScopeTwo = tracer.createBaggageInScope("from_span_in_scope 2",
+                    "value 2");) {
+                then(baggageForSpanInScopeTwo.get()).as("[In scope] Baggage 2").isEqualTo("value 2");
+                then(tracer.getBaggage("from_span_in_scope 2").get()).as("[In scope] Baggage 2").isEqualTo("value 2");
+            }
+        }
+
+        // Assuming that you have a handle to the span
+        try (BaggageInScope baggageForExplicitSpan = tracer.createBaggageInScope(span.context(), "from_span",
+                "value 3")) {
+            then(baggageForExplicitSpan.get(span.context())).as("[Span passed explicitly] Baggage 3")
+                    .isEqualTo("value 3");
+            then(tracer.getBaggage("from_span").get(span.context())).as("[Span passed explicitly] Baggage 3")
+                    .isEqualTo("value 3");
+        }
+
+        // Assuming that there's no span in scope
+        // When there's no span in scope, there will never be any baggage - even if you
+        // make it current
+        try (BaggageInScope baggageFour = tracer.createBaggageInScope("from_span_in_scope 1", "value 1");) {
             then(baggageFour.get()).as("[Out of span scope] Baggage 1").isNull();
             then(tracer.getBaggage("from_span_in_scope 1").get()).as("[Out of span scope] Baggage 1").isNull();
         }

micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageInScope.java

@@ -75,6 +75,7 @@ public String get(TraceContext traceContext) {
     }
 
     @Override
+    @Deprecated
     public io.micrometer.tracing.Baggage set(String value) {
         return doSet(this.currentTraceContext.context(), value);
     }
@@ -112,10 +113,21 @@ private Entry entry() {
     }
 
     @Override
+    @Deprecated
     public io.micrometer.tracing.Baggage set(TraceContext traceContext, String value) {
         return doSet(traceContext, value);
     }
 
+    @Override
+    public BaggageInScope makeCurrent(String value) {
+        return doSet(currentTraceContext.context(), value).makeCurrent();
+    }
+
+    @Override
+    public BaggageInScope makeCurrent(TraceContext traceContext, String value) {
+        return doSet(traceContext, value).makeCurrent();
+    }
+
     @Override
     public BaggageInScope makeCurrent() {
         Entry entry = entry();