From fe7d4f775fee777cc53078fb7d95abcc45bcdb36 Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Tue, 31 May 2022 17:19:29 +0200 Subject: [PATCH] Ensures that executors are removed from cache on shutdown; fixes gh-2171 --- .../async/TraceableExecutorService.java | 16 ++++++++-- .../async/issues/issue410/Issue410Tests.java | 5 +-- ...TraceScheduledThreadPoolExecutorTests.java | 1 + .../TraceThreadPoolTaskSchedulerTests.java | 8 +++-- .../async/TraceableExecutorServiceTests.java | 32 +++++++++++++------ 5 files changed, 45 insertions(+), 17 deletions(-) diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java index cc3ad4aad..6f073f8c4 100644 --- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java +++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java @@ -42,7 +42,7 @@ import org.springframework.cloud.sleuth.internal.ContextUtil; // public as most types in this package were documented for use public class TraceableExecutorService implements ExecutorService { - private static final Map CACHE = new ConcurrentHashMap<>(); + static final Map CACHE = new ConcurrentHashMap<>(); final ExecutorService delegate; @@ -93,12 +93,22 @@ public class TraceableExecutorService implements ExecutorService { @Override public void shutdown() { - this.delegate.shutdown(); + try { + this.delegate.shutdown(); + } + finally { + CACHE.remove(this.delegate); + } } @Override public List shutdownNow() { - return this.delegate.shutdownNow(); + try { + return this.delegate.shutdownNow(); + } + finally { + CACHE.remove(this.delegate); + } } @Override diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/issues/issue410/Issue410Tests.java b/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/issues/issue410/Issue410Tests.java index 0491c44b9..1a51df74b 100644 --- a/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/issues/issue410/Issue410Tests.java +++ b/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/issues/issue410/Issue410Tests.java @@ -20,6 +20,7 @@ import java.util.Date; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.atomic.AtomicReference; @@ -258,8 +259,8 @@ class AppConfig { return new RestTemplate(); } - @Bean("taskScheduler") - public Executor myScheduler() { + @Bean(value = "taskScheduler", destroyMethod = "shutdown") + public ExecutorService myScheduler() { return Executors.newSingleThreadExecutor(); } diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceScheduledThreadPoolExecutorTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceScheduledThreadPoolExecutorTests.java index 989ed5a90..258aaefa6 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceScheduledThreadPoolExecutorTests.java +++ b/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() { diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceThreadPoolTaskSchedulerTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceThreadPoolTaskSchedulerTests.java index 5c2716f76..81f6253a6 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceThreadPoolTaskSchedulerTests.java +++ b/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(); diff --git a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java b/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java index 5327e6cce..08b6f0579 100644 --- a/tests/common/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java +++ b/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>> 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() {