Ensures that executors are removed from cache on shutdown; fixes gh-2171

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceScheduledThreadPoolExecutorTests.java

@@ -64,6 +64,7 @@ public abstract class TraceScheduledThreadPoolExecutorTests implements TestTraci
 	@AfterEach
 	void clear() {
 		this.delegate.shutdown();
+		this.traceThreadPoolTaskExecutor.shutdown();
 	}
 
 	private BeanFactory beanFactory() {

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceThreadPoolTaskSchedulerTests.java

@@ -20,6 +20,7 @@ import java.sql.Date;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.RejectedExecutionHandler;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -66,6 +67,7 @@ public abstract class TraceThreadPoolTaskSchedulerTests implements TestTracingAw
 	@AfterEach
 	void clear() {
 		this.delegate.shutdown();
+		this.traceThreadPoolTaskExecutor.shutdown();
 	}
 
 	private BeanFactory beanFactory() {
@@ -81,7 +83,7 @@ public abstract class TraceThreadPoolTaskSchedulerTests implements TestTracingAw
 		Span span = tracerTest().tracing().tracer().nextSpan().name("foo");
 
 		try (Tracer.SpanInScope ws = tracerTest().tracing().tracer().withSpan(span.start())) {
-			this.traceThreadPoolTaskExecutor.initializeExecutor(new ThreadFactory() {
+			ExecutorService executorService = this.traceThreadPoolTaskExecutor.initializeExecutor(new ThreadFactory() {
 				@Override
 				public Thread newThread(Runnable r) {
 					return new Thread(r);
@@ -91,7 +93,9 @@ public abstract class TraceThreadPoolTaskSchedulerTests implements TestTracingAw
 				public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
 
 				}
-			}).submit(aRunnable(executed, span)).get(1, TimeUnit.SECONDS);
+			});
+			executorService.submit(aRunnable(executed, span)).get(1, TimeUnit.SECONDS);
+			executorService.shutdown();
 		}
 		finally {
 			span.end();

tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java

@@ -28,7 +28,6 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 
-import org.assertj.core.api.BDDAssertions;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -50,6 +49,7 @@ import org.springframework.cloud.sleuth.internal.SleuthContextListenerAccessor;
 import org.springframework.cloud.sleuth.test.TestTracingAwareSupplier;
 
 import static java.util.stream.Collectors.toList;
+import static org.assertj.core.api.BDDAssertions.then;
 
 @ExtendWith(MockitoExtension.class)
 public abstract class TraceableExecutorServiceTests implements TestTracingAwareSupplier {
@@ -61,7 +61,7 @@ public abstract class TraceableExecutorServiceTests implements TestTracingAwareS
 
 	ExecutorService executorService = Executors.newFixedThreadPool(3);
 
-	ExecutorService traceManagerableExecutorService;
+	TraceableExecutorService traceManagerableExecutorService;
 
 	SpanVerifyingRunnable spanVerifyingRunnable = new SpanVerifyingRunnable();
 
@@ -93,9 +93,8 @@ public abstract class TraceableExecutorServiceTests implements TestTracingAwareS
 			span.end();
 		}
 
-		BDDAssertions.then(this.spanVerifyingRunnable.traceIds.stream().distinct().collect(toList())).hasSize(1);
-		BDDAssertions.then(this.spanVerifyingRunnable.spanIds.stream().distinct().collect(toList()))
-				.hasSize(TOTAL_THREADS);
+		then(this.spanVerifyingRunnable.traceIds.stream().distinct().collect(toList())).hasSize(1);
+		then(this.spanVerifyingRunnable.spanIds.stream().distinct().collect(toList())).hasSize(TOTAL_THREADS);
 	}
 
 	@Test
@@ -150,7 +149,7 @@ public abstract class TraceableExecutorServiceTests implements TestTracingAwareS
 	private ArgumentMatcher<Collection<? extends Callable<Object>>> withSpanContinuingTraceCallablesOnly() {
 		return argument -> {
 			try {
-				BDDAssertions.then(argument).flatExtracting(Object::getClass)
+				then(argument).flatExtracting(Object::getClass)
 						.containsOnlyElementsOf(Collections.singletonList(TraceCallable.class));
 			}
 			catch (AssertionError e) {
@@ -180,8 +179,21 @@ public abstract class TraceableExecutorServiceTests implements TestTracingAwareS
 				"calculateTax"));
 		// end::completablefuture[]
 
-		BDDAssertions.then(completableFuture.get()).isEqualTo(1_000_000L);
-		BDDAssertions.then(this.tracer.currentSpan()).isNull();
+		then(completableFuture.get()).isEqualTo(1_000_000L);
+		then(this.tracer.currentSpan()).isNull();
+	}
+
+	@Test
+	public void should_remove_entries_from_cache_when_executor_service_shutsdown() throws Exception {
+		then(TraceableExecutorService.CACHE).doesNotContainKey(executorService);
+
+		TraceableExecutorService.wrap(beanFactory, executorService, "foo").shutdown();
+
+		then(TraceableExecutorService.CACHE).doesNotContainKey(executorService);
+
+		TraceableExecutorService.wrap(beanFactory, executorService, "foo").shutdownNow();
+
+		then(TraceableExecutorService.CACHE).doesNotContainKey(executorService);
 	}
 
 	@Test
@@ -194,8 +206,8 @@ public abstract class TraceableExecutorServiceTests implements TestTracingAwareS
 			return 1_000_000L;
 		}, new TraceableExecutorService(beanFactory, executorService, "calculateTax"));
 
-		BDDAssertions.then(completableFuture.get()).isEqualTo(1_000_000L);
-		BDDAssertions.then(this.tracer.currentSpan()).isNull();
+		then(completableFuture.get()).isEqualTo(1_000_000L);
+		then(this.tracer.currentSpan()).isNull();
 	}
 
 	private CompletableFuture<?>[] runnablesExecutedViaTraceManagerableExecutorService() {