Hide trace scheduler from spring sleuth (#2158) (#2166)

Co-authored-by: JBodkin <[email protected]>

spring-cloud-gcp-autoconfigure/src/main/java/org/springframework/cloud/gcp/autoconfigure/trace/StackdriverTraceAutoConfiguration.java

@@ -17,8 +17,11 @@
 package org.springframework.cloud.gcp.autoconfigure.trace;
 
 import java.io.IOException;
+import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
+import javax.annotation.PreDestroy;
+
 import brave.http.HttpClientParser;
 import brave.http.HttpServerParser;
 import brave.propagation.Propagation;
@@ -101,6 +104,8 @@ public class StackdriverTraceAutoConfiguration {
 
 	private UserAgentHeaderProvider headerProvider = new UserAgentHeaderProvider(this.getClass());
 
+	private ThreadPoolTaskScheduler defaultTraceSenderThreadPool;
+
 	public StackdriverTraceAutoConfiguration(GcpProjectIdProvider gcpProjectIdProvider,
 			CredentialsProvider credentialsProvider,
 			GcpTraceProperties gcpTraceProperties) throws IOException {
@@ -121,19 +126,21 @@ public SleuthProperties stackdriverSleuthProperties(SleuthProperties sleuthPrope
 		return sleuthProperties;
 	}
 
-	@Bean
-	@ConditionalOnMissingBean(name = "traceSenderThreadPool")
-	public ThreadPoolTaskScheduler traceSenderThreadPool(GcpTraceProperties traceProperties) {
-		ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
-		scheduler.setPoolSize(traceProperties.getNumExecutorThreads());
-		scheduler.setThreadNamePrefix("gcp-trace-sender");
-		scheduler.setDaemon(true);
-		return scheduler;
-	}
-
 	@Bean
 	@ConditionalOnMissingBean(name = "traceExecutorProvider")
-	public ExecutorProvider traceExecutorProvider(@Qualifier("traceSenderThreadPool") ThreadPoolTaskScheduler scheduler) {
+	public ExecutorProvider traceExecutorProvider(GcpTraceProperties traceProperties, @Qualifier("traceSenderThreadPool") Optional<ThreadPoolTaskScheduler> userProvidedScheduler) {
+		ThreadPoolTaskScheduler scheduler;
+		if (userProvidedScheduler.isPresent()) {
+			scheduler = userProvidedScheduler.get();
+		}
+		else {
+			this.defaultTraceSenderThreadPool = new ThreadPoolTaskScheduler();
+			scheduler = this.defaultTraceSenderThreadPool;
+			scheduler.setPoolSize(traceProperties.getNumExecutorThreads());
+			scheduler.setThreadNamePrefix("gcp-trace-sender");
+			scheduler.setDaemon(true);
+			scheduler.initialize();
+		}
 		return FixedExecutorProvider.create(scheduler.getScheduledExecutor());
 	}
 
@@ -218,6 +225,13 @@ public Propagation.Factory stackdriverPropagation() {
 		return StackdriverTracePropagation.FACTORY;
 	}
 
+	@PreDestroy
+	public void closeScheduler() {
+		if (this.defaultTraceSenderThreadPool != null) {
+			this.defaultTraceSenderThreadPool.shutdown();
+		}
+	}
+
 	/**
 	 * Configuration for Sleuth.
 	 */
@@ -242,4 +256,5 @@ HttpServerParser stackdriverHttpServerParser() {
 			return new StackdriverHttpServerParser();
 		}
 	}
+
 }

spring-cloud-gcp-autoconfigure/src/test/java/org/springframework/cloud/gcp/autoconfigure/trace/StackdriverTraceAutoConfigurationTests.java

@@ -21,12 +21,14 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
 import brave.Tracing;
 import brave.http.HttpClientParser;
 import brave.http.HttpServerParser;
 import com.google.api.gax.core.CredentialsProvider;
+import com.google.api.gax.core.ExecutorProvider;
 import com.google.auth.Credentials;
 import com.google.auth.RequestMetadataCallback;
 import com.google.devtools.cloudtrace.v2.BatchWriteSpansRequest;
@@ -57,6 +59,7 @@
 import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
 import org.springframework.cloud.sleuth.log.SleuthLogAutoConfiguration;
 import org.springframework.context.annotation.Bean;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.awaitility.Awaitility.await;
@@ -164,6 +167,29 @@ public void testAsyncReporterHealthCheck() {
 				});
 	}
 
+	@Test
+	public void defaultSchedulerUsedWhenNoneProvided() {
+		this.contextRunner
+				.run(context -> {
+					final ExecutorProvider executorProvider = context.getBean("traceExecutorProvider", ExecutorProvider.class);
+					assertThat(executorProvider.getExecutor()).isNotNull();
+				});
+	}
+
+	@Test
+	public void customSchedulerUsedWhenAvailable() {
+		ThreadPoolTaskScheduler threadPoolTaskSchedulerMock = mock(ThreadPoolTaskScheduler.class);
+		ScheduledExecutorService scheduledExecutorServiceMock = mock(ScheduledExecutorService.class);
+		when(threadPoolTaskSchedulerMock.getScheduledExecutor()).thenReturn(scheduledExecutorServiceMock);
+
+		this.contextRunner
+				.withBean("traceSenderThreadPool", ThreadPoolTaskScheduler.class, () -> threadPoolTaskSchedulerMock)
+				.run(context -> {
+					final ExecutorProvider executorProvider = context.getBean("traceExecutorProvider", ExecutorProvider.class);
+					assertThat(executorProvider.getExecutor()).isEqualTo(scheduledExecutorServiceMock);
+				});
+	}
+
 	/**
 	 * Spring config for tests.
 	 */