diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncAutoConfiguration.java index 78bc480b8..029e6b82a 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncAutoConfiguration.java @@ -17,13 +17,12 @@ package org.springframework.cloud.sleuth.instrument.async; import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; -import org.springframework.scheduling.annotation.AsyncConfigurer; /** * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration - * Auto-configuration} that wraps an existing custom {@link AsyncConfigurer} in a - * {@link LazyTraceAsyncCustomizer}. + * Auto-configuration} for asynchronous communication. * * @author Jesus Alonso * @since 2.1.0 @@ -32,4 +31,9 @@ import org.springframework.scheduling.annotation.AsyncConfigurer; @EnableConfigurationProperties(SleuthAsyncProperties.class) public class AsyncAutoConfiguration { + @Bean + ContextRefreshedListener traceContextRefreshedListener() { + return new ContextRefreshedListener(); + } + } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncDefaultAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncDefaultAutoConfiguration.java index 24b4d5bd8..b52bd5b5f 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncDefaultAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/AsyncDefaultAutoConfiguration.java @@ -20,6 +20,7 @@ import java.util.concurrent.Executor; import brave.Tracer; import brave.Tracing; + import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.config.BeanDefinition; diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ContextRefreshedListener.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ContextRefreshedListener.java new file mode 100644 index 000000000..17d68a89a --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ContextRefreshedListener.java @@ -0,0 +1,46 @@ +/* + * Copyright 2013-2018 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.async; + +import java.util.concurrent.atomic.AtomicBoolean; + +import org.springframework.context.ApplicationEvent; +import org.springframework.context.event.ContextRefreshedEvent; +import org.springframework.context.event.SmartApplicationListener; + +class ContextRefreshedListener extends AtomicBoolean implements SmartApplicationListener { + + ContextRefreshedListener(boolean initialValue) { + super(initialValue); + } + + ContextRefreshedListener() { + } + + @Override + public boolean supportsEventType(Class eventType) { + return ContextRefreshedEvent.class.isAssignableFrom(eventType); + } + + @Override + public void onApplicationEvent(ApplicationEvent event) { + if (event instanceof ContextRefreshedEvent) { + set(true); + } + } + +} \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ContextUtil.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ContextUtil.java new file mode 100644 index 000000000..2e0060525 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ContextUtil.java @@ -0,0 +1,54 @@ +/* + * Copyright 2013-2018 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.sleuth.instrument.async; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import org.springframework.beans.factory.BeanFactory; + +/** + * Utility class that verifies that context is in creation. + * + * @author Marcin Grzejszczak + * @since 2.1.0 + */ +class ContextUtil { + + private static final Log log = LogFactory.getLog(ContextUtil.class); + + private static Map CACHE = new ConcurrentHashMap<>(); + + static boolean isContextInCreation(BeanFactory beanFactory) { + ContextRefreshedListener bean = CACHE.compute(beanFactory, + (beanFactory1, contextRefreshedListener) -> { + if (contextRefreshedListener != null) { + return contextRefreshedListener; + } + return beanFactory.getBean(ContextRefreshedListener.class); + }); + boolean contextRefreshed = bean.get(); + if (!contextRefreshed && log.isDebugEnabled()) { + log.debug("Context is not ready yet"); + } + return !contextRefreshed; + } + +} \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java index f8e911b90..d7562c7be 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceExecutor.java @@ -21,13 +21,14 @@ import java.util.concurrent.Executor; import brave.Tracing; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; + import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.NoSuchBeanDefinitionException; import org.springframework.cloud.sleuth.DefaultSpanNamer; import org.springframework.cloud.sleuth.SpanNamer; /** - * {@link Executor} that wraps {@link Runnable} in a trace representation. + * {@link Executor} that wraps {@link Runnable} in a trace representation * * @author Dave Syer * @since 1.0.0 @@ -51,11 +52,15 @@ public class LazyTraceExecutor implements Executor { @Override public void execute(Runnable command) { + if (ContextUtil.isContextInCreation(this.beanFactory)) { + this.delegate.execute(command); + return; + } if (this.tracing == null) { try { this.tracing = this.beanFactory.getBean(Tracing.class); } - catch (NoSuchBeanDefinitionException ex) { + catch (NoSuchBeanDefinitionException e) { this.delegate.execute(command); return; } @@ -69,7 +74,7 @@ public class LazyTraceExecutor implements Executor { try { this.spanNamer = this.beanFactory.getBean(SpanNamer.class); } - catch (NoSuchBeanDefinitionException ex) { + catch (NoSuchBeanDefinitionException e) { log.warn( "SpanNamer bean not found - will provide a manually created instance"); return new DefaultSpanNamer(); diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java index 6a426d182..0cdc52ada 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskExecutor.java @@ -25,6 +25,7 @@ import java.util.concurrent.ThreadPoolExecutor; import brave.Tracing; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; + import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.NoSuchBeanDefinitionException; import org.springframework.cloud.sleuth.DefaultSpanNamer; @@ -34,7 +35,7 @@ import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; import org.springframework.util.concurrent.ListenableFuture; /** - * Trace representation of {@link ThreadPoolTaskExecutor}. + * Trace representation of {@link ThreadPoolTaskExecutor} * * @author Marcin Grzejszczak * @since 1.0.10 @@ -61,35 +62,40 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { @Override public void execute(Runnable task) { - this.delegate.execute(new TraceRunnable(tracing(), spanNamer(), task)); + this.delegate.execute(ContextUtil.isContextInCreation(this.beanFactory) ? task + : new TraceRunnable(tracing(), spanNamer(), task)); } @Override public void execute(Runnable task, long startTimeout) { - this.delegate.execute(new TraceRunnable(tracing(), spanNamer(), task), - startTimeout); + this.delegate.execute(ContextUtil.isContextInCreation(this.beanFactory) ? task + : new TraceRunnable(tracing(), spanNamer(), task), startTimeout); } @Override public Future submit(Runnable task) { - return this.delegate.submit(new TraceRunnable(tracing(), spanNamer(), task)); + return this.delegate.submit(ContextUtil.isContextInCreation(this.beanFactory) + ? task : new TraceRunnable(tracing(), spanNamer(), task)); } @Override public Future submit(Callable task) { - return this.delegate.submit(new TraceCallable<>(tracing(), spanNamer(), task)); + return this.delegate.submit(ContextUtil.isContextInCreation(this.beanFactory) + ? task : new TraceCallable<>(tracing(), spanNamer(), task)); } @Override public ListenableFuture submitListenable(Runnable task) { return this.delegate - .submitListenable(new TraceRunnable(tracing(), spanNamer(), task)); + .submitListenable(ContextUtil.isContextInCreation(this.beanFactory) ? task + : new TraceRunnable(tracing(), spanNamer(), task)); } @Override public ListenableFuture submitListenable(Callable task) { return this.delegate - .submitListenable(new TraceCallable<>(tracing(), spanNamer(), task)); + .submitListenable(ContextUtil.isContextInCreation(this.beanFactory) ? task + : new TraceCallable<>(tracing(), spanNamer(), task)); } @Override @@ -102,6 +108,11 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { this.delegate.setThreadFactory(threadFactory); } + @Override + public void setThreadNamePrefix(String threadNamePrefix) { + this.delegate.setThreadNamePrefix(threadNamePrefix); + } + @Override public void setRejectedExecutionHandler( RejectedExecutionHandler rejectedExecutionHandler) { @@ -174,8 +185,8 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { } @Override - public void setThreadNamePrefix(String threadNamePrefix) { - this.delegate.setThreadNamePrefix(threadNamePrefix); + public void setThreadPriority(int threadPriority) { + this.delegate.setThreadPriority(threadPriority); } @Override @@ -184,8 +195,8 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { } @Override - public void setThreadPriority(int threadPriority) { - this.delegate.setThreadPriority(threadPriority); + public void setDaemon(boolean daemon) { + this.delegate.setDaemon(daemon); } @Override @@ -193,34 +204,24 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { return this.delegate.isDaemon(); } - @Override - public void setDaemon(boolean daemon) { - this.delegate.setDaemon(daemon); - } - @Override public void setThreadGroupName(String name) { this.delegate.setThreadGroupName(name); } - @Override - public ThreadGroup getThreadGroup() { - return this.delegate.getThreadGroup(); - } - @Override public void setThreadGroup(ThreadGroup threadGroup) { this.delegate.setThreadGroup(threadGroup); } @Override - public Thread createThread(Runnable runnable) { - return this.delegate.createThread(runnable); + public ThreadGroup getThreadGroup() { + return this.delegate.getThreadGroup(); } @Override - public int getCorePoolSize() { - return this.delegate.getCorePoolSize(); + public Thread createThread(Runnable runnable) { + return this.delegate.createThread(runnable); } @Override @@ -229,8 +230,8 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { } @Override - public int getMaxPoolSize() { - return this.delegate.getMaxPoolSize(); + public int getCorePoolSize() { + return this.delegate.getCorePoolSize(); } @Override @@ -239,8 +240,8 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { } @Override - public int getKeepAliveSeconds() { - return this.delegate.getKeepAliveSeconds(); + public int getMaxPoolSize() { + return this.delegate.getMaxPoolSize(); } @Override @@ -248,6 +249,11 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { this.delegate.setKeepAliveSeconds(keepAliveSeconds); } + @Override + public int getKeepAliveSeconds() { + return this.delegate.getKeepAliveSeconds(); + } + @Override public void setQueueCapacity(int queueCapacity) { this.delegate.setQueueCapacity(queueCapacity); @@ -275,7 +281,7 @@ public class LazyTraceThreadPoolTaskExecutor extends ThreadPoolTaskExecutor { try { this.spanNamer = this.beanFactory.getBean(SpanNamer.class); } - catch (NoSuchBeanDefinitionException ex) { + catch (NoSuchBeanDefinitionException e) { log.warn( "SpanNamer bean not found - will provide a manually created instance"); return new DefaultSpanNamer(); diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java index eb9e033e5..f646729a7 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorService.java @@ -26,6 +26,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import brave.Tracing; + import org.springframework.beans.factory.BeanFactory; import org.springframework.cloud.sleuth.SpanNamer; @@ -61,9 +62,8 @@ public class TraceableExecutorService implements ExecutorService { @Override public void execute(Runnable command) { - final Runnable r = new TraceRunnable(tracing(), spanNamer(), command, - this.spanName); - this.delegate.execute(r); + this.delegate.submit(ContextUtil.isContextInCreation(this.beanFactory) ? command + : new TraceRunnable(tracing(), spanNamer(), command, this.spanName)); } @Override @@ -94,45 +94,52 @@ public class TraceableExecutorService implements ExecutorService { @Override public Future submit(Callable task) { - Callable c = new TraceCallable<>(tracing(), spanNamer(), task, this.spanName); - return this.delegate.submit(c); + return this.delegate.submit(ContextUtil.isContextInCreation(this.beanFactory) + ? task + : new TraceCallable<>(tracing(), spanNamer(), task, this.spanName)); } @Override public Future submit(Runnable task, T result) { - Runnable r = new TraceRunnable(tracing(), spanNamer(), task, this.spanName); - return this.delegate.submit(r, result); + return this.delegate.submit( + ContextUtil.isContextInCreation(this.beanFactory) ? task + : new TraceRunnable(tracing(), spanNamer(), task, this.spanName), + result); } @Override public Future submit(Runnable task) { - Runnable r = new TraceRunnable(tracing(), spanNamer(), task, this.spanName); - return this.delegate.submit(r); + return this.delegate.submit(ContextUtil.isContextInCreation(this.beanFactory) + ? task : new TraceRunnable(tracing(), spanNamer(), task, this.spanName)); } @Override public List> invokeAll(Collection> tasks) throws InterruptedException { - return this.delegate.invokeAll(wrapCallableCollection(tasks)); + return this.delegate.invokeAll(ContextUtil.isContextInCreation(this.beanFactory) + ? tasks : wrapCallableCollection(tasks)); } @Override public List> invokeAll(Collection> tasks, long timeout, TimeUnit unit) throws InterruptedException { - return this.delegate.invokeAll(wrapCallableCollection(tasks), timeout, unit); + return this.delegate.invokeAll(ContextUtil.isContextInCreation(this.beanFactory) + ? tasks : wrapCallableCollection(tasks), timeout, unit); } @Override public T invokeAny(Collection> tasks) throws InterruptedException, ExecutionException { - return this.delegate.invokeAny(wrapCallableCollection(tasks)); + return this.delegate.invokeAny(ContextUtil.isContextInCreation(this.beanFactory) + ? tasks : wrapCallableCollection(tasks)); } @Override public T invokeAny(Collection> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { - return this.delegate.invokeAny(wrapCallableCollection(tasks), timeout, unit); + return this.delegate.invokeAny(ContextUtil.isContextInCreation(this.beanFactory) + ? tasks : wrapCallableCollection(tasks), timeout, unit); } private Collection> wrapCallableCollection( diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java index 010a64859..a0e015ddf 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorService.java @@ -44,31 +44,37 @@ public class TraceableScheduledExecutorService extends TraceableExecutorService @Override public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { - Runnable r = new TraceRunnable(tracing(), spanNamer(), command); - return getScheduledExecutorService().schedule(r, delay, unit); + return getScheduledExecutorService().schedule( + ContextUtil.isContextInCreation(this.beanFactory) ? command + : new TraceRunnable(tracing(), spanNamer(), command), + delay, unit); } @Override public ScheduledFuture schedule(Callable callable, long delay, TimeUnit unit) { - Callable c = new TraceCallable<>(tracing(), spanNamer(), callable); - return getScheduledExecutorService().schedule(c, delay, unit); + return getScheduledExecutorService().schedule( + ContextUtil.isContextInCreation(this.beanFactory) ? callable + : new TraceCallable<>(tracing(), spanNamer(), callable), + delay, unit); } @Override public ScheduledFuture scheduleAtFixedRate(Runnable command, long initialDelay, long period, TimeUnit unit) { - Runnable r = new TraceRunnable(tracing(), spanNamer(), command); - return getScheduledExecutorService().scheduleAtFixedRate(r, initialDelay, period, - unit); + return getScheduledExecutorService().scheduleAtFixedRate( + ContextUtil.isContextInCreation(this.beanFactory) ? command + : new TraceRunnable(tracing(), spanNamer(), command), + initialDelay, period, unit); } @Override public ScheduledFuture scheduleWithFixedDelay(Runnable command, long initialDelay, long delay, TimeUnit unit) { - Runnable r = new TraceRunnable(tracing(), spanNamer(), command); - return getScheduledExecutorService().scheduleWithFixedDelay(r, initialDelay, - delay, unit); + return getScheduledExecutorService().scheduleWithFixedDelay( + ContextUtil.isContextInCreation(this.beanFactory) ? command + : new TraceRunnable(tracing(), spanNamer(), command), + initialDelay, delay, unit); } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessorTests.java index cc7eb3d01..921fc497e 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessorTests.java @@ -30,7 +30,6 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; -import brave.Tracer; import brave.Tracing; import org.aopalliance.aop.Advice; import org.junit.After; @@ -41,14 +40,13 @@ import org.mockito.BDDMockito; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.junit.MockitoJUnitRunner; + import org.springframework.aop.framework.AopConfigException; import org.springframework.aop.framework.ProxyFactoryBean; import org.springframework.beans.factory.BeanFactory; import org.springframework.cloud.sleuth.DefaultSpanNamer; -import org.springframework.cloud.sleuth.SpanName; import org.springframework.cloud.sleuth.SpanNamer; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; -import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.util.ClassUtils; import static org.assertj.core.api.BDDAssertions.then; @@ -76,6 +74,8 @@ public class ExecutorBeanPostProcessorTests { Mockito.when(this.beanFactory.getBean(Tracing.class)).thenReturn(this.tracing); Mockito.when(this.beanFactory.getBean(SpanNamer.class)) .thenReturn(new DefaultSpanNamer()); + Mockito.when(this.beanFactory.getBean(ContextRefreshedListener.class)) + .thenReturn(new ContextRefreshedListener(true)); } @After diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java index b45d84a8f..bce07ed75 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableExecutorServiceTests.java @@ -77,8 +77,8 @@ public class TraceableExecutorServiceTests { @Before public void setup() { - this.traceManagerableExecutorService = new TraceableExecutorService(beanFactory(), - this.executorService); + this.traceManagerableExecutorService = new TraceableExecutorService( + beanFactory(true), this.executorService); this.reporter.clear(); this.spanVerifyingRunnable.clear(); } @@ -116,7 +116,7 @@ public class TraceableExecutorServiceTests { throws Exception { ExecutorService executorService = Mockito.mock(ExecutorService.class); TraceableExecutorService traceExecutorService = new TraceableExecutorService( - beanFactory(), executorService); + beanFactory(true), executorService); traceExecutorService.invokeAll(callables()); BDDMockito.then(executorService).should() @@ -137,6 +137,33 @@ public class TraceableExecutorServiceTests { BDDMockito.eq(1L), BDDMockito.eq(TimeUnit.DAYS)); } + @Test + @SuppressWarnings("unchecked") + public void should_not_wrap_methods_in_trace_representation_only_for_non_tracing_callables_when_context_not_ready() + throws Exception { + ExecutorService executorService = Mockito.mock(ExecutorService.class); + TraceableExecutorService traceExecutorService = new TraceableExecutorService( + beanFactory(false), executorService); + + traceExecutorService.invokeAll(callables()); + BDDMockito.then(executorService).should(BDDMockito.never()) + .invokeAll(BDDMockito.argThat(withSpanContinuingTraceCallablesOnly())); + + traceExecutorService.invokeAll(callables(), 1L, TimeUnit.DAYS); + BDDMockito.then(executorService).should(BDDMockito.never()).invokeAll( + BDDMockito.argThat(withSpanContinuingTraceCallablesOnly()), + BDDMockito.eq(1L), BDDMockito.eq(TimeUnit.DAYS)); + + traceExecutorService.invokeAny(callables()); + BDDMockito.then(executorService).should(BDDMockito.never()) + .invokeAny(BDDMockito.argThat(withSpanContinuingTraceCallablesOnly())); + + traceExecutorService.invokeAny(callables(), 1L, TimeUnit.DAYS); + BDDMockito.then(executorService).should(BDDMockito.never()).invokeAny( + BDDMockito.argThat(withSpanContinuingTraceCallablesOnly()), + BDDMockito.eq(1L), BDDMockito.eq(TimeUnit.DAYS)); + } + private ArgumentMatcher>> withSpanContinuingTraceCallablesOnly() { return argument -> { try { @@ -162,7 +189,7 @@ public class TraceableExecutorServiceTests { public void should_propagate_trace_info_when_compleable_future_is_used() throws Exception { ExecutorService executorService = this.executorService; - BeanFactory beanFactory = beanFactory(); + BeanFactory beanFactory = beanFactory(true); // tag::completablefuture[] CompletableFuture completableFuture = CompletableFuture.supplyAsync(() -> { // perform some logic @@ -176,6 +203,20 @@ public class TraceableExecutorServiceTests { then(this.tracer.currentSpan()).isNull(); } + @Test + public void should_not_propagate_trace_info_when_compleable_future_is_used_when_context_not_refreshed() + throws Exception { + ExecutorService executorService = this.executorService; + BeanFactory beanFactory = beanFactory(false); + CompletableFuture completableFuture = CompletableFuture.supplyAsync(() -> { + // perform some logic + return 1_000_000L; + }, new TraceableExecutorService(beanFactory, executorService, "calculateTax")); + + then(completableFuture.get()).isEqualTo(1_000_000L); + then(this.tracer.currentSpan()).isNull(); + } + private CompletableFuture[] runnablesExecutedViaTraceManagerableExecutorService() { List> futures = new ArrayList<>(); for (int i = 0; i < TOTAL_THREADS; i++) { @@ -185,11 +226,13 @@ public class TraceableExecutorServiceTests { return futures.toArray(new CompletableFuture[futures.size()]); } - BeanFactory beanFactory() { + BeanFactory beanFactory(boolean refreshed) { BDDMockito.given(this.beanFactory.getBean(Tracing.class)) .willReturn(this.tracing); BDDMockito.given(this.beanFactory.getBean(SpanNamer.class)) .willReturn(new DefaultSpanNamer()); + BDDMockito.given(this.beanFactory.getBean(ContextRefreshedListener.class)) + .willReturn(new ContextRefreshedListener(refreshed)); return this.beanFactory; } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorServiceTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorServiceTest.java index ce5cd8f02..2c84df2c0 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorServiceTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/TraceableScheduledExecutorServiceTest.java @@ -39,6 +39,7 @@ import org.springframework.cloud.sleuth.SpanNamer; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.BDDMockito.then; +import static org.mockito.Mockito.never; /** * @author Marcin Grzejszczak @@ -107,6 +108,60 @@ public class TraceableScheduledExecutorServiceTest { anyLong(), anyLong(), any(TimeUnit.class)); } + @Test + public void should_not_schedule_a_trace_runnable_when_context_not_ready() + throws Exception { + BDDMockito.given(this.beanFactory.getBean(ContextRefreshedListener.class)) + .willReturn(new ContextRefreshedListener(false)); + this.traceableScheduledExecutorService.schedule(aRunnable(), 1L, TimeUnit.DAYS); + + then(this.scheduledExecutorService).should(never()).schedule( + BDDMockito.argThat( + matcher(Runnable.class, instanceOf(TraceRunnable.class))), + anyLong(), any(TimeUnit.class)); + } + + @Test + public void should_not_schedule_a_trace_callable_when_context_not_ready() + throws Exception { + BDDMockito.given(this.beanFactory.getBean(ContextRefreshedListener.class)) + .willReturn(new ContextRefreshedListener(false)); + this.traceableScheduledExecutorService.schedule(aCallable(), 1L, TimeUnit.DAYS); + + then(this.scheduledExecutorService).should(never()).schedule( + BDDMockito.argThat( + matcher(Callable.class, instanceOf(TraceCallable.class))), + anyLong(), any(TimeUnit.class)); + } + + @Test + public void should_not_schedule_at_fixed_rate_a_trace_runnable_when_context_not_ready() + throws Exception { + BDDMockito.given(this.beanFactory.getBean(ContextRefreshedListener.class)) + .willReturn(new ContextRefreshedListener(false)); + this.traceableScheduledExecutorService.scheduleAtFixedRate(aRunnable(), 1L, 1L, + TimeUnit.DAYS); + + then(this.scheduledExecutorService).should(never()).scheduleAtFixedRate( + BDDMockito.argThat( + matcher(Runnable.class, instanceOf(TraceRunnable.class))), + anyLong(), anyLong(), any(TimeUnit.class)); + } + + @Test + public void should_not_schedule_with_fixed_delay_a_trace_runnable_when_context_not_ready() + throws Exception { + BDDMockito.given(this.beanFactory.getBean(ContextRefreshedListener.class)) + .willReturn(new ContextRefreshedListener(false)); + this.traceableScheduledExecutorService.scheduleWithFixedDelay(aRunnable(), 1L, 1L, + TimeUnit.DAYS); + + then(this.scheduledExecutorService).should(never()).scheduleWithFixedDelay( + BDDMockito.argThat( + matcher(Runnable.class, instanceOf(TraceRunnable.class))), + anyLong(), anyLong(), any(TimeUnit.class)); + } + Predicate instanceOf(Class clazz) { return (argument) -> argument.getClass().isAssignableFrom(clazz); } @@ -129,6 +184,8 @@ public class TraceableScheduledExecutorServiceTest { .willReturn(this.tracing); BDDMockito.given(this.beanFactory.getBean(SpanNamer.class)) .willReturn(new DefaultSpanNamer()); + BDDMockito.given(this.beanFactory.getBean(ContextRefreshedListener.class)) + .willReturn(new ContextRefreshedListener(true)); return this.beanFactory; } diff --git a/spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/ZipkinAutoConfigurationTests.java b/spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/ZipkinAutoConfigurationTests.java index 622a18f91..0d30d2a27 100644 --- a/spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/ZipkinAutoConfigurationTests.java +++ b/spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/ZipkinAutoConfigurationTests.java @@ -69,7 +69,8 @@ public class ZipkinAutoConfigurationTests { @Test public void defaultsToV2Endpoint() throws Exception { this.context = new AnnotationConfigApplicationContext(); - environment().setProperty("spring.zipkin.base-url", this.server.url("/").toString()); + environment().setProperty("spring.zipkin.base-url", + this.server.url("/").toString()); this.context.register(ZipkinAutoConfiguration.class, PropertyPlaceholderAutoConfiguration.class, TraceAutoConfiguration.class, Config.class); @@ -79,8 +80,8 @@ public class ZipkinAutoConfigurationTests { span.finish(); - Awaitility.await() - .untilAsserted(() -> then(this.server.getRequestCount()).isGreaterThan(0)); + Awaitility.await().untilAsserted( + () -> then(this.server.getRequestCount()).isGreaterThan(0)); RecordedRequest request = this.server.takeRequest(); then(request.getPath()).isEqualTo("/api/v2/spans"); then(request.getBody().readUtf8()).contains("localEndpoint"); @@ -94,7 +95,8 @@ public class ZipkinAutoConfigurationTests { @Test public void encoderDirectsEndpoint() throws Exception { this.context = new AnnotationConfigApplicationContext(); - environment().setProperty("spring.zipkin.base-url", this.server.url("/").toString()); + environment().setProperty("spring.zipkin.base-url", + this.server.url("/").toString()); environment().setProperty("spring.zipkin.encoder", "JSON_V1"); this.context.register(ZipkinAutoConfiguration.class, PropertyPlaceholderAutoConfiguration.class, TraceAutoConfiguration.class, @@ -105,8 +107,8 @@ public class ZipkinAutoConfigurationTests { span.finish(); - Awaitility.await() - .untilAsserted(() -> then(this.server.getRequestCount()).isGreaterThan(0)); + Awaitility.await().untilAsserted( + () -> then(this.server.getRequestCount()).isGreaterThan(0)); RecordedRequest request = this.server.takeRequest(); then(request.getPath()).isEqualTo("/api/v1/spans"); then(request.getBody().readUtf8()).contains("binaryAnnotations");