WIP

Author: mhalbritter

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

@@ -16,6 +16,8 @@
 
 package org.springframework.boot.actuate.autoconfigure.tracing;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import brave.CurrentSpanCustomizer;
@@ -34,18 +36,21 @@
 import brave.baggage.CorrelationScopeDecorator;
 import brave.context.slf4j.MDCScopeDecorator;
 import brave.handler.SpanHandler;
+import brave.propagation.B3Propagation;
 import brave.propagation.CurrentTraceContext;
 import brave.propagation.CurrentTraceContext.ScopeDecorator;
 import brave.propagation.CurrentTraceContextCustomizer;
 import brave.propagation.Propagation.Factory;
 import brave.propagation.ThreadLocalCurrentTraceContext;
 import brave.sampler.Sampler;
+import io.micrometer.tracing.BaggageManager;
 import io.micrometer.tracing.brave.bridge.BraveBaggageManager;
 import io.micrometer.tracing.brave.bridge.BraveCurrentTraceContext;
 import io.micrometer.tracing.brave.bridge.BravePropagator;
 import io.micrometer.tracing.brave.bridge.BraveSpanCustomizer;
 import io.micrometer.tracing.brave.bridge.BraveTracer;
 import io.micrometer.tracing.brave.bridge.CompositeSpanHandler;
+import io.micrometer.tracing.brave.bridge.W3CPropagation;
 import io.micrometer.tracing.exporter.SpanExportingPredicate;
 import io.micrometer.tracing.exporter.SpanFilter;
 import io.micrometer.tracing.exporter.SpanReporter;
@@ -76,8 +81,6 @@
 @ConditionalOnEnabledTracing
 public class BraveAutoConfiguration {
 
-	private static final BraveBaggageManager BRAVE_BAGGAGE_MANAGER = new BraveBaggageManager();
-
 	/**
 	 * Default value for application name if {@code spring.application.name} is not set.
 	 */
@@ -92,19 +95,25 @@ CompositeSpanHandler compositeSpanHandler(ObjectProvider<SpanExportingPredicate>
 				filters.orderedStream().toList());
 	}
 
+	@Bean
+	@ConditionalOnMissingBean
+	BraveBaggageManager braveBaggageManager() {
+		return new BraveBaggageManager();
+	}
+
 	@Bean
 	@ConditionalOnMissingBean
 	public Tracing braveTracing(Environment environment, List<SpanHandler> spanHandlers,
 			List<TracingCustomizer> tracingCustomizers, CurrentTraceContext currentTraceContext,
 			Factory propagationFactory, Sampler sampler) {
 		String applicationName = environment.getProperty("spring.application.name", DEFAULT_APPLICATION_NAME);
 		Builder builder = Tracing.newBuilder()
-			.currentTraceContext(currentTraceContext)
-			.traceId128Bit(true)
-			.supportsJoin(false)
-			.propagationFactory(propagationFactory)
-			.sampler(sampler)
-			.localServiceName(applicationName);
+				.currentTraceContext(currentTraceContext)
+				.traceId128Bit(true)
+				.supportsJoin(false)
+				.propagationFactory(propagationFactory)
+				.sampler(sampler)
+				.localServiceName(applicationName);
 		spanHandlers.forEach(builder::addSpanHandler);
 		for (TracingCustomizer tracingCustomizer : tracingCustomizers) {
 			tracingCustomizer.customize(builder);
@@ -138,8 +147,9 @@ public Sampler braveSampler(TracingProperties properties) {
 
 	@Bean
 	@ConditionalOnMissingBean(io.micrometer.tracing.Tracer.class)
-	BraveTracer braveTracerBridge(brave.Tracer tracer, CurrentTraceContext currentTraceContext) {
-		return new BraveTracer(tracer, new BraveCurrentTraceContext(currentTraceContext), BRAVE_BAGGAGE_MANAGER);
+	BraveTracer braveTracerBridge(brave.Tracer tracer, CurrentTraceContext currentTraceContext,
+			BraveBaggageManager braveBaggageManager) {
+		return new BraveTracer(tracer, new BraveCurrentTraceContext(currentTraceContext), braveBaggageManager);
 	}
 
 	@Bean
@@ -166,8 +176,20 @@ static class BraveNoBaggageConfiguration {
 
 		@Bean
 		@ConditionalOnMissingBean
-		Factory propagationFactory(TracingProperties tracing) {
-			return new CompositePropagationFactory(tracing.getPropagation().getType(), null, injectorFactories, extractorFactories);
+		ExtractorFactories extractorFactories() {
+			return ExtractorFactories.of(PropagationFactoryFactory.factoriesFor(PropagationType.values()));
+		}
+
+		@Bean
+		@ConditionalOnMissingBean
+		InjectorFactories injectorFactories(TracingProperties tracing) {
+			return InjectorFactories.of(PropagationFactoryFactory.factoriesFor(tracing.getPropagation().getType()));
+		}
+
+		@Bean
+		@ConditionalOnMissingBean
+		Factory propagationFactory(InjectorFactories injectorFactories, ExtractorFactories extractorFactories) {
+			return new CompositePropagationFactory(injectorFactories, extractorFactories);
 		}
 
 	}
@@ -182,11 +204,26 @@ static class BraveBaggageConfiguration {
 			this.tracingProperties = tracingProperties;
 		}
 
+		@Bean
+		@ConditionalOnMissingBean
+		ExtractorFactories extractorFactories(BraveBaggageManager braveBaggageManager) {
+			return ExtractorFactories
+					.of(PropagationFactoryFactory.factoriesFor(braveBaggageManager, PropagationType.values()));
+		}
+
+		@Bean
+		@ConditionalOnMissingBean
+		InjectorFactories injectorFactories(BraveBaggageManager braveBaggageManager) {
+			return InjectorFactories.of(PropagationFactoryFactory.factoriesFor(braveBaggageManager,
+					this.tracingProperties.getPropagation().getType()));
+		}
+
 		@Bean
 		@ConditionalOnMissingBean
 		BaggagePropagation.FactoryBuilder propagationFactoryBuilder(
-				ObjectProvider<BaggagePropagationCustomizer> baggagePropagationCustomizers) {
-			Factory delegate = new CompositePropagationFactory(this.tracingProperties.getPropagation().getType(), BRAVE_BAGGAGE_MANAGER, injectorFactories, extractorFactories);
+				ObjectProvider<BaggagePropagationCustomizer> baggagePropagationCustomizers,
+				InjectorFactories injectorFactories, ExtractorFactories extractorFactories) {
+			Factory delegate = new CompositePropagationFactory(injectorFactories, extractorFactories);
 			FactoryBuilder builder = BaggagePropagation.newFactoryBuilder(delegate);
 			baggagePropagationCustomizers.orderedStream().forEach((customizer) -> customizer.customize(builder));
 			return builder;
@@ -227,8 +264,8 @@ CorrelationScopeCustomizer correlationFieldsCorrelationScopeCustomizer() {
 				List<String> correlationFields = this.tracingProperties.getBaggage().getCorrelation().getFields();
 				for (String field : correlationFields) {
 					builder.add(CorrelationScopeConfig.SingleCorrelationField.newBuilder(BaggageField.create(field))
-						.flushOnUpdate()
-						.build());
+							.flushOnUpdate()
+							.build());
 				}
 			};
 		}
@@ -241,4 +278,105 @@ ScopeDecorator correlationScopeDecorator(CorrelationScopeDecorator.Builder build
 
 	}
 
+	/**
+	 * A collection of {@link Factory factories} used for trace context extraction.
+	 *
+	 * @author Moritz Halbritter
+	 * @since 3.0.8
+	 */
+	public interface ExtractorFactories extends Iterable<Factory> {
+
+		/**
+		 * Creates an {@link ExtractorFactories} object from the given factories.
+		 * @param factories factories
+		 * @return {@link ExtractorFactories} object
+		 */
+		static ExtractorFactories of(Iterable<Factory> factories) {
+			return factories::iterator;
+		}
+
+		static ExtractorFactories of(PropagationType... types) {
+			return of(PropagationFactoryFactory.factoriesFor(types));
+		}
+
+		static ExtractorFactories of(BaggageManager baggageManager, PropagationType... types) {
+			return of(PropagationFactoryFactory.factoriesFor(baggageManager, types));
+		}
+
+	}
+
+	/**
+	 * A collection of {@link Factory factories} used for trace context injection.
+	 *
+	 * @author Moritz Halbritter
+	 * @since 3.0.8
+	 */
+	public interface InjectorFactories extends Iterable<Factory> {
+
+		/**
+		 * Creates an {@link InjectorFactories} object from the given factories.
+		 * @param factories factories
+		 * @return {@link InjectorFactories} object
+		 */
+		static InjectorFactories of(Iterable<Factory> factories) {
+			return factories::iterator;
+		}
+
+		static InjectorFactories of(PropagationType... types) {
+			return of(PropagationFactoryFactory.factoriesFor(types));
+		}
+
+		static InjectorFactories of(BaggageManager baggageManager, PropagationType... types) {
+			return of(PropagationFactoryFactory.factoriesFor(baggageManager, types));
+		}
+
+	}
+
+	/**
+	 * Factory for {@link Factory propagation factories}.
+	 */
+	private static class PropagationFactoryFactory {
+
+		private PropagationFactoryFactory() {
+		}
+
+		/**
+		 * Creates a new B3 propagation factory.
+		 * @return B3 propagation factory.
+		 */
+		private static Factory b3() {
+			return B3Propagation.newFactoryBuilder().injectFormat(B3Propagation.Format.SINGLE_NO_PARENT).build();
+		}
+
+		/**
+		 * Creates a new W3C propagation factory.
+		 * @param baggageManager baggage manager to use, or {@code null}
+		 * @return W3C Propagation propagation factory
+		 */
+		private static W3CPropagation w3c(BaggageManager baggageManager) {
+			return (baggageManager != null) ? new W3CPropagation(baggageManager, Collections.emptyList())
+					: new W3CPropagation();
+		}
+
+		private static Factory forType(BaggageManager baggageManager, PropagationType type) {
+			return switch (type) {
+				case B3 -> PropagationFactoryFactory.b3();
+				case W3C -> PropagationFactoryFactory.w3c(baggageManager);
+			};
+		}
+
+		private static List<Factory> factoriesFor(PropagationType... types) {
+			return factoriesFor(null, types);
+		}
+
+		private static List<Factory> factoriesFor(BaggageManager baggageManager, PropagationType... types) {
+			List<Factory> result = new ArrayList<>(types.length);
+			for (PropagationType type : types) {
+				result.add(forType(baggageManager, type));
+			}
+			return result;
+		}
+
+	}
+
 }