MDC trace and span IDs for Brave and OTel

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/BraveAutoConfiguration.java

@@ -21,6 +21,9 @@
 import brave.Tracing;
 import brave.Tracing.Builder;
 import brave.TracingCustomizer;
+import brave.baggage.CorrelationScopeCustomizer;
+import brave.baggage.CorrelationScopeDecorator;
+import brave.context.slf4j.MDCScopeDecorator;
 import brave.handler.SpanHandler;
 import brave.http.HttpClientHandler;
 import brave.http.HttpClientRequest;
@@ -41,9 +44,12 @@
 import io.micrometer.tracing.brave.bridge.BraveHttpClientHandler;
 import io.micrometer.tracing.brave.bridge.BraveHttpServerHandler;
 import io.micrometer.tracing.brave.bridge.BraveTracer;
+import org.slf4j.MDC;
 
+import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
@@ -166,6 +172,28 @@ BraveHttpClientHandler braveHttpClientHandler(
 			return new BraveHttpClientHandler(httpClientHandler);
 		}
 
+		@Configuration(proxyBeanMethods = false)
+		static class CorrelationScopeDecoratorConfiguration {
+
+			@Bean
+			@ConditionalOnMissingBean
+			@ConditionalOnClass(MDC.class)
+			CorrelationScopeDecorator.Builder mdcCorrelationScopeDecoratorBuilder() {
+				return MDCScopeDecorator.newBuilder();
+			}
+
+			@Bean
+			@ConditionalOnMissingBean(CorrelationScopeDecorator.class)
+			@ConditionalOnBean(CorrelationScopeDecorator.Builder.class)
+			ScopeDecorator correlationScopeDecorator(CorrelationScopeDecorator.Builder builder,
+					ObjectProvider<List<CorrelationScopeCustomizer>> correlationScopeCustomizers) {
+				correlationScopeCustomizers.ifAvailable(
+						(customizers) -> customizers.forEach((customizer) -> customizer.customize(builder)));
+				return builder.build();
+			}
+
+		}
+
 	}
 
 }

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/OpenTelemetryConfigurations.java

@@ -23,15 +23,19 @@
 import io.micrometer.tracing.SamplerFunction;
 import io.micrometer.tracing.otel.bridge.DefaultHttpClientAttributesGetter;
 import io.micrometer.tracing.otel.bridge.DefaultHttpServerAttributesExtractor;
+import io.micrometer.tracing.otel.bridge.EventListener;
+import io.micrometer.tracing.otel.bridge.EventPublishingContextWrapper;
 import io.micrometer.tracing.otel.bridge.OtelBaggageManager;
 import io.micrometer.tracing.otel.bridge.OtelCurrentTraceContext;
 import io.micrometer.tracing.otel.bridge.OtelHttpClientHandler;
 import io.micrometer.tracing.otel.bridge.OtelHttpServerHandler;
 import io.micrometer.tracing.otel.bridge.OtelTracer;
 import io.micrometer.tracing.otel.bridge.OtelTracer.EventPublisher;
+import io.micrometer.tracing.otel.bridge.Slf4JEventListener;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.ContextStorage;
 import io.opentelemetry.context.propagation.ContextPropagators;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
@@ -43,6 +47,7 @@
 import io.opentelemetry.sdk.trace.export.SpanExporter;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import io.opentelemetry.semconv.resource.attributes.ResourceAttributes;
+import org.slf4j.MDC;
 
 import org.springframework.boot.SpringBootVersion;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
@@ -141,14 +146,14 @@ OtelTracer micrometerOtelTracer(Tracer tracer, EventPublisher eventPublisher,
 
 		@Bean
 		@ConditionalOnMissingBean
-		EventPublisher otelTracerEventPublisher() {
-			return (event) -> {
-			};
+		EventPublisher otelTracerEventPublisher(List<EventListener> eventListeners) {
+			return new OTelEventPublisher(eventListeners);
 		}
 
 		@Bean
 		@ConditionalOnMissingBean
-		OtelCurrentTraceContext otelCurrentTraceContext() {
+		OtelCurrentTraceContext otelCurrentTraceContext(EventPublisher publisher) {
+			ContextStorage.addWrapper(new EventPublishingContextWrapper(publisher));
 			return new OtelCurrentTraceContext();
 		}
 
@@ -168,6 +173,35 @@ OtelHttpServerHandler otelHttpServerHandler(OpenTelemetry openTelemetry) {
 					new DefaultHttpServerAttributesExtractor());
 		}
 
+		static class OTelEventPublisher implements EventPublisher {
+
+			private final List<EventListener> listeners;
+
+			OTelEventPublisher(List<EventListener> listeners) {
+				this.listeners = listeners;
+			}
+
+			@Override
+			public void publishEvent(Object event) {
+				for (EventListener listener : this.listeners) {
+					listener.onEvent(event);
+				}
+			}
+
+		}
+
+		@Configuration(proxyBeanMethods = false)
+		@ConditionalOnClass(MDC.class)
+		static class Slf4jConfiguration {
+
+			@Bean
+			@ConditionalOnMissingBean
+			Slf4JEventListener otelSlf4JEventListener() {
+				return new Slf4JEventListener();
+			}
+
+		}
+
 	}
 
 }

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/BraveAutoConfigurationTests.java

@@ -152,6 +152,17 @@ void shouldNotSupplyBraveBeansIfTracingIsDisabled() {
 		});
 	}
 
+	@Test
+	void shouldNotSupplyMdcCorrelationScopeWhenMdcNotOnClasspath() {
+		this.contextRunner.withClassLoader(new FilteredClassLoader("org.slf4j"))
+				.run((context) -> assertThat(context).doesNotHaveBean("mdcCorrelationScopeDecoratorBuilder"));
+	}
+
+	@Test
+	void shouldSupplyMdcCorrelationScopeDecoratorWhenMdcOnClasspath() {
+		this.contextRunner.run((context) -> assertThat(context).hasBean("mdcCorrelationScopeDecoratorBuilder"));
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	private static class CustomBraveConfiguration {
 

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/MdcAutoConfigurationTests.java

@@ -0,0 +1,139 @@
+/*
+ * Copyright 2012-2022 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.boot.actuate.autoconfigure.tracing;
+
+import java.util.function.Supplier;
+
+import io.micrometer.tracing.Span;
+import io.micrometer.tracing.Tracer;
+import io.opentelemetry.context.Context;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.slf4j.MDC;
+
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.ApplicationContext;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for Baggage configuration.
+ *
+ * @author Marcin Grzejszczak
+ */
+class MdcAutoConfigurationTests {
+
+	@BeforeEach
+	@AfterEach
+	void setup() {
+		MDC.clear();
+	}
+
+	@ParameterizedTest
+	@EnumSource(AutoConfig.class)
+	void shouldSetEntriesToMdcFromSpan(AutoConfig autoConfig) {
+		autoConfig.get().run((context) -> {
+			Tracer tracer = tracer(context);
+			Span span = createSpan(tracer);
+			assertThatTracingContextIsInitialized(autoConfig);
+			try (Tracer.SpanInScope scope = tracer.withSpan(span.start())) {
+				assertThat(MDC.get("traceId")).isEqualTo(span.context().traceId());
+			}
+			finally {
+				span.end();
+			}
+			assertThatMdcContainsUnsetTraceId();
+		});
+	}
+
+	@ParameterizedTest
+	@EnumSource(AutoConfig.class)
+	void shouldRemoveEntriesFromMdcFromNullSpan(AutoConfig autoConfig) {
+		autoConfig.get().run((context) -> {
+			Tracer tracer = tracer(context);
+			Span span = createSpan(tracer);
+			assertThatTracingContextIsInitialized(autoConfig);
+			// can't use NOOP as it is special cased
+			try (Tracer.SpanInScope scope = tracer.withSpan(span.start())) {
+				assertThat(MDC.get("traceId")).isEqualTo(span.context().traceId());
+
+				// can't use NOOP as it is special cased
+				try (Tracer.SpanInScope scope2 = tracer.withSpan(null)) {
+					assertThatMdcContainsUnsetTraceId();
+				}
+
+				assertThat(MDC.get("traceId")).isEqualTo(span.context().traceId());
+			}
+			finally {
+				span.end();
+			}
+			assertThatMdcContainsUnsetTraceId();
+
+		});
+
+	}
+
+	private Span createSpan(Tracer tracer) {
+		return tracer.nextSpan().name("span");
+	}
+
+	private Tracer tracer(ApplicationContext context) {
+		return context.getBean(Tracer.class);
+	}
+
+	private void assertThatTracingContextIsInitialized(AutoConfig autoConfig) {
+		if (autoConfig == AutoConfig.OTEL) {
+			assertThat(Context.current()).isEqualTo(Context.root());
+		}
+	}
+
+	private void assertThatMdcContainsUnsetTraceId() {
+		assertThat(isInvalidBraveTraceId() || isInvalidOtelTraceId()).isTrue();
+	}
+
+	private boolean isInvalidBraveTraceId() {
+		return MDC.get("traceId") == null;
+	}
+
+	private boolean isInvalidOtelTraceId() {
+		return MDC.get("traceId").equals("00000000000000000000000000000000");
+	}
+
+	enum AutoConfig implements Supplier<ApplicationContextRunner> {
+
+		BRAVE {
+			@Override
+			public ApplicationContextRunner get() {
+				return new ApplicationContextRunner()
+						.withConfiguration(AutoConfigurations.of(BraveAutoConfiguration.class));
+			}
+		},
+
+		OTEL {
+			@Override
+			public ApplicationContextRunner get() {
+				return new ApplicationContextRunner()
+						.withConfiguration(AutoConfigurations.of(OpenTelemetryAutoConfiguration.class));
+			}
+		}
+
+	}
+
+}

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/OpenTelemetryConfigurationsMicrometerConfigurationTests.java

@@ -21,6 +21,7 @@
 import io.micrometer.tracing.otel.bridge.OtelHttpServerHandler;
 import io.micrometer.tracing.otel.bridge.OtelTracer;
 import io.micrometer.tracing.otel.bridge.OtelTracer.EventPublisher;
+import io.micrometer.tracing.otel.bridge.Slf4JEventListener;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.api.trace.Tracer;
 import org.junit.jupiter.api.Test;
@@ -99,6 +100,19 @@ void shouldBackOffOnCustomBeans() {
 		});
 	}
 
+	@Test
+	void shouldSupplySlf4jEventListenersWhenMdcOnClasspath() {
+		this.contextRunner.withUserConfiguration(TracerConfiguration.class)
+				.run((context) -> assertThat(context).hasSingleBean(Slf4JEventListener.class));
+	}
+
+	@Test
+	void shouldNotSupplySlf4jEventListenersWhenMdcNotOnClasspath() {
+		this.contextRunner.withClassLoader(new FilteredClassLoader("org.slf4j"))
+				.withUserConfiguration(TracerConfiguration.class)
+				.run((context) -> assertThat(context).doesNotHaveBean(Slf4JEventListener.class));
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	private static class CustomConfiguration {