diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageInScope.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageInScope.java index 97ddffc2..c2284b0f 100755 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageInScope.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageInScope.java @@ -16,6 +16,9 @@ package io.micrometer.tracing.brave.bridge; import brave.baggage.BaggageField; +import io.micrometer.common.lang.Nullable; +import io.micrometer.common.util.internal.logging.InternalLogger; +import io.micrometer.common.util.internal.logging.InternalLoggerFactory; import io.micrometer.tracing.Baggage; import io.micrometer.tracing.BaggageInScope; import io.micrometer.tracing.TraceContext; @@ -28,10 +31,19 @@ */ class BraveBaggageInScope implements Baggage, BaggageInScope { + private static final InternalLogger logger = InternalLoggerFactory.getInstance(BraveBaggageInScope.class); + private final BaggageField delegate; - BraveBaggageInScope(BaggageField delegate) { + private final String previousBaggage; + + @Nullable + private brave.propagation.TraceContext traceContext; + + BraveBaggageInScope(BaggageField delegate, @Nullable brave.propagation.TraceContext traceContext) { this.delegate = delegate; + this.traceContext = traceContext; + this.previousBaggage = delegate.getValue(); } @Override @@ -41,7 +53,7 @@ public String name() { @Override public String get() { - return this.delegate.getValue(); + return this.traceContext != null ? this.delegate.getValue(traceContext) : this.delegate.getValue(); } @Override @@ -52,17 +64,34 @@ public String get(TraceContext traceContext) { @Override @Deprecated public Baggage set(String value) { - this.delegate.updateValue(value); + if (this.traceContext != null) { + this.delegate.updateValue(this.traceContext, value); + } + else { + this.delegate.updateValue(value); + } return this; } @Override @Deprecated public Baggage set(TraceContext traceContext, String value) { - this.delegate.updateValue(BraveTraceContext.toBrave(traceContext), value); + brave.propagation.TraceContext braveContext = updateBraveTraceContext(traceContext); + this.delegate.updateValue(braveContext, value); return this; } + private brave.propagation.TraceContext updateBraveTraceContext(TraceContext traceContext) { + brave.propagation.TraceContext braveContext = BraveTraceContext.toBrave(traceContext); + if (this.traceContext != braveContext) { + logger.debug( + "Create on baggage was called on <{}> but now you want to set baggage on <{}>. That's unexpected.", + this.traceContext, traceContext); + this.traceContext = braveContext; + } + return braveContext; + } + @Override public BaggageInScope makeCurrent() { return this; @@ -70,18 +99,27 @@ public BaggageInScope makeCurrent() { @Override public BaggageInScope makeCurrent(String value) { - this.delegate.updateValue(value); + if (this.traceContext != null) { + this.delegate.updateValue(this.traceContext, value); + } + else { + this.delegate.updateValue(value); + } return makeCurrent(); } @Override public BaggageInScope makeCurrent(TraceContext traceContext, String value) { - this.delegate.updateValue(BraveTraceContext.toBrave(traceContext), value); + brave.propagation.TraceContext braveContext = updateBraveTraceContext(traceContext); + this.delegate.updateValue(braveContext, value); return makeCurrent(); } @Override public void close() { + if (this.traceContext != null) { + this.delegate.updateValue(this.traceContext, this.previousBaggage); + } } } diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageManager.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageManager.java index 0ffe0532..95f80dfa 100755 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageManager.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveBaggageManager.java @@ -15,6 +15,7 @@ */ package io.micrometer.tracing.brave.bridge; +import brave.Tracing; import brave.baggage.BaggageField; import io.micrometer.tracing.Baggage; import io.micrometer.tracing.BaggageInScope; @@ -23,7 +24,6 @@ import java.io.Closeable; import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; /** * Brave implementation of a {@link BaggageManager}. @@ -33,13 +33,19 @@ */ public class BraveBaggageManager implements Closeable, BaggageManager { - private static final Map CACHE = new ConcurrentHashMap<>(); - @Override public Map getAllBaggage() { return BaggageField.getAllValues(); } + @Override + public Map getAllBaggage(TraceContext traceContext) { + if (traceContext == null) { + return getAllBaggage(); + } + return BaggageField.getAllValues(BraveTraceContext.toBrave(traceContext)); + } + @Override public Baggage getBaggage(String name) { return createBaggage(name); @@ -51,7 +57,7 @@ public Baggage getBaggage(TraceContext traceContext, String name) { if (baggageField == null) { return null; } - return new BraveBaggageInScope(baggageField); + return new BraveBaggageInScope(baggageField, BraveTraceContext.toBrave(traceContext)); } @Override @@ -60,8 +66,18 @@ public Baggage createBaggage(String name) { return baggage(name); } + private BraveBaggageInScope baggage(String name, TraceContext traceContext) { + return new BraveBaggageInScope(BaggageField.create(name), BraveTraceContext.toBrave(traceContext)); + } + private BraveBaggageInScope baggage(String name) { - return CACHE.computeIfAbsent(name, s -> new BraveBaggageInScope(BaggageField.create(s))); + return new BraveBaggageInScope(BaggageField.create(name), currentTraceContext()); + } + + // Taken from BraveField + private static brave.propagation.TraceContext currentTraceContext() { + Tracing tracing = Tracing.current(); + return tracing != null ? tracing.currentTraceContext().get() : null; } @Override @@ -77,12 +93,12 @@ public BaggageInScope createBaggageInScope(String name, String value) { @Override public BaggageInScope createBaggageInScope(TraceContext traceContext, String name, String value) { - return baggage(name).makeCurrent(traceContext, value); + return baggage(name, traceContext).makeCurrent(traceContext, value); } @Override public void close() { - CACHE.clear(); + } } diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveTracer.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveTracer.java index a5f43098..d0974f54 100755 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveTracer.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/main/java/io/micrometer/tracing/brave/bridge/BraveTracer.java @@ -106,6 +106,11 @@ public TraceContext.Builder traceContextBuilder() { return new BraveTraceContextBuilder(); } + @Override + public Map getAllBaggage(TraceContext traceContext) { + return this.braveBaggageManager.getAllBaggage(traceContext); + } + @Override public Map getAllBaggage() { return this.braveBaggageManager.getAllBaggage(); diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveTracingApiTests.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveTracingApiTests.java index 90adf50b..4e81b033 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveTracingApiTests.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/bridge/BraveTracingApiTests.java @@ -211,9 +211,8 @@ void should_work_with_baggage_with_legacy_api() { 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"); + // Baggage is present only within the scope + then(tracer.getBaggage("from_span").get(span.context())).as("[Out of scope - with context] Baggage 3").isNull(); } @Test @@ -256,9 +255,8 @@ void should_work_with_baggage() { 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"); + // Baggage is present only within the scope + then(tracer.getBaggage("from_span").get(span.context())).as("[Out of scope - with context] Baggage 3").isNull(); } @Test diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/contextpropagation/ScopesTests.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/contextpropagation/ScopesTests.java index 9732292b..ae8f478e 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/contextpropagation/ScopesTests.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-brave/src/test/java/io/micrometer/tracing/brave/contextpropagation/ScopesTests.java @@ -16,23 +16,25 @@ package io.micrometer.tracing.brave.contextpropagation; import brave.Tracing; -import brave.baggage.BaggageField; -import brave.baggage.CorrelationFlushScopeArrayReader; -import brave.baggage.CorrelationScopeConfig; +import brave.baggage.*; import brave.context.slf4j.MDCScopeDecorator; +import brave.propagation.B3Propagation; import brave.propagation.CurrentTraceContext; import brave.propagation.ThreadLocalCurrentTraceContext; +import brave.sampler.Sampler; import io.micrometer.common.util.internal.logging.InternalLogger; import io.micrometer.common.util.internal.logging.InternalLoggerFactory; import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor; +import io.micrometer.tracing.BaggageInScope; import io.micrometer.tracing.Span; import io.micrometer.tracing.Tracer; import io.micrometer.tracing.brave.bridge.BraveBaggageManager; import io.micrometer.tracing.brave.bridge.BraveCurrentTraceContext; import io.micrometer.tracing.brave.bridge.BraveTracer; import io.micrometer.tracing.handler.DefaultTracingObservationHandler; +import io.micrometer.tracing.handler.TracingObservationHandler.TracingContext; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -40,6 +42,7 @@ import reactor.core.publisher.Mono; import reactor.util.context.Context; +import java.util.Map; import java.util.concurrent.atomic.AtomicReference; import static org.assertj.core.api.BDDAssertions.then; @@ -50,13 +53,23 @@ class ScopesTests { CurrentTraceContext currentTraceContext = ThreadLocalCurrentTraceContext.newBuilder() .addScopeDecorator(MDCScopeDecorator.newBuilder() - .add(CorrelationScopeConfig.SingleCorrelationField.newBuilder(BaggageField.create("foo")) + .add(CorrelationScopeConfig.SingleCorrelationField + .newBuilder(BaggageField.create("foo")) .flushOnUpdate() .build()) .build()) .build(); - Tracing tracing = Tracing.newBuilder().currentTraceContext(currentTraceContext).build(); + Tracing tracing = Tracing.newBuilder() + .currentTraceContext(currentTraceContext) + .supportsJoin(false) + .traceId128Bit(true) + .propagationFactory(BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY) + .add(BaggagePropagationConfig.SingleBaggageField + .remote(BaggageField.create("foo"))) + .build()) + .sampler(Sampler.ALWAYS_SAMPLE) + .build(); brave.Tracer braveTracer = tracing.tracer(); @@ -133,6 +146,66 @@ void should_open_and_close_scopes_with_reactor() { logger.info("SPAN AFTER CLOSE 1 [" + tracer.currentSpan() + "]"); } + @Test + void should_open_and_close_scopes_with_reactor_with_baggage() { + Observation obs1 = Observation.start("1", observationRegistry); + + Observation.Scope scope = obs1.openScope(); + Span span1 = tracer.currentSpan(); + BaggageInScope baggageInScope1 = tracer.createBaggageInScope("foo", "span1"); + then(tracer.getAllBaggage().get("foo")).isEqualTo("span1"); + then(tracer.getAllBaggage(span1.context()).get("foo")).isEqualTo("span1"); + + Observation obs2 = Observation.start("2", observationRegistry); + Observation.Scope scope2 = obs2.openScope(); + Span span2 = tracer.currentSpan(); + BaggageInScope baggageInScope2 = tracer.createBaggageInScope(span2.context(), "foo", "span2"); + + then(tracer.getAllBaggage().get("foo")).isEqualTo("span2"); + then(tracer.getAllBaggage(span1.context()).get("foo")).isEqualTo("span1"); + then(tracer.getAllBaggage(span2.context()).get("foo")).isEqualTo("span2"); + + AtomicReference errorInFlatMap = new AtomicReference<>(); + AtomicReference errorInOnNext = new AtomicReference<>(); + + Mono.just(1).flatMap(integer -> { + return Mono.just(2).doOnNext(integer1 -> { + Map baggageInEmpty = tracer.getAllBaggage(); + logger.info("\n\n[2] BAGGAGE IN EMPTY [" + baggageInEmpty + "]"); + assertBaggageInReactor(errorInFlatMap, baggageInEmpty, null); + }).contextWrite(context -> Context.empty()); + }).doOnNext(integer -> { + Map baggageInOnNext = tracer.getAllBaggage(); + logger.info("\n\n[1] SPAN IN ON NEXT [" + baggageInOnNext + "]"); + assertBaggageInReactor(errorInOnNext, baggageInOnNext, "span2"); + }).contextWrite(context -> context.put(ObservationThreadLocalAccessor.KEY, obs2)).block(); + + logger.info("Checking if there were no errors in reactor"); + then(errorInFlatMap).hasValue(null); + then(errorInOnNext).hasValue(null); + + then(tracer.getBaggage(span2.context(), "foo").get()).isEqualTo("span2"); + then(tracer.getBaggage(span1.context(), "foo").get()).isEqualTo("span1"); + + baggageInScope2.close(); + scope2.close(); + obs2.stop(); + TracingContext tracingContext2 = obs2.getContext().get(TracingContext.class); + then(tracingContext2.getBaggage()).isNullOrEmpty(); + + then(tracer.getBaggage("foo").get()).isEqualTo("span1"); + then(tracer.getBaggage(span1.context(), "foo").get()).isEqualTo("span1"); + then(tracer.getBaggage(span2.context(), "foo").get()).isEqualTo("span1"); + + baggageInScope1.close(); + scope.close(); + obs1.stop(); + TracingContext tracingContext1 = obs1.getContext().get(TracingContext.class); + then(tracingContext1.getBaggage()).isNullOrEmpty(); + + then(tracer.getAllBaggage()).isEmpty(); + } + private static void assertInReactor(AtomicReference errors, Span spanWOnNext, Span expectedSpan) { try { then(spanWOnNext).isEqualTo(expectedSpan); @@ -142,4 +215,13 @@ private static void assertInReactor(AtomicReference errors, Span } } + private static void assertBaggageInReactor(AtomicReference errors, Map baggageMap, String expectedValue) { + try { + then(baggageMap.get("foo")).isEqualTo(expectedValue); + } + catch (AssertionError er) { + errors.set(er); + } + } + } diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageInScope.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageInScope.java index 5b5783bb..116060a9 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageInScope.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageInScope.java @@ -43,6 +43,10 @@ class OtelBaggageInScope implements io.micrometer.tracing.Baggage, BaggageInScop private final AtomicReference entry = new AtomicReference<>(); + private final AtomicReference contextWithoutBaggage = new AtomicReference<>(null); + + private final AtomicReference mutatedTraceContext = new AtomicReference<>(null); + private final AtomicReference contextWithBaggage = new AtomicReference<>(null); private final AtomicReference scope = new AtomicReference<>(); @@ -88,7 +92,9 @@ private io.micrometer.tracing.Baggage doSet(TraceContext context, String value) Span currentSpan = Span.current(); io.opentelemetry.api.baggage.Baggage baggage; OtelTraceContext ctx = (OtelTraceContext) context; + mutatedTraceContext.set(ctx); Context storedCtx = ctx.context(); + contextWithoutBaggage.set(storedCtx); Baggage fromContext = Baggage.fromContext(storedCtx); BaggageBuilder newBaggageBuilder = fromContext.toBuilder(); @@ -152,6 +158,10 @@ public void close() { if (scope != null) { this.scope.set(null); scope.close(); + OtelTraceContext traceContext = this.mutatedTraceContext.get(); + if (traceContext != null) { + traceContext.updateContext(this.contextWithoutBaggage.get()); + } } } diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageManager.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageManager.java index 1e7a6b7f..48cc59de 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageManager.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelBaggageManager.java @@ -68,13 +68,28 @@ public OtelBaggageManager(CurrentTraceContext currentTraceContext, List @Override public Map getAllBaggage() { + return toMap(currentBaggage()); + } + + private Map toMap(CompositeBaggage compositeBaggage) { Map baggage = new HashMap<>(); - currentBaggage().getEntries().forEach(entry -> baggage.put(entry.getKey(), entry.getValue())); + compositeBaggage.getEntries().forEach(entry -> baggage.put(entry.getKey(), entry.getValue())); return baggage; } + @Override + public Map getAllBaggage(TraceContext traceContext) { + if (traceContext == null) { + return getAllBaggage(); + } + return toMap(baggage((OtelTraceContext) traceContext)); + } + CompositeBaggage currentBaggage() { - OtelTraceContext traceContext = (OtelTraceContext) currentTraceContext.context(); + return baggage((OtelTraceContext) currentTraceContext.context()); + } + + private CompositeBaggage baggage(OtelTraceContext traceContext) { Context context = Context.current(); Deque stack = new ArrayDeque<>(); if (traceContext != null) { diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelTracer.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelTracer.java index 710a7081..db66bab2 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelTracer.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/main/java/io/micrometer/tracing/otel/bridge/OtelTracer.java @@ -152,6 +152,11 @@ public Map getAllBaggage() { return this.otelBaggageManager.getAllBaggage(); } + @Override + public Map getAllBaggage(TraceContext traceContext) { + return this.otelBaggageManager.getAllBaggage(traceContext); + } + @Override public Baggage getBaggage(String name) { return this.otelBaggageManager.getBaggage(name); diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/bridge/OtelTracingApiTests.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/bridge/OtelTracingApiTests.java index 4c18f45b..cdd684f9 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/bridge/OtelTracingApiTests.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/bridge/OtelTracingApiTests.java @@ -211,9 +211,8 @@ void should_work_with_baggage_with_legacy_api() { 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"); + // Baggage is present only within the scope + then(tracer.getBaggage("from_span").get(span.context())).as("[Out of scope - with context] Baggage 3").isNull(); } @Test @@ -256,9 +255,8 @@ void should_work_with_baggage() { 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"); + // Baggage is present only within the scope + then(tracer.getBaggage("from_span").get(span.context())).as("[Out of scope - with context] Baggage 3").isNull(); } @Test diff --git a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/contextpropagation/ScopesTests.java b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/contextpropagation/ScopesTests.java index 36a4674e..702799fa 100644 --- a/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/contextpropagation/ScopesTests.java +++ b/micrometer-tracing-bridges/micrometer-tracing-bridge-otel/src/test/java/io/micrometer/tracing/otel/contextpropagation/ScopesTests.java @@ -20,9 +20,11 @@ import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor; +import io.micrometer.tracing.BaggageInScope; import io.micrometer.tracing.Span; import io.micrometer.tracing.Tracer; import io.micrometer.tracing.handler.DefaultTracingObservationHandler; +import io.micrometer.tracing.handler.TracingObservationHandler.TracingContext; import io.micrometer.tracing.otel.bridge.ArrayListSpanProcessor; import io.micrometer.tracing.otel.bridge.OtelBaggageManager; import io.micrometer.tracing.otel.bridge.OtelCurrentTraceContext; @@ -39,6 +41,7 @@ import reactor.util.context.Context; import java.util.Collections; +import java.util.Map; import java.util.concurrent.atomic.AtomicReference; import static org.assertj.core.api.BDDAssertions.then; @@ -62,7 +65,7 @@ class ScopesTests { io.opentelemetry.api.trace.Tracer otelTracer = openTelemetrySdk.getTracer("io.micrometer.micrometer-tracing"); Tracer tracer = new OtelTracer(otelTracer, new OtelCurrentTraceContext(), event -> { - }, new OtelBaggageManager(new OtelCurrentTraceContext(), Collections.emptyList(), Collections.emptyList())); + }, new OtelBaggageManager(new OtelCurrentTraceContext(), Collections.singletonList("foo"), Collections.emptyList())); DefaultTracingObservationHandler handler = new DefaultTracingObservationHandler(tracer); @@ -122,6 +125,66 @@ void should_open_and_close_scopes_with_reactor() { logger.info("SPAN AFTER CLOSE 1 [" + tracer.currentSpan() + "]"); } + @Test + void should_open_and_close_scopes_with_reactor_with_baggage() { + Observation obs1 = Observation.start("1", observationRegistry); + + Observation.Scope scope = obs1.openScope(); + Span span1 = tracer.currentSpan(); + BaggageInScope baggageInScope1 = tracer.createBaggageInScope("foo", "span1"); + then(tracer.getAllBaggage().get("foo")).isEqualTo("span1"); + then(tracer.getAllBaggage(span1.context()).get("foo")).isEqualTo("span1"); + + Observation obs2 = Observation.start("2", observationRegistry); + Observation.Scope scope2 = obs2.openScope(); + Span span2 = tracer.currentSpan(); + BaggageInScope baggageInScope2 = tracer.createBaggageInScope(span2.context(), "foo", "span2"); + + then(tracer.getAllBaggage().get("foo")).isEqualTo("span2"); + then(tracer.getAllBaggage(span1.context()).get("foo")).isEqualTo("span1"); + then(tracer.getAllBaggage(span2.context()).get("foo")).isEqualTo("span2"); + + AtomicReference errorInFlatMap = new AtomicReference<>(); + AtomicReference errorInOnNext = new AtomicReference<>(); + + Mono.just(1).flatMap(integer -> { + return Mono.just(2).doOnNext(integer1 -> { + Map baggageInEmpty = tracer.getAllBaggage(); + logger.info("\n\n[2] BAGGAGE IN EMPTY [" + baggageInEmpty + "]"); + assertBaggageInReactor(errorInFlatMap, baggageInEmpty, null); + }).contextWrite(context -> Context.empty()); + }).doOnNext(integer -> { + Map baggageInOnNext = tracer.getAllBaggage(); + logger.info("\n\n[1] SPAN IN ON NEXT [" + baggageInOnNext + "]"); + assertBaggageInReactor(errorInOnNext, baggageInOnNext, "span2"); + }).contextWrite(context -> context.put(ObservationThreadLocalAccessor.KEY, obs2)).block(); + + logger.info("Checking if there were no errors in reactor"); + then(errorInFlatMap).hasValue(null); + then(errorInOnNext).hasValue(null); + + then(tracer.getBaggage(span2.context(), "foo").get()).isEqualTo("span2"); + then(tracer.getBaggage(span1.context(), "foo").get()).isEqualTo("span1"); + + baggageInScope2.close(); + scope2.close(); + obs2.stop(); + TracingContext tracingContext2 = obs2.getContext().get(TracingContext.class); + then(tracingContext2.getBaggage()).isNullOrEmpty(); + + then(tracer.getBaggage("foo").get()).isEqualTo("span1"); + then(tracer.getBaggage(span1.context(), "foo").get()).isEqualTo("span1"); + then(tracer.getBaggage(span2.context(), "foo").get()).isEqualTo("span1"); + + baggageInScope1.close(); + scope.close(); + obs1.stop(); + TracingContext tracingContext1 = obs1.getContext().get(TracingContext.class); + then(tracingContext1.getBaggage()).isNullOrEmpty(); + + then(tracer.getAllBaggage()).isEmpty(); + } + private static void assertInReactor(AtomicReference errors, Span spanWOnNext, Span expectedSpan) { try { then(spanWOnNext).isEqualTo(expectedSpan); @@ -131,4 +194,13 @@ private static void assertInReactor(AtomicReference errors, Span } } + private static void assertBaggageInReactor(AtomicReference errors, Map baggageMap, String expectedValue) { + try { + then(baggageMap.get("foo")).isEqualTo(expectedValue); + } + catch (AssertionError er) { + errors.set(er); + } + } + } diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/BaggageManager.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/BaggageManager.java index 858ecad9..a6c7b399 100644 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/BaggageManager.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/BaggageManager.java @@ -75,6 +75,18 @@ public BaggageInScope createBaggageInScope(TraceContext traceContext, String nam */ Map getAllBaggage(); + /** + * @param traceContext trace context with baggage. If {@code null} will try to get all + * baggage from current available context + * @return mapping of all baggage entries from the given scope + */ + default Map getAllBaggage(@Nullable TraceContext traceContext) { + if (traceContext == null) { + return getAllBaggage(); + } + return Collections.emptyMap(); + } + /** * Retrieves {@link Baggage} for the given name. * @param name baggage name diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/DefaultTracingObservationHandler.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/DefaultTracingObservationHandler.java index 351d48b6..58922d6b 100755 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/DefaultTracingObservationHandler.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/DefaultTracingObservationHandler.java @@ -50,7 +50,7 @@ public void onStop(Observation.Context context) { Span span = getRequiredSpan(context); span.name(getSpanName(context)); tagSpan(context, span); - span.end(); + endSpan(context, span); } @Override diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingReceiverTracingObservationHandler.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingReceiverTracingObservationHandler.java index c2c7b053..ce0b8f3e 100644 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingReceiverTracingObservationHandler.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingReceiverTracingObservationHandler.java @@ -98,7 +98,7 @@ public void onStop(T context) { tagSpan(context, span); customizeReceiverSpan(context, span); span.name(context.getContextualName() != null ? context.getContextualName() : context.getName()); - span.end(); + endSpan(context, span); } /** diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingSenderTracingObservationHandler.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingSenderTracingObservationHandler.java index c03e4e93..612c1c79 100644 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingSenderTracingObservationHandler.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/PropagatingSenderTracingObservationHandler.java @@ -101,7 +101,7 @@ public void onStop(T context) { tagSpan(context, span); customizeSenderSpan(context, span); span.name(context.getContextualName() != null ? context.getContextualName() : context.getName()); - span.end(); + endSpan(context, span); } /** diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/RevertingScope.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/RevertingScope.java index 8fa66a8a..308709f6 100644 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/RevertingScope.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/RevertingScope.java @@ -16,7 +16,10 @@ package io.micrometer.tracing.handler; import io.micrometer.tracing.CurrentTraceContext; +import io.micrometer.tracing.CurrentTraceContext.Scope; +import io.micrometer.tracing.handler.TracingObservationHandler.TracingContext; +import java.util.Map; import java.util.Objects; class RevertingScope implements CurrentTraceContext.Scope { @@ -27,17 +30,21 @@ class RevertingScope implements CurrentTraceContext.Scope { private final CurrentTraceContext.Scope previousScope; - RevertingScope(TracingObservationHandler.TracingContext tracingContext, CurrentTraceContext.Scope currentScope, - CurrentTraceContext.Scope previousScope) { + private final Map previousBaggage; + + RevertingScope(TracingContext tracingContext, Scope currentScope, Scope previousScope, + Map previousBaggage) { this.tracingContext = tracingContext; this.currentScope = currentScope; this.previousScope = previousScope; + this.previousBaggage = previousBaggage; } @Override public void close() { this.currentScope.close(); this.tracingContext.setScope(this.previousScope); + this.tracingContext.setBaggage(this.previousBaggage); } @Override @@ -55,12 +62,13 @@ public boolean equals(Object o) { } RevertingScope that = (RevertingScope) o; return Objects.equals(tracingContext, that.tracingContext) && Objects.equals(currentScope, that.currentScope) - && Objects.equals(previousScope, that.previousScope); + && Objects.equals(previousScope, that.previousScope) + && Objects.equals(previousBaggage, that.previousBaggage); } @Override public int hashCode() { - return Objects.hash(tracingContext, currentScope, previousScope); + return Objects.hash(tracingContext, currentScope, previousScope, previousBaggage); } } diff --git a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/TracingObservationHandler.java b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/TracingObservationHandler.java index cddb9ad7..abd216e5 100755 --- a/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/TracingObservationHandler.java +++ b/micrometer-tracing/src/main/java/io/micrometer/tracing/handler/TracingObservationHandler.java @@ -29,6 +29,7 @@ import io.micrometer.tracing.Tracer; import io.micrometer.tracing.internal.SpanNameUtil; +import java.util.Collections; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -93,10 +94,17 @@ default void onScopeOpened(T context) { default void setMaybeScopeOnTracingContext(TracingContext tracingContext, @Nullable Span newSpan) { Span spanFromThisObservation = tracingContext.getSpan(); TraceContext newContext = newSpan != null ? newSpan.context() : null; + Map previousBaggage = Collections.emptyMap(); + if (spanFromThisObservation != null) { + previousBaggage = getTracer() + .getAllBaggage(spanFromThisObservation.context()); + } CurrentTraceContext.Scope scope = getTracer().currentTraceContext().maybeScope(newContext); CurrentTraceContext.Scope previousScopeOnThisObservation = tracingContext.getScope(); + Map newBaggage = getTracer().getAllBaggage(newContext); + tracingContext.setBaggage(newBaggage); tracingContext.setSpanAndScope(spanFromThisObservation, - new RevertingScope(tracingContext, scope, previousScopeOnThisObservation)); + new RevertingScope(tracingContext, scope, previousScopeOnThisObservation, previousBaggage)); } @Override @@ -197,6 +205,18 @@ default Span getRequiredSpan(T context) { return span; } + /** + * Ends span and clears resources. + * + * @param context context + * @param span span to end + */ + default void endSpan(T context, Span span) { + TracingContext tracingContext = getTracingContext(context); + tracingContext.close(); + span.end(); + } + /** * Returns the {@link Tracer}. * @return tracer @@ -209,12 +229,14 @@ default Span getRequiredSpan(T context) { * @author Marcin Grzejszczak * @since 1.0.0 */ - class TracingContext { + class TracingContext implements AutoCloseable { private Span span; private Map scopes = new ConcurrentHashMap<>(); + private Map> baggage = new ConcurrentHashMap<>(); + /** * Returns the span. * @return span @@ -252,6 +274,27 @@ public void setScope(CurrentTraceContext.Scope scope) { } } + /** + * Returns the baggage corresponding to this span. + * @return baggage attached to the span + */ + public Map getBaggage() { + return this.baggage.get(Thread.currentThread()); + } + + /** + * Sets the baggage + * @param baggage baggage to set + */ + public void setBaggage(Map baggage) { + if (baggage == null) { + this.baggage.remove(Thread.currentThread()); + } + else { + this.baggage.put(Thread.currentThread(), baggage); + } + } + /** * Convenience method to set both span and scope. * @param span span to set @@ -262,9 +305,15 @@ public void setSpanAndScope(Span span, CurrentTraceContext.Scope scope) { setScope(scope); } + @Override + public void close() { + this.baggage.clear(); + this.scopes.clear(); + } + @Override public String toString() { - return "TracingContext{" + "span=" + traceContextFromSpan() + '}'; + return "TracingContext{" + "span=" + traceContextFromSpan() + ",baggage=" + baggage + '}'; } private String traceContextFromSpan() { diff --git a/micrometer-tracing/src/test/java/io/micrometer/tracing/handler/TracingObservationHandlerTests.java b/micrometer-tracing/src/test/java/io/micrometer/tracing/handler/TracingObservationHandlerTests.java index a229321c..57c9280a 100644 --- a/micrometer-tracing/src/test/java/io/micrometer/tracing/handler/TracingObservationHandlerTests.java +++ b/micrometer-tracing/src/test/java/io/micrometer/tracing/handler/TracingObservationHandlerTests.java @@ -24,6 +24,8 @@ import org.mockito.BDDMockito; import org.mockito.InOrder; +import java.util.Collections; + import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.BDDAssertions.thenNoException; import static org.assertj.core.api.BDDAssertions.thenThrownBy; @@ -58,8 +60,9 @@ void spanShouldBeClearedOnScopeReset() { TracingObservationHandler.TracingContext tracingContext = new TracingObservationHandler.TracingContext(); CurrentTraceContext.Scope scope1 = mock(CurrentTraceContext.Scope.class); CurrentTraceContext.Scope scope2 = mock(CurrentTraceContext.Scope.class); - RevertingScope revertingScope1 = new RevertingScope(tracingContext, scope1, null); - RevertingScope revertingScope2 = new RevertingScope(tracingContext, scope2, revertingScope1); + RevertingScope revertingScope1 = new RevertingScope(tracingContext, scope1, null, Collections.emptyMap()); + RevertingScope revertingScope2 = new RevertingScope(tracingContext, scope2, revertingScope1, + Collections.emptyMap()); tracingContext.setScope(revertingScope2); context.put(TracingObservationHandler.TracingContext.class, tracingContext);