From 4afb5836e83c8ac76a0598e9d20f05fdb37c709a Mon Sep 17 00:00:00 2001
From: Brian Devins-Suresh <badevins@gmail.com>
Date: Fri, 14 Dec 2018 03:33:40 -0500
Subject: [PATCH] Allow customizing the base Propagation.Factory (#1156)

* Add a customizer interface, not a huge fan of the class names

* Allow an optional factory builder to be passed

* Put imports back

* Add autoconfig test

* Faster tests and better coverage

* Better class name
---
 .../autoconfig/TraceAutoConfiguration.java    | 13 ++++-
 ...gurationPropagationCustomizationTests.java | 58 +++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfigurationPropagationCustomizationTests.java

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfiguration.java
index c71a79899..ba37b2e7f 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfiguration.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfiguration.java
@@ -71,6 +71,9 @@ public class TraceAutoConfiguration {
 	@Autowired(required = false)
 	List<CurrentTraceContext.ScopeDecorator> scopeDecorators = new ArrayList<>();
 
+	@Autowired(required = false)
+	ExtraFieldPropagation.FactoryBuilder extraFieldPropagationFactoryBuilder;
+
 	@Bean
 	@ConditionalOnMissingBean
 	// NOTE: stable bean name as might be used outside sleuth
@@ -126,8 +129,14 @@ public class TraceAutoConfiguration {
 				&& sleuthProperties.getPropagationKeys().isEmpty()) {
 			return B3Propagation.FACTORY;
 		}
-		ExtraFieldPropagation.FactoryBuilder factoryBuilder = ExtraFieldPropagation
-				.newFactoryBuilder(B3Propagation.FACTORY);
+		ExtraFieldPropagation.FactoryBuilder factoryBuilder;
+		if (extraFieldPropagationFactoryBuilder != null) {
+			factoryBuilder = extraFieldPropagationFactoryBuilder;
+		}
+		else {
+			factoryBuilder = ExtraFieldPropagation
+					.newFactoryBuilder(B3Propagation.FACTORY);
+		}
 		if (!sleuthProperties.getBaggageKeys().isEmpty()) {
 			factoryBuilder = factoryBuilder
 					// for HTTP
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfigurationPropagationCustomizationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfigurationPropagationCustomizationTests.java
new file mode 100644
index 000000000..590ccb959
--- /dev/null
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/autoconfig/TraceAutoConfigurationPropagationCustomizationTests.java
@@ -0,0 +1,58 @@
+package org.springframework.cloud.sleuth.autoconfig;
+
+import brave.propagation.B3Propagation;
+import brave.propagation.B3SinglePropagation;
+import brave.propagation.ExtraFieldPropagation;
+import brave.propagation.Propagation;
+import org.assertj.core.api.BDDAssertions;
+import org.junit.Test;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+public class TraceAutoConfigurationPropagationCustomizationTests {
+
+	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+			.withConfiguration(AutoConfigurations.of(TraceAutoConfiguration.class));
+
+	@Test
+	public void stillCreatesDefault() {
+		this.contextRunner.run((context) -> {
+					BDDAssertions.then(context.getBean(Propagation.Factory.class))
+							.isEqualTo(B3Propagation.FACTORY);
+				});
+	}
+
+	@Test
+	public void allowsCustomization() {
+		this.contextRunner
+				.withPropertyValues("spring.sleuth.baggage-keys=my-baggage")
+				.run((context) -> {
+					BDDAssertions.then(context.getBean(Propagation.Factory.class))
+							.hasFieldOrPropertyWithValue("delegate", B3Propagation.FACTORY);
+				});
+	}
+
+	@Test
+	public void allowsCustomizationOfBuilder() {
+		this.contextRunner
+				.withPropertyValues("spring.sleuth.baggage-keys=my-baggage")
+				.withUserConfiguration(CustomPropagationFactoryBuilderConfig.class)
+				.run((context) -> {
+					BDDAssertions.then(context.getBean(Propagation.Factory.class))
+							.hasFieldOrPropertyWithValue("delegate", B3SinglePropagation.FACTORY);
+		});
+	}
+
+	@Configuration
+	static class CustomPropagationFactoryBuilderConfig {
+
+		@Bean
+		public ExtraFieldPropagation.FactoryBuilder factoryBuilder() {
+			return ExtraFieldPropagation.newFactoryBuilder(B3SinglePropagation.FACTORY);
+		}
+
+	}
+
+}