Refactor injector and extractor code into CompositeTextMapPropagator

Author: mhalbritter

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

@@ -16,16 +16,22 @@
 
 package org.springframework.boot.actuate.autoconfigure.tracing;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
+import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.context.propagation.TextMapSetter;
+import io.opentelemetry.extension.trace.propagation.B3Propagator;
 
 /**
  * {@link TextMapPropagator} which supports multiple tracing formats. It is able to
@@ -114,13 +120,85 @@ public <C> Context extract(Context context, C carrier, TextMapGetter<C> getter)
 		return currentContext;
 	}
 
+	/**
+	 * Creates a new {@link CompositeTextMapPropagator}, which uses the given
+	 * {@code injectionTypes} for injection and all supported types for extraction.
+	 * @param injectionTypes the propagation types for injection
+	 * @return the {@link CompositeTextMapPropagator}
+	 */
+	static TextMapPropagator create(List<TracingProperties.Propagation.PropagationType> injectionTypes) {
+		return create(null, injectionTypes);
+	}
+
+	/**
+	 * Creates a new {@link CompositeTextMapPropagator}, which uses the given
+	 * {@code injectionTypes} for injection and all supported types for extraction.
+	 * @param baggagePropagator the baggage propagator to use, or {@code null}
+	 * @param injectionTypes the propagation types for injection
+	 * @return the {@link CompositeTextMapPropagator}
+	 */
+	static CompositeTextMapPropagator create(TextMapPropagator baggagePropagator,
+			Iterable<TracingProperties.Propagation.PropagationType> injectionTypes) {
+		List<TextMapPropagator> injectors = stream(injectionTypes)
+			.map((injection) -> forType(injection, baggagePropagator != null))
+			.collect(Collectors.toCollection(ArrayList::new));
+		if (baggagePropagator != null) {
+			injectors.add(baggagePropagator);
+		}
+		List<TextMapPropagator> extractors = Arrays.stream(TracingProperties.Propagation.PropagationType.values())
+			.map((extraction) -> forType(extraction, baggagePropagator != null))
+			.toList();
+		return new CompositeTextMapPropagator(injectors, extractors,
+				(baggagePropagator != null) ? List.of(baggagePropagator) : Collections.emptyList());
+	}
+
 	@SafeVarargs
-	private static <T> Stream<T> concat(Iterable<T>... iterable) {
+	private static <T> Stream<T> concat(Iterable<T>... iterables) {
 		Stream<T> result = Stream.empty();
-		for (Iterable<T> anIterable : iterable) {
-			result = Stream.concat(result, StreamSupport.stream(anIterable.spliterator(), false));
+		for (Iterable<T> iterable : iterables) {
+			result = Stream.concat(result, stream(iterable));
 		}
 		return result;
 	}
 
+	private static <T> Stream<T> stream(Iterable<T> iterable) {
+		return StreamSupport.stream(iterable.spliterator(), false);
+	}
+
+	/**
+	 * Creates a new B3 propagator using a single B3 header.
+	 * @return the B3 propagator
+	 */
+	private static TextMapPropagator b3Single() {
+		return B3Propagator.injectingSingleHeader();
+	}
+
+	/**
+	 * Creates a new B3 propagator using multiple B3 headers.
+	 * @return the B3 propagator
+	 */
+	private static TextMapPropagator b3Multi() {
+		return B3Propagator.injectingMultiHeaders();
+	}
+
+	/**
+	 * Creates a new W3C propagator.
+	 * @param baggage whether baggage propagation should be supported
+	 * @return the W3C propagator
+	 */
+	private static TextMapPropagator w3c(boolean baggage) {
+		if (!baggage) {
+			return W3CTraceContextPropagator.getInstance();
+		}
+		return TextMapPropagator.composite(W3CTraceContextPropagator.getInstance(), W3CBaggagePropagator.getInstance());
+	}
+
+	private static TextMapPropagator forType(TracingProperties.Propagation.PropagationType type, boolean baggage) {
+		return switch (type) {
+			case B3 -> b3Single();
+			case B3_MULTI -> b3Multi();
+			case W3C -> w3c(baggage);
+		};
+	}
+
 }

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

@@ -16,9 +16,6 @@
 
 package org.springframework.boot.actuate.autoconfigure.tracing;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 
@@ -35,14 +32,11 @@
 import io.micrometer.tracing.otel.bridge.Slf4JEventListener;
 import io.micrometer.tracing.otel.propagation.BaggageTextMapPropagator;
 import io.opentelemetry.api.OpenTelemetry;
-import io.opentelemetry.api.baggage.propagation.W3CBaggagePropagator;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Tracer;
-import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.ContextStorage;
 import io.opentelemetry.context.propagation.ContextPropagators;
 import io.opentelemetry.context.propagation.TextMapPropagator;
-import io.opentelemetry.extension.trace.propagation.B3Propagator;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.resources.Resource;
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
@@ -55,7 +49,6 @@
 
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.SpringBootVersion;
-import org.springframework.boot.actuate.autoconfigure.tracing.TracingProperties.Propagation.PropagationType;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -194,11 +187,8 @@ TextMapPropagator textMapPropagatorWithBaggage(OtelCurrentTraceContext otelCurre
 			List<String> remoteFields = this.tracingProperties.getBaggage().getRemoteFields();
 			BaggageTextMapPropagator baggagePropagator = new BaggageTextMapPropagator(remoteFields,
 					new OtelBaggageManager(otelCurrentTraceContext, remoteFields, Collections.emptyList()));
-			List<TextMapPropagator> injectors = new ArrayList<>(
-					TextMapPropagatorFactory.forTypes(this.tracingProperties.getPropagation().getType(), true));
-			injectors.add(baggagePropagator);
-			List<TextMapPropagator> extractors = TextMapPropagatorFactory.forTypes(PropagationType.values(), true);
-			return new CompositeTextMapPropagator(injectors, extractors, List.of(baggagePropagator));
+			return CompositeTextMapPropagator.create(baggagePropagator,
+					this.tracingProperties.getPropagation().getType());
 		}
 
 		@Bean
@@ -217,10 +207,7 @@ static class NoBaggageConfiguration {
 
 		@Bean
 		TextMapPropagator textMapPropagator(TracingProperties properties) {
-			List<TextMapPropagator> injectors = TextMapPropagatorFactory.forTypes(properties.getPropagation().getType(),
-					false);
-			List<TextMapPropagator> extractors = TextMapPropagatorFactory.forTypes(PropagationType.values(), false);
-			return new CompositeTextMapPropagator(injectors, extractors);
+			return CompositeTextMapPropagator.create(properties.getPropagation().getType());
 		}
 
 	}
@@ -242,63 +229,4 @@ public void publishEvent(Object event) {
 
 	}
 
-	/**
-	 * Factory for {@link TextMapPropagator TextMapPropagators}.
-	 */
-	private static final class TextMapPropagatorFactory {
-
-		private TextMapPropagatorFactory() {
-		}
-
-		/**
-		 * Creates a new B3 propagator using a single B3 header.
-		 * @return the B3 propagator
-		 */
-		private static TextMapPropagator b3Single() {
-			return B3Propagator.injectingSingleHeader();
-		}
-
-		/**
-		 * Creates a new B3 propagator using multiple B3 headers.
-		 * @return the B3 propagator
-		 */
-		private static TextMapPropagator b3Multi() {
-			return B3Propagator.injectingMultiHeaders();
-		}
-
-		/**
-		 * Creates a new W3C propagator.
-		 * @param baggage whether baggage propagation should be supported
-		 * @return the W3C propagator
-		 */
-		private static TextMapPropagator w3c(boolean baggage) {
-			if (!baggage) {
-				return W3CTraceContextPropagator.getInstance();
-			}
-			return TextMapPropagator.composite(W3CTraceContextPropagator.getInstance(),
-					W3CBaggagePropagator.getInstance());
-		}
-
-		private static TextMapPropagator forType(PropagationType type, boolean baggage) {
-			return switch (type) {
-				case B3 -> b3Single();
-				case B3_MULTI -> b3Multi();
-				case W3C -> w3c(baggage);
-			};
-		}
-
-		private static List<TextMapPropagator> forTypes(Collection<PropagationType> types, boolean baggage) {
-			List<TextMapPropagator> result = new ArrayList<>(types.size());
-			for (PropagationType type : types) {
-				result.add(forType(type, baggage));
-			}
-			return result;
-		}
-
-		private static List<TextMapPropagator> forTypes(PropagationType[] types, boolean baggage) {
-			return forTypes(Arrays.stream(types).toList(), baggage);
-		}
-
-	}
-
 }