[Observability] MDC, Context Propagation (via B3 and W3C) and Baggage

Author: marcingrzejszczak

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

@@ -21,6 +21,14 @@
 import brave.Tracing;
 import brave.Tracing.Builder;
 import brave.TracingCustomizer;
+import brave.baggage.BaggageField;
+import brave.baggage.BaggagePropagation;
+import brave.baggage.BaggagePropagationConfig;
+import brave.baggage.BaggagePropagationCustomizer;
+import brave.baggage.CorrelationScopeConfig;
+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;
@@ -33,19 +41,27 @@
 import brave.propagation.CurrentTraceContext;
 import brave.propagation.CurrentTraceContext.ScopeDecorator;
 import brave.propagation.CurrentTraceContextCustomizer;
+import brave.propagation.Propagation;
 import brave.propagation.Propagation.Factory;
 import brave.propagation.ThreadLocalCurrentTraceContext;
 import brave.sampler.Sampler;
 import io.micrometer.tracing.brave.bridge.BraveBaggageManager;
 import io.micrometer.tracing.brave.bridge.BraveCurrentTraceContext;
 import io.micrometer.tracing.brave.bridge.BraveHttpClientHandler;
 import io.micrometer.tracing.brave.bridge.BraveHttpServerHandler;
+import io.micrometer.tracing.brave.bridge.BravePropagator;
 import io.micrometer.tracing.brave.bridge.BraveTracer;
+import io.micrometer.tracing.brave.bridge.W3CPropagation;
+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.autoconfigure.condition.ConditionalOnMissingClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -55,6 +71,7 @@
  * {@link EnableAutoConfiguration Auto-configuration} for Brave.
  *
  * @author Moritz Halbritter
+ * @author Marcin Grzejszczak
  * @since 3.0.0
  */
 @AutoConfiguration(before = MicrometerTracingAutoConfiguration.class)
@@ -105,12 +122,6 @@ public CurrentTraceContext braveCurrentTraceContext(List<CurrentTraceContext.Sco
 		return builder.build();
 	}
 
-	@Bean
-	@ConditionalOnMissingBean
-	public Factory bravePropagationFactory() {
-		return B3Propagation.newFactoryBuilder().injectFormat(B3Propagation.Format.SINGLE_NO_PARENT).build();
-	}
-
 	@Bean
 	@ConditionalOnMissingBean
 	public Sampler braveSampler(TracingProperties properties) {
@@ -135,21 +146,35 @@ public HttpClientHandler<HttpClientRequest, HttpClientResponse> httpClientHandle
 		return HttpClientHandler.create(httpTracing);
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnMissingClass("io.micrometer.tracing.brave.bridge.BraveTracer")
+	static class BraveMicrometerMissing {
+
+		@Bean
+		@ConditionalOnMissingBean
+		@ConditionalOnProperty(value = "management.tracing.propagation.type", havingValue = "B3", matchIfMissing = true)
+		Factory bravePropagationFactory() {
+			return B3Propagation.newFactoryBuilder().injectFormat(B3Propagation.Format.SINGLE_NO_PARENT).build();
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	@ConditionalOnClass(BraveTracer.class)
 	static class BraveMicrometer {
 
+		private static final BraveBaggageManager BRAVE_BAGGAGE_MANAGER = new BraveBaggageManager();
+
 		@Bean
 		@ConditionalOnMissingBean
-		BraveTracer braveTracerBridge(brave.Tracer tracer, CurrentTraceContext currentTraceContext,
-				BraveBaggageManager braveBaggageManager) {
-			return new BraveTracer(tracer, new BraveCurrentTraceContext(currentTraceContext), braveBaggageManager);
+		BraveTracer braveTracerBridge(brave.Tracer tracer, CurrentTraceContext currentTraceContext) {
+			return new BraveTracer(tracer, new BraveCurrentTraceContext(currentTraceContext), BRAVE_BAGGAGE_MANAGER);
 		}
 
 		@Bean
 		@ConditionalOnMissingBean
-		BraveBaggageManager braveBaggageManager() {
-			return new BraveBaggageManager();
+		BravePropagator bravePropagator(Tracing tracing) {
+			return new BravePropagator(tracing);
 		}
 
 		@Bean
@@ -166,6 +191,123 @@ BraveHttpClientHandler braveHttpClientHandler(
 			return new BraveHttpClientHandler(httpClientHandler);
 		}
 
+		@Configuration(proxyBeanMethods = false)
+		@ConditionalOnProperty(value = "management.tracing.baggage.enabled", havingValue = "false",
+				matchIfMissing = true)
+		static class BraveNoBaggageConfiguration {
+
+			@Bean
+			@ConditionalOnMissingBean
+			@ConditionalOnProperty(value = "management.tracing.propagation.type", havingValue = "W3C",
+					matchIfMissing = true)
+			Factory w3cPropagationNoBaggageFactory() {
+				return new W3CPropagation(BRAVE_BAGGAGE_MANAGER, List.of()); // TODO: Use
+																				// snapshots
+																				// of
+																				// tracing
+																				// to not
+																				// use
+																				// baggage
+																				// for W3C
+			}
+
+			@Bean
+			@ConditionalOnMissingBean
+			@ConditionalOnProperty(value = "management.tracing.propagation.type", havingValue = "B3")
+			Factory b3PropagationNoBaggageFactory() {
+				return B3Propagation.newFactoryBuilder().injectFormat(B3Propagation.Format.SINGLE_NO_PARENT).build();
+			}
+
+		}
+
+		@Configuration(proxyBeanMethods = false)
+		@ConditionalOnProperty(value = "management.tracing.baggage.enabled", matchIfMissing = true)
+		static class BraveBaggageConfiguration {
+
+			@Bean
+			@ConditionalOnMissingBean
+			@ConditionalOnProperty(value = "management.tracing.propagation.type", havingValue = "W3C",
+					matchIfMissing = true)
+			BaggagePropagation.FactoryBuilder w3cPropagationFactory() {
+				return BaggagePropagation.newFactoryBuilder(new W3CPropagation(BRAVE_BAGGAGE_MANAGER, List.of()));
+			}
+
+			@Bean
+			@ConditionalOnMissingBean
+			@ConditionalOnProperty(value = "management.tracing.propagation.type", havingValue = "B3")
+			BaggagePropagation.FactoryBuilder b3PropagationFactory() {
+				return BaggagePropagation.newFactoryBuilder(
+						B3Propagation.newFactoryBuilder().injectFormat(B3Propagation.Format.SINGLE_NO_PARENT).build());
+			}
+
+			@Bean
+			@ConditionalOnMissingBean
+			Propagation.Factory micrometerTracingPropagationWithBaggage(
+					BaggagePropagation.FactoryBuilder factoryBuilder, TracingProperties tracingProperties,
+					ObjectProvider<List<BaggagePropagationCustomizer>> baggagePropagationCustomizers) {
+				List<String> remoteFields = tracingProperties.getBaggage().getRemoteFields();
+				for (String fieldName : remoteFields) {
+					factoryBuilder
+							.add(BaggagePropagationConfig.SingleBaggageField.remote(BaggageField.create(fieldName)));
+				}
+				baggagePropagationCustomizers.ifAvailable(
+						(customizers) -> customizers.forEach((customizer) -> customizer.customize(factoryBuilder)));
+				return factoryBuilder.build();
+			}
+
+			@Bean
+			@ConditionalOnMissingBean(CorrelationScopeDecorator.class)
+			@ConditionalOnBean(CorrelationScopeDecorator.Builder.class)
+			@ConditionalOnProperty(value = "management.tracing.baggage.correlation-enabled", matchIfMissing = true)
+			ScopeDecorator correlationFieldsCorrelationScopeDecorator(TracingProperties properties,
+					ObjectProvider<List<CorrelationScopeCustomizer>> correlationScopeCustomizers,
+					CorrelationScopeDecorator.Builder builder) {
+				List<String> correlationFields = properties.getBaggage().getCorrelationFields();
+				for (String field : correlationFields) {
+					builder.add(CorrelationScopeConfig.SingleCorrelationField.newBuilder(BaggageField.create(field))
+							.flushOnUpdate().build());
+				}
+				correlationScopeCustomizers.ifAvailable(
+						(customizers) -> customizers.forEach((customizer) -> customizer.customize(builder)));
+				return builder.build();
+			}
+
+			@Bean
+			@ConditionalOnMissingBean(CorrelationScopeDecorator.class)
+			@ConditionalOnBean(CorrelationScopeDecorator.Builder.class)
+			@ConditionalOnProperty(value = "management.tracing.baggage.correlation-enabled", havingValue = "false")
+			ScopeDecorator noCorrelationFieldsCorrelationScopeDecorator(CorrelationScopeDecorator.Builder builder,
+					ObjectProvider<List<CorrelationScopeCustomizer>> correlationScopeCustomizers) {
+				correlationScopeCustomizers.ifAvailable(
+						(customizers) -> customizers.forEach((customizer) -> customizer.customize(builder)));
+				return builder.build();
+			}
+
+		}
+
+	}
+
+	@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)
+		@ConditionalOnMissingClass("io.micrometer.tracing.brave.bridge.BraveTracer")
+		ScopeDecorator correlationScopeDecorator(CorrelationScopeDecorator.Builder builder,
+				ObjectProvider<List<CorrelationScopeCustomizer>> correlationScopeCustomizers) {
+			correlationScopeCustomizers
+					.ifAvailable((customizers) -> customizers.forEach((customizer) -> customizer.customize(builder)));
+			return builder.build();
+		}
+
 	}
 
 }