From 7f0d4f1dc34d9730d9e9bdbe083e1bf12c4d622a Mon Sep 17 00:00:00 2001 From: Marcin Grzejszczak Date: Wed, 10 Apr 2019 15:02:14 +0200 Subject: [PATCH] Added a fix for special cases of executors that fail to be wrapped in a cglib proxy fixes gh-1232 --- .../async/ExecutorBeanPostProcessor.java | 82 ++- .../LazyTraceScheduledThreadPoolExecutor.java | 570 ++++++++++++++++++ .../LazyTraceThreadPoolTaskScheduler.java | 439 ++++++++++++++ .../scheduling/TraceSchedulingAspect.java | 3 + ...LazyTraceThreadPoolTaskSchedulerTests.java | 462 ++++++++++++++ .../async/issues/issue410/Issue410Tests.java | 106 ++++ .../scheduling/TracingOnScheduledTests.java | 1 + 7 files changed, 1652 insertions(+), 11 deletions(-) create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java create mode 100644 spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java index 374cebb7e..443a82041 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java @@ -21,6 +21,7 @@ import java.lang.reflect.Method; import java.lang.reflect.Modifier; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; +import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.function.Supplier; import org.aopalliance.aop.Advice; @@ -36,6 +37,7 @@ import org.springframework.beans.factory.BeanFactory; import org.springframework.beans.factory.config.BeanPostProcessor; import org.springframework.core.task.AsyncTaskExecutor; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; +import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.util.ReflectionUtils; /** @@ -133,6 +135,22 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { return createThreadPoolTaskExecutorProxy(bean, cglibProxy, executor); } + private ProxyFactoryBean wrapThreadPoolTaskScheduler(Object bean) { + boolean classFinal = Modifier.isFinal(bean.getClass().getModifiers()); + boolean cglibProxy = !classFinal; + ThreadPoolTaskScheduler executor = (ThreadPoolTaskScheduler) bean; + return proxyFactoryBean(bean, cglibProxy, executor, + createThreadPoolTaskSchedulerProxy(executor)); + } + + private ProxyFactoryBean wrapScheduledThreadPoolExecutor(Object bean) { + boolean classFinal = Modifier.isFinal(bean.getClass().getModifiers()); + boolean cglibProxy = !classFinal; + ScheduledThreadPoolExecutor executor = (ScheduledThreadPoolExecutor) bean; + return proxyFactoryBean(bean, cglibProxy, executor, + createScheduledThreadPoolExecutorProxy(executor)); + } + private Object wrapExecutorService(Object bean) { boolean classFinal = Modifier.isFinal(bean.getClass().getModifiers()); boolean cglibProxy = !classFinal; @@ -158,6 +176,17 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { () -> new LazyTraceThreadPoolTaskExecutor(this.beanFactory, executor)); } + Supplier createThreadPoolTaskSchedulerProxy(ThreadPoolTaskScheduler executor) { + return () -> new LazyTraceThreadPoolTaskScheduler(this.beanFactory, executor); + } + + Supplier createScheduledThreadPoolExecutorProxy( + ScheduledThreadPoolExecutor executor) { + return () -> new LazyTraceScheduledThreadPoolExecutor(executor.getCorePoolSize(), + executor.getThreadFactory(), executor.getRejectedExecutionHandler(), + this.beanFactory, executor); + } + Object createExecutorServiceProxy(Object bean, boolean cglibProxy, ExecutorService executor) { return getProxiedObject(bean, cglibProxy, executor, @@ -172,6 +201,47 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { private Object getProxiedObject(Object bean, boolean cglibProxy, Executor executor, Supplier supplier) { + ProxyFactoryBean factory = proxyFactoryBean(bean, cglibProxy, executor, supplier); + try { + return getObject(factory); + } + catch (Exception ex) { + if (log.isDebugEnabled()) { + log.debug( + "Exception occurred while trying to get a proxy. Will fallback to a different implementation", + ex); + } + try { + if (bean instanceof ThreadPoolTaskScheduler) { + if (log.isDebugEnabled()) { + log.debug( + "Will wrap ThreadPoolTaskScheduler in its tracing representation due to previous errors"); + } + return createThreadPoolTaskSchedulerProxy( + (ThreadPoolTaskScheduler) bean).get(); + } + else if (bean instanceof ScheduledThreadPoolExecutor) { + if (log.isDebugEnabled()) { + log.debug( + "Will wrap ScheduledThreadPoolExecutor in its tracing representation due to previous errors"); + } + return createScheduledThreadPoolExecutorProxy( + (ScheduledThreadPoolExecutor) bean).get(); + } + } + catch (Exception ex2) { + if (log.isDebugEnabled()) { + log.debug( + "Fallback for special wrappers failed, will try the tracing representation instead", + ex2); + } + } + return supplier.get(); + } + } + + private ProxyFactoryBean proxyFactoryBean(Object bean, boolean cglibProxy, + Executor executor, Supplier supplier) { ProxyFactoryBean factory = new ProxyFactoryBean(); factory.setProxyTargetClass(cglibProxy); factory.addAdvice( @@ -182,17 +252,7 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor { } }); factory.setTarget(bean); - try { - return getObject(factory); - } - catch (Exception e) { - if (log.isDebugEnabled()) { - log.debug( - "Exception occurred while trying to get a proxy. Will fallback to a different implementation", - e); - } - return supplier.get(); - } + return factory; } Object getObject(ProxyFactoryBean factory) { diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java new file mode 100644 index 000000000..ce29e6eb0 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceScheduledThreadPoolExecutor.java @@ -0,0 +1,570 @@ +/* + * 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 + * + * https://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.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionHandler; +import java.util.concurrent.RunnableFuture; +import java.util.concurrent.RunnableScheduledFuture; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +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; +import org.springframework.util.ReflectionUtils; + +/** + * Trace representation of {@link ScheduledThreadPoolExecutor}. Should be used only * as + * last resort, when any other approaches fail. + * + * @author Marcin Grzejszczak + * @since 2.0.4 + */ +// TODO: Think of a better solution than this +class LazyTraceScheduledThreadPoolExecutor extends ScheduledThreadPoolExecutor { + + private static final Log log = LogFactory + .getLog(LazyTraceScheduledThreadPoolExecutor.class); + + private final BeanFactory beanFactory; + + private final ScheduledThreadPoolExecutor delegate; + + private final Method decorateTaskRunnable; + + private final Method decorateTaskCallable; + + private final Method finalize; + + private final Method beforeExecute; + + private final Method afterExecute; + + private final Method terminated; + + private final Method newTaskForRunnable; + + private final Method newTaskForCallable; + + private Tracing tracing; + + private SpanNamer spanNamer; + + LazyTraceScheduledThreadPoolExecutor(int corePoolSize, BeanFactory beanFactory, + ScheduledThreadPoolExecutor delegate) { + super(corePoolSize); + this.beanFactory = beanFactory; + this.delegate = delegate; + this.decorateTaskRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTask", Runnable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskRunnable); + this.decorateTaskCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTaskCallable", Callable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskCallable); + this.finalize = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "finalize", null); + makeAccessibleIfNotNull(this.finalize); + this.beforeExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "beforeExecute", null); + makeAccessibleIfNotNull(this.beforeExecute); + this.afterExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "afterExecute", null); + makeAccessibleIfNotNull(this.afterExecute); + this.terminated = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "terminated", null); + makeAccessibleIfNotNull(this.terminated); + this.newTaskForRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Runnable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForRunnable); + this.newTaskForCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Callable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForCallable); + } + + private void makeAccessibleIfNotNull(Method method) { + if (method != null) { + ReflectionUtils.makeAccessible(method); + } + } + + LazyTraceScheduledThreadPoolExecutor(int corePoolSize, ThreadFactory threadFactory, + BeanFactory beanFactory, ScheduledThreadPoolExecutor delegate) { + super(corePoolSize, threadFactory); + this.beanFactory = beanFactory; + this.delegate = delegate; + this.decorateTaskRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTask", Runnable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskRunnable); + this.decorateTaskCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTaskCallable", Callable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskCallable); + this.finalize = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "finalize"); + makeAccessibleIfNotNull(this.finalize); + this.beforeExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "beforeExecute"); + makeAccessibleIfNotNull(this.beforeExecute); + this.afterExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "afterExecute", null); + makeAccessibleIfNotNull(this.afterExecute); + this.terminated = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "terminated", null); + makeAccessibleIfNotNull(this.terminated); + this.newTaskForRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Runnable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForRunnable); + this.newTaskForCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Callable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForCallable); + } + + LazyTraceScheduledThreadPoolExecutor(int corePoolSize, + RejectedExecutionHandler handler, BeanFactory beanFactory, + ScheduledThreadPoolExecutor delegate) { + super(corePoolSize, handler); + this.beanFactory = beanFactory; + this.delegate = delegate; + this.decorateTaskRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTask", Runnable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskRunnable); + this.decorateTaskCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTaskCallable", Callable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskCallable); + this.finalize = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "finalize", null); + makeAccessibleIfNotNull(this.finalize); + this.beforeExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "beforeExecute", null); + makeAccessibleIfNotNull(this.beforeExecute); + this.afterExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "afterExecute", null); + makeAccessibleIfNotNull(this.afterExecute); + this.terminated = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "terminated", null); + makeAccessibleIfNotNull(this.terminated); + this.newTaskForRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Runnable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForRunnable); + this.newTaskForCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Callable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForCallable); + } + + LazyTraceScheduledThreadPoolExecutor(int corePoolSize, ThreadFactory threadFactory, + RejectedExecutionHandler handler, BeanFactory beanFactory, + ScheduledThreadPoolExecutor delegate) { + super(corePoolSize, threadFactory, handler); + this.beanFactory = beanFactory; + this.delegate = delegate; + this.decorateTaskRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTask", Runnable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskRunnable); + this.decorateTaskCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "decorateTaskCallable", Callable.class, + RunnableScheduledFuture.class); + makeAccessibleIfNotNull(this.decorateTaskCallable); + this.finalize = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "finalize", null); + makeAccessibleIfNotNull(this.finalize); + this.beforeExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "beforeExecute", null); + makeAccessibleIfNotNull(this.beforeExecute); + this.afterExecute = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "afterExecute", null); + makeAccessibleIfNotNull(this.afterExecute); + this.terminated = ReflectionUtils.findMethod(ScheduledThreadPoolExecutor.class, + "terminated"); + makeAccessibleIfNotNull(this.terminated); + this.newTaskForRunnable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Runnable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForRunnable); + this.newTaskForCallable = ReflectionUtils.findMethod( + ScheduledThreadPoolExecutor.class, "newTaskFor", Callable.class, + Object.class); + makeAccessibleIfNotNull(this.newTaskForCallable); + } + + @Override + @SuppressWarnings("unchecked") + public RunnableScheduledFuture decorateTask(Runnable runnable, + RunnableScheduledFuture task) { + return (RunnableScheduledFuture) ReflectionUtils.invokeMethod( + this.decorateTaskRunnable, this.delegate, + new TraceRunnable(tracing(), spanNamer(), runnable), task); + } + + @Override + @SuppressWarnings("unchecked") + public RunnableScheduledFuture decorateTask(Callable callable, + RunnableScheduledFuture task) { + return (RunnableScheduledFuture) ReflectionUtils.invokeMethod( + this.decorateTaskCallable, this.delegate, + new TraceCallable<>(tracing(), spanNamer(), callable), task); + } + + @Override + public ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { + return this.delegate.schedule(new TraceRunnable(tracing(), spanNamer(), command), + delay, unit); + } + + @Override + public ScheduledFuture schedule(Callable callable, long delay, + TimeUnit unit) { + return this.delegate.schedule( + new TraceCallable<>(tracing(), spanNamer(), callable), delay, unit); + } + + @Override + public ScheduledFuture scheduleAtFixedRate(Runnable command, long initialDelay, + long period, TimeUnit unit) { + return this.delegate.scheduleAtFixedRate( + new TraceRunnable(tracing(), spanNamer(), command), initialDelay, period, + unit); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay(Runnable command, long initialDelay, + long delay, TimeUnit unit) { + return this.delegate.scheduleWithFixedDelay( + new TraceRunnable(tracing(), spanNamer(), command), initialDelay, delay, + unit); + } + + @Override + public void execute(Runnable command) { + this.delegate.execute(new TraceRunnable(tracing(), spanNamer(), command)); + } + + @Override + public Future submit(Runnable task) { + return this.delegate.submit(new TraceRunnable(tracing(), spanNamer(), task)); + } + + @Override + public Future submit(Runnable task, T result) { + return this.delegate.submit(new TraceRunnable(tracing(), spanNamer(), task), + result); + } + + @Override + public Future submit(Callable task) { + return this.delegate.submit(new TraceCallable<>(tracing(), spanNamer(), task)); + } + + @Override + public void setContinueExistingPeriodicTasksAfterShutdownPolicy(boolean value) { + this.delegate.setContinueExistingPeriodicTasksAfterShutdownPolicy(value); + } + + @Override + public boolean getContinueExistingPeriodicTasksAfterShutdownPolicy() { + return this.delegate.getContinueExistingPeriodicTasksAfterShutdownPolicy(); + } + + @Override + public void setExecuteExistingDelayedTasksAfterShutdownPolicy(boolean value) { + this.delegate.setExecuteExistingDelayedTasksAfterShutdownPolicy(value); + } + + @Override + public boolean getExecuteExistingDelayedTasksAfterShutdownPolicy() { + return this.delegate.getExecuteExistingDelayedTasksAfterShutdownPolicy(); + } + + @Override + public void setRemoveOnCancelPolicy(boolean value) { + this.delegate.setRemoveOnCancelPolicy(value); + } + + @Override + public boolean getRemoveOnCancelPolicy() { + return this.delegate.getRemoveOnCancelPolicy(); + } + + @Override + public void shutdown() { + this.delegate.shutdown(); + } + + @Override + public List shutdownNow() { + return this.delegate.shutdownNow(); + } + + @Override + public BlockingQueue getQueue() { + return this.delegate.getQueue(); + } + + @Override + public boolean isShutdown() { + return this.delegate.isShutdown(); + } + + @Override + public boolean isTerminating() { + return this.delegate.isTerminating(); + } + + @Override + public boolean isTerminated() { + return this.delegate.isTerminated(); + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) + throws InterruptedException { + return this.delegate.awaitTermination(timeout, unit); + } + + @Override + @Deprecated + public void finalize() { + ReflectionUtils.invokeMethod(this.finalize, this.delegate); + } + + @Override + public void setThreadFactory(ThreadFactory threadFactory) { + this.delegate.setThreadFactory(threadFactory); + } + + @Override + public ThreadFactory getThreadFactory() { + return this.delegate.getThreadFactory(); + } + + @Override + public void setRejectedExecutionHandler(RejectedExecutionHandler handler) { + this.delegate.setRejectedExecutionHandler(handler); + } + + @Override + public RejectedExecutionHandler getRejectedExecutionHandler() { + return this.delegate.getRejectedExecutionHandler(); + } + + @Override + public void setCorePoolSize(int corePoolSize) { + this.delegate.setCorePoolSize(corePoolSize); + } + + @Override + public int getCorePoolSize() { + return this.delegate.getCorePoolSize(); + } + + @Override + public boolean prestartCoreThread() { + return this.delegate.prestartCoreThread(); + } + + @Override + public int prestartAllCoreThreads() { + return this.delegate.prestartAllCoreThreads(); + } + + @Override + public boolean allowsCoreThreadTimeOut() { + return this.delegate.allowsCoreThreadTimeOut(); + } + + @Override + public void allowCoreThreadTimeOut(boolean value) { + this.delegate.allowCoreThreadTimeOut(value); + } + + @Override + public void setMaximumPoolSize(int maximumPoolSize) { + this.delegate.setMaximumPoolSize(maximumPoolSize); + } + + @Override + public int getMaximumPoolSize() { + return this.delegate.getMaximumPoolSize(); + } + + @Override + public void setKeepAliveTime(long time, TimeUnit unit) { + this.delegate.setKeepAliveTime(time, unit); + } + + @Override + public long getKeepAliveTime(TimeUnit unit) { + return this.delegate.getKeepAliveTime(unit); + } + + @Override + public boolean remove(Runnable task) { + return this.delegate.remove(task); + } + + @Override + public void purge() { + this.delegate.purge(); + } + + @Override + public int getPoolSize() { + return this.delegate.getPoolSize(); + } + + @Override + public int getActiveCount() { + return this.delegate.getActiveCount(); + } + + @Override + public int getLargestPoolSize() { + return this.delegate.getLargestPoolSize(); + } + + @Override + public long getTaskCount() { + return this.delegate.getTaskCount(); + } + + @Override + public long getCompletedTaskCount() { + return this.delegate.getCompletedTaskCount(); + } + + @Override + public String toString() { + return this.delegate.toString(); + } + + @Override + public void beforeExecute(Thread t, Runnable r) { + ReflectionUtils.invokeMethod(this.beforeExecute, this.delegate, t, + new TraceRunnable(tracing(), spanNamer(), r)); + } + + @Override + public void afterExecute(Runnable r, Throwable t) { + ReflectionUtils.invokeMethod(this.afterExecute, this.delegate, + new TraceRunnable(tracing(), spanNamer(), r), t); + } + + @Override + public void terminated() { + ReflectionUtils.invokeMethod(this.terminated, this.delegate); + } + + @Override + @SuppressWarnings("unchecked") + public RunnableFuture newTaskFor(Runnable runnable, T value) { + return (RunnableFuture) ReflectionUtils.invokeMethod(this.newTaskForRunnable, + this.delegate, new TraceRunnable(tracing(), spanNamer(), runnable), + value); + } + + @Override + @SuppressWarnings("unchecked") + public RunnableFuture newTaskFor(Callable callable) { + return (RunnableFuture) ReflectionUtils.invokeMethod(this.newTaskForCallable, + this.delegate, new TraceCallable<>(tracing(), spanNamer(), callable)); + } + + @Override + public T invokeAny(Collection> tasks) + throws InterruptedException, ExecutionException { + return this.delegate.invokeAny(wrapCallableCollection(tasks)); + } + + private Collection> wrapCallableCollection( + Collection> tasks) { + List> ts = new ArrayList<>(); + for (Callable task : tasks) { + if (!(task instanceof TraceCallable)) { + ts.add(new TraceCallable<>(tracing(), spanNamer(), task)); + } + } + return ts; + } + + @Override + public T invokeAny(Collection> tasks, long timeout, + TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + return this.delegate.invokeAny(wrapCallableCollection(tasks), timeout, unit); + } + + @Override + public List> invokeAll(Collection> tasks) + throws InterruptedException { + return this.delegate.invokeAll(wrapCallableCollection(tasks)); + } + + @Override + public List> invokeAll(Collection> tasks, + long timeout, TimeUnit unit) throws InterruptedException { + return this.delegate.invokeAll(wrapCallableCollection(tasks), timeout, unit); + } + + private Tracing tracing() { + if (this.tracing == null) { + this.tracing = this.beanFactory.getBean(Tracing.class); + } + return this.tracing; + } + + private SpanNamer spanNamer() { + if (this.spanNamer == null) { + try { + this.spanNamer = this.beanFactory.getBean(SpanNamer.class); + } + catch (NoSuchBeanDefinitionException e) { + log.warn( + "SpanNamer bean not found - will provide a manually created instance"); + return new DefaultSpanNamer(); + } + } + return this.spanNamer; + } + +} diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java new file mode 100644 index 000000000..f16031d30 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskScheduler.java @@ -0,0 +1,439 @@ +/* + * Copyright 2013-2019 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 + * + * https://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.lang.reflect.Method; +import java.time.Duration; +import java.time.Instant; +import java.util.Date; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionHandler; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.ThreadFactory; + +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; +import org.springframework.lang.Nullable; +import org.springframework.scheduling.Trigger; +import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; +import org.springframework.util.CustomizableThreadCreator; +import org.springframework.util.ErrorHandler; +import org.springframework.util.ReflectionUtils; +import org.springframework.util.concurrent.ListenableFuture; + +/** + * Trace representation of {@link ThreadPoolTaskScheduler}. Should be used only as last + * resort, when any other approaches fail. + * + * @author Marcin Grzejszczak + * @since 2.0.4 + */ +// TODO: Think of a better solution than this +class LazyTraceThreadPoolTaskScheduler extends ThreadPoolTaskScheduler { + + private static final Log log = LogFactory + .getLog(LazyTraceThreadPoolTaskScheduler.class); + + private final BeanFactory beanFactory; + + private final ThreadPoolTaskScheduler delegate; + + private final Method initializeExecutor; + + private final Method createExecutor; + + private final Method cancelRemainingTask; + + private final Method nextThreadName; + + private final Method getDefaultThreadNamePrefix; + + private Tracing tracing; + + private SpanNamer spanNamer; + + LazyTraceThreadPoolTaskScheduler(BeanFactory beanFactory, + ThreadPoolTaskScheduler delegate) { + this.beanFactory = beanFactory; + this.delegate = delegate; + this.initializeExecutor = ReflectionUtils + .findMethod(ThreadPoolTaskScheduler.class, "initializeExecutor", null); + makeAccessibleIfNotNull(this.initializeExecutor); + this.createExecutor = ReflectionUtils.findMethod(ThreadPoolTaskScheduler.class, + "createExecutor", null); + makeAccessibleIfNotNull(this.createExecutor); + this.cancelRemainingTask = ReflectionUtils + .findMethod(ThreadPoolTaskScheduler.class, "cancelRemainingTask", null); + makeAccessibleIfNotNull(this.cancelRemainingTask); + this.nextThreadName = ReflectionUtils.findMethod(ThreadPoolTaskScheduler.class, + "nextThreadName", null); + makeAccessibleIfNotNull(this.nextThreadName); + this.getDefaultThreadNamePrefix = ReflectionUtils.findMethod( + CustomizableThreadCreator.class, "getDefaultThreadNamePrefix", null); + makeAccessibleIfNotNull(this.getDefaultThreadNamePrefix); + } + + private void makeAccessibleIfNotNull(Method method) { + if (method != null) { + ReflectionUtils.makeAccessible(method); + } + } + + @Override + public void setPoolSize(int poolSize) { + this.delegate.setPoolSize(poolSize); + } + + @Override + public void setRemoveOnCancelPolicy(boolean removeOnCancelPolicy) { + this.delegate.setRemoveOnCancelPolicy(removeOnCancelPolicy); + } + + @Override + public void setErrorHandler(ErrorHandler errorHandler) { + this.delegate.setErrorHandler(errorHandler); + } + + @Override + public ExecutorService initializeExecutor(ThreadFactory threadFactory, + RejectedExecutionHandler rejectedExecutionHandler) { + ExecutorService executorService = (ExecutorService) ReflectionUtils.invokeMethod( + this.initializeExecutor, this.delegate, traceThreadFactory(threadFactory), + rejectedExecutionHandler); + if (executorService instanceof TraceableScheduledExecutorService) { + return executorService; + } + return new TraceableExecutorService(this.beanFactory, executorService); + } + + private ThreadFactory traceThreadFactory(ThreadFactory threadFactory) { + return r -> threadFactory.newThread(new TraceRunnable(tracing(), spanNamer(), r)); + } + + @Override + public ScheduledExecutorService createExecutor(int poolSize, + ThreadFactory threadFactory, + RejectedExecutionHandler rejectedExecutionHandler) { + ScheduledExecutorService executorService = (ScheduledExecutorService) ReflectionUtils + .invokeMethod(this.createExecutor, this.delegate, poolSize, + traceThreadFactory(threadFactory), rejectedExecutionHandler); + if (executorService instanceof TraceableScheduledExecutorService) { + return executorService; + } + return new TraceableScheduledExecutorService(this.beanFactory, executorService); + } + + @Override + public ScheduledExecutorService getScheduledExecutor() throws IllegalStateException { + ScheduledExecutorService executor = this.delegate.getScheduledExecutor(); + return executor instanceof TraceableScheduledExecutorService ? executor + : new TraceableScheduledExecutorService(this.beanFactory, executor); + } + + @Override + public ScheduledThreadPoolExecutor getScheduledThreadPoolExecutor() + throws IllegalStateException { + ScheduledThreadPoolExecutor executor = this.delegate + .getScheduledThreadPoolExecutor(); + if (executor instanceof LazyTraceScheduledThreadPoolExecutor) { + return executor; + } + return new LazyTraceScheduledThreadPoolExecutor(executor.getCorePoolSize(), + executor.getThreadFactory(), executor.getRejectedExecutionHandler(), + this.beanFactory, executor); + } + + @Override + public int getPoolSize() { + return this.delegate.getPoolSize(); + } + + @Override + public boolean isRemoveOnCancelPolicy() { + return this.delegate.isRemoveOnCancelPolicy(); + } + + @Override + public int getActiveCount() { + return this.delegate.getActiveCount(); + } + + @Override + public void execute(Runnable task) { + this.delegate.execute(task); + } + + @Override + public void execute(Runnable task, long startTimeout) { + this.delegate.execute(new TraceRunnable(tracing(), spanNamer(), task), + startTimeout); + } + + @Override + public Future submit(Runnable task) { + return this.delegate.submit(new TraceRunnable(tracing(), spanNamer(), task)); + } + + @Override + public Future submit(Callable task) { + return this.delegate.submit(new TraceCallable<>(tracing(), spanNamer(), task)); + } + + @Override + public ListenableFuture submitListenable(Runnable task) { + return this.delegate + .submitListenable(new TraceRunnable(tracing(), spanNamer(), task)); + } + + @Override + public ListenableFuture submitListenable(Callable task) { + return this.delegate + .submitListenable(new TraceCallable<>(tracing(), spanNamer(), task)); + } + + @Override + public void cancelRemainingTask(Runnable task) { + ReflectionUtils.invokeMethod(this.cancelRemainingTask, this.delegate, + new TraceRunnable(tracing(), spanNamer(), task)); + } + + @Override + public boolean prefersShortLivedTasks() { + return this.delegate.prefersShortLivedTasks(); + } + + @Override + @Nullable + public ScheduledFuture schedule(Runnable task, Trigger trigger) { + return this.delegate.schedule(new TraceRunnable(tracing(), spanNamer(), task), + trigger); + } + + @Override + public ScheduledFuture schedule(Runnable task, Date startTime) { + return this.delegate.schedule(new TraceRunnable(tracing(), spanNamer(), task), + startTime); + } + + @Override + public ScheduledFuture scheduleAtFixedRate(Runnable task, Date startTime, + long period) { + return this.delegate.scheduleAtFixedRate( + new TraceRunnable(tracing(), spanNamer(), task), startTime, period); + } + + @Override + public ScheduledFuture scheduleAtFixedRate(Runnable task, long period) { + return this.delegate.scheduleAtFixedRate( + new TraceRunnable(tracing(), spanNamer(), task), period); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay(Runnable task, Date startTime, + long delay) { + return this.delegate.scheduleWithFixedDelay( + new TraceRunnable(tracing(), spanNamer(), task), startTime, delay); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay(Runnable task, long delay) { + return this.delegate.scheduleWithFixedDelay( + new TraceRunnable(tracing(), spanNamer(), task), delay); + } + + @Override + public void setThreadFactory(ThreadFactory threadFactory) { + this.delegate.setThreadFactory(threadFactory); + } + + @Override + public void setThreadNamePrefix(String threadNamePrefix) { + this.delegate.setThreadNamePrefix(threadNamePrefix); + } + + @Override + public void setRejectedExecutionHandler( + RejectedExecutionHandler rejectedExecutionHandler) { + this.delegate.setRejectedExecutionHandler(rejectedExecutionHandler); + } + + @Override + public void setWaitForTasksToCompleteOnShutdown( + boolean waitForJobsToCompleteOnShutdown) { + this.delegate + .setWaitForTasksToCompleteOnShutdown(waitForJobsToCompleteOnShutdown); + } + + @Override + public void setAwaitTerminationSeconds(int awaitTerminationSeconds) { + this.delegate.setAwaitTerminationSeconds(awaitTerminationSeconds); + } + + @Override + public void setBeanName(String name) { + this.delegate.setBeanName(name); + } + + @Override + public void afterPropertiesSet() { + this.delegate.afterPropertiesSet(); + } + + @Override + public void initialize() { + this.delegate.initialize(); + } + + @Override + public void destroy() { + this.delegate.destroy(); + } + + @Override + public void shutdown() { + this.delegate.shutdown(); + } + + @Override + public Thread newThread(Runnable runnable) { + return this.delegate.newThread(runnable); + } + + @Override + public String getThreadNamePrefix() { + return this.delegate.getThreadNamePrefix(); + } + + @Override + public void setThreadPriority(int threadPriority) { + this.delegate.setThreadPriority(threadPriority); + } + + @Override + public int getThreadPriority() { + return this.delegate.getThreadPriority(); + } + + @Override + public void setDaemon(boolean daemon) { + this.delegate.setDaemon(daemon); + } + + @Override + public boolean isDaemon() { + return this.delegate.isDaemon(); + } + + @Override + public void setThreadGroupName(String name) { + this.delegate.setThreadGroupName(name); + } + + @Override + public void setThreadGroup(ThreadGroup threadGroup) { + this.delegate.setThreadGroup(threadGroup); + } + + @Override + @Nullable + public ThreadGroup getThreadGroup() { + return this.delegate.getThreadGroup(); + } + + @Override + public Thread createThread(Runnable runnable) { + return this.delegate.createThread(runnable); + } + + @Override + public String nextThreadName() { + return (String) ReflectionUtils.invokeMethod(this.nextThreadName, this.delegate); + } + + @Override + public String getDefaultThreadNamePrefix() { + if (this.delegate == null) { + return super.getDefaultThreadNamePrefix(); + } + return (String) ReflectionUtils.invokeMethod(this.getDefaultThreadNamePrefix, + this.delegate); + } + + @Override + public ScheduledFuture schedule(Runnable task, Instant startTime) { + return this.delegate.schedule(new TraceRunnable(tracing(), spanNamer(), task), + startTime); + } + + @Override + public ScheduledFuture scheduleAtFixedRate(Runnable task, Instant startTime, + Duration period) { + return this.delegate.scheduleAtFixedRate( + new TraceRunnable(tracing(), spanNamer(), task), startTime, period); + } + + @Override + public ScheduledFuture scheduleAtFixedRate(Runnable task, Duration period) { + return this.delegate.scheduleAtFixedRate( + new TraceRunnable(tracing(), spanNamer(), task), period); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay(Runnable task, Instant startTime, + Duration delay) { + return this.delegate.scheduleWithFixedDelay( + new TraceRunnable(tracing(), spanNamer(), task), startTime, delay); + } + + @Override + public ScheduledFuture scheduleWithFixedDelay(Runnable task, Duration delay) { + return this.delegate.scheduleWithFixedDelay( + new TraceRunnable(tracing(), spanNamer(), task), delay); + } + + private Tracing tracing() { + if (this.tracing == null) { + this.tracing = this.beanFactory.getBean(Tracing.class); + } + return this.tracing; + } + + private SpanNamer spanNamer() { + if (this.spanNamer == null) { + try { + this.spanNamer = this.beanFactory.getBean(SpanNamer.class); + } + catch (NoSuchBeanDefinitionException e) { + log.warn( + "SpanNamer bean not found - will provide a manually created instance"); + return new DefaultSpanNamer(); + } + } + return this.spanNamer; + } + +} diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAspect.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAspect.java index aef5712be..12c37b6e6 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAspect.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAspect.java @@ -60,6 +60,9 @@ public class TraceSchedulingAspect { @Around("execution (@org.springframework.scheduling.annotation.Scheduled * *.*(..))") public Object traceBackgroundThread(final ProceedingJoinPoint pjp) throws Throwable { if (this.skipPattern.matcher(pjp.getTarget().getClass().getName()).matches()) { + // we might have a span in context due to wrapping of runnables + // we want to clear that context + this.tracer.withSpanInScope(null); return pjp.proceed(); } String spanName = SpanNameUtil.toLowerHyphen(pjp.getSignature().getName()); diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java new file mode 100644 index 000000000..9c01b8ba5 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/LazyTraceThreadPoolTaskSchedulerTests.java @@ -0,0 +1,462 @@ +/* + * Copyright 2013-2019 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 + * + * https://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.time.Duration; +import java.time.Instant; +import java.util.Date; +import java.util.concurrent.Callable; +import java.util.concurrent.RejectedExecutionHandler; +import java.util.concurrent.ThreadFactory; + +import brave.Tracing; +import brave.propagation.StrictScopeDecorator; +import brave.propagation.ThreadLocalCurrentTraceContext; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.BDDMockito; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnitRunner; + +import org.springframework.beans.factory.BeanFactory; +import org.springframework.cloud.sleuth.DefaultSpanNamer; +import org.springframework.cloud.sleuth.SpanNamer; +import org.springframework.scheduling.Trigger; +import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; +import org.springframework.util.ErrorHandler; + +@RunWith(MockitoJUnitRunner.class) +public class LazyTraceThreadPoolTaskSchedulerTests { + + Tracing tracing = Tracing.newBuilder() + .currentTraceContext(ThreadLocalCurrentTraceContext.newBuilder() + .addScopeDecorator(StrictScopeDecorator.create()).build()) + .build(); + + @Mock + BeanFactory beanFactory; + + @Mock + ThreadPoolTaskScheduler delegate; + + LazyTraceThreadPoolTaskScheduler executor; + + @Before + public void setup() { + this.executor = new LazyTraceThreadPoolTaskScheduler(beanFactory(), + this.delegate); + } + + BeanFactory beanFactory() { + BDDMockito.given(this.beanFactory.getBean(Tracing.class)) + .willReturn(this.tracing); + BDDMockito.given(this.beanFactory.getBean(SpanNamer.class)) + .willReturn(new DefaultSpanNamer()); + ContextRefreshedListenerAccessor.set(this.beanFactory, true); + return this.beanFactory; + } + + @Test + public void setPoolSize() { + this.executor.setPoolSize(10); + + BDDMockito.then(this.delegate).should().setPoolSize(10); + } + + @Test + public void setRemoveOnCancelPolicy() { + this.executor.setRemoveOnCancelPolicy(true); + + BDDMockito.then(this.delegate).should().setRemoveOnCancelPolicy(true); + } + + @Test + public void setErrorHandler() { + ErrorHandler handler = (throwable) -> { + }; + this.executor.setErrorHandler(handler); + + BDDMockito.then(this.delegate).should().setErrorHandler(handler); + } + + @Test + public void getScheduledExecutor() { + this.executor.getScheduledExecutor(); + + BDDMockito.then(this.delegate).should().getScheduledExecutor(); + } + + @Test + public void getPoolSize() { + this.executor.getPoolSize(); + + BDDMockito.then(this.delegate).should().getPoolSize(); + } + + @Test + public void isRemoveOnCancelPolicy() { + this.executor.setRemoveOnCancelPolicy(true); + + BDDMockito.then(this.delegate).should().setRemoveOnCancelPolicy(true); + } + + @Test + public void getActiveCount() { + this.executor.getActiveCount(); + + BDDMockito.then(this.delegate).should().getActiveCount(); + } + + @Test + public void execute() { + Runnable r = () -> { + }; + this.executor.execute(r); + + BDDMockito.then(this.delegate).should().execute(r); + } + + @Test + public void execute1() { + Runnable r = () -> { + }; + this.executor.execute(r, 10L); + + BDDMockito.then(this.delegate).should() + .execute(BDDMockito.any(TraceRunnable.class), BDDMockito.eq(10L)); + } + + @Test + public void submit() { + Runnable c = () -> { + }; + this.executor.submit(c); + + BDDMockito.then(this.delegate).should() + .submit(BDDMockito.any(TraceRunnable.class)); + } + + @Test + public void submit1() { + Callable c = () -> null; + this.executor.submit(c); + + BDDMockito.then(this.delegate).should() + .submit(BDDMockito.any(TraceCallable.class)); + } + + @Test + public void submitListenable() { + Runnable c = () -> { + }; + this.executor.submitListenable(c); + + BDDMockito.then(this.delegate).should() + .submitListenable(BDDMockito.any(TraceRunnable.class)); + } + + @Test + public void submitListenable1() { + Callable c = () -> null; + this.executor.submitListenable(c); + + BDDMockito.then(this.delegate).should() + .submitListenable(BDDMockito.any(TraceCallable.class)); + } + + @Test + public void prefersShortLivedTasks() { + this.executor.prefersShortLivedTasks(); + + BDDMockito.then(this.delegate).should().prefersShortLivedTasks(); + } + + @Test + public void schedule() { + Runnable c = () -> { + }; + Trigger trigger = triggerContext -> null; + this.executor.schedule(c, trigger); + + BDDMockito.then(this.delegate).should() + .schedule(BDDMockito.any(TraceRunnable.class), BDDMockito.eq(trigger)); + } + + @Test + public void schedule1() { + Runnable c = () -> { + }; + Date date = new Date(); + this.executor.schedule(c, date); + + BDDMockito.then(this.delegate).should() + .schedule(BDDMockito.any(TraceRunnable.class), BDDMockito.eq(date)); + } + + @Test + public void scheduleAtFixedRate() { + Runnable c = () -> { + }; + Date date = new Date(); + this.executor.scheduleAtFixedRate(c, date, 10L); + + BDDMockito.then(this.delegate).should().scheduleAtFixedRate( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(date), + BDDMockito.eq(10L)); + } + + @Test + public void scheduleAtFixedRate1() { + Runnable c = () -> { + }; + this.executor.scheduleAtFixedRate(c, 10L); + + BDDMockito.then(this.delegate).should().scheduleAtFixedRate( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(10L)); + } + + @Test + public void scheduleWithFixedDelay() { + Runnable c = () -> { + }; + Date date = new Date(); + this.executor.scheduleWithFixedDelay(c, date, 10L); + + BDDMockito.then(this.delegate).should().scheduleWithFixedDelay( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(date), + BDDMockito.eq(10L)); + } + + @Test + public void scheduleWithFixedDelay1() { + Runnable c = () -> { + }; + this.executor.scheduleWithFixedDelay(c, 10L); + + BDDMockito.then(this.delegate).should().scheduleWithFixedDelay( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(10L)); + } + + @Test + public void scheduleWithFixedDelay2() { + Runnable c = () -> { + }; + Instant instant = Instant.now(); + Duration duration = Duration.ZERO; + this.executor.scheduleWithFixedDelay(c, instant, duration); + + BDDMockito.then(this.delegate).should().scheduleWithFixedDelay( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(instant), + BDDMockito.eq(duration)); + } + + @Test + public void scheduleWithFixedDelay3() { + Runnable c = () -> { + }; + Duration duration = Duration.ZERO; + this.executor.scheduleWithFixedDelay(c, duration); + + BDDMockito.then(this.delegate).should().scheduleWithFixedDelay( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(duration)); + } + + @Test + public void setThreadFactory() { + ThreadFactory threadFactory = r -> null; + this.executor.setThreadFactory(threadFactory); + + BDDMockito.then(this.delegate).should().setThreadFactory(threadFactory); + } + + @Test + public void setThreadNamePrefix() { + this.executor.setThreadNamePrefix("foo"); + + BDDMockito.then(this.delegate).should().setThreadNamePrefix("foo"); + } + + @Test + public void setRejectedExecutionHandler() { + RejectedExecutionHandler handler = (r, executor1) -> { + }; + this.executor.setRejectedExecutionHandler(handler); + + BDDMockito.then(this.delegate).should().setRejectedExecutionHandler(handler); + } + + @Test + public void setWaitForTasksToCompleteOnShutdown() { + this.executor.setWaitForTasksToCompleteOnShutdown(true); + + BDDMockito.then(this.delegate).should().setWaitForTasksToCompleteOnShutdown(true); + } + + @Test + public void setAwaitTerminationSeconds() { + this.executor.setAwaitTerminationSeconds(10); + + BDDMockito.then(this.delegate).should().setAwaitTerminationSeconds(10); + } + + @Test + public void setBeanName() { + this.executor.setBeanName("foo"); + + BDDMockito.then(this.delegate).should().setBeanName("foo"); + } + + @Test + public void afterPropertiesSet() { + this.executor.afterPropertiesSet(); + + BDDMockito.then(this.delegate).should().afterPropertiesSet(); + } + + @Test + public void initialize() { + this.executor.initialize(); + + BDDMockito.then(this.delegate).should().initialize(); + } + + @Test + public void destroy() { + this.executor.destroy(); + + BDDMockito.then(this.delegate).should().destroy(); + } + + @Test + public void shutdown() { + this.executor.shutdown(); + + BDDMockito.then(this.delegate).should().shutdown(); + } + + @Test + public void newThread() { + Runnable runnable = () -> { + }; + this.executor.newThread(runnable); + + BDDMockito.then(this.delegate).should().newThread(runnable); + } + + @Test + public void getThreadNamePrefix() { + this.executor.getThreadNamePrefix(); + + BDDMockito.then(this.delegate).should().getThreadNamePrefix(); + } + + @Test + public void setThreadPriority() { + this.executor.setThreadPriority(10); + + BDDMockito.then(this.delegate).should().setThreadPriority(10); + } + + @Test + public void getThreadPriority() { + this.executor.getThreadPriority(); + + BDDMockito.then(this.delegate).should().getThreadPriority(); + } + + @Test + public void setDaemon() { + this.executor.setDaemon(true); + + BDDMockito.then(this.delegate).should().setDaemon(true); + } + + @Test + public void isDaemon() { + this.executor.isDaemon(); + + BDDMockito.then(this.delegate).should().isDaemon(); + } + + @Test + public void setThreadGroupName() { + this.executor.setThreadGroupName("foo"); + + BDDMockito.then(this.delegate).should().setThreadGroupName("foo"); + } + + @Test + public void setThreadGroup() { + ThreadGroup threadGroup = new ThreadGroup("foo"); + this.executor.setThreadGroup(threadGroup); + + BDDMockito.then(this.delegate).should().setThreadGroup(threadGroup); + } + + @Test + public void getThreadGroup() { + this.executor.getThreadGroup(); + + BDDMockito.then(this.delegate).should().getThreadGroup(); + } + + @Test + public void createThread() { + Runnable r = () -> { + }; + this.executor.createThread(r); + + BDDMockito.then(this.delegate).should().createThread(r); + } + + @Test + public void schedule2() { + Runnable r = () -> { + }; + Instant instant = Instant.now(); + this.executor.schedule(r, instant); + + BDDMockito.then(this.delegate).should() + .schedule(BDDMockito.any(TraceRunnable.class), BDDMockito.eq(instant)); + } + + @Test + public void scheduleAtFixedRate2() { + Runnable r = () -> { + }; + Instant instant = Instant.now(); + Duration duration = Duration.ZERO; + this.executor.scheduleAtFixedRate(r, instant, duration); + + BDDMockito.then(this.delegate).should().scheduleAtFixedRate( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(instant), + BDDMockito.eq(duration)); + } + + @Test + public void scheduleAtFixedRate3() { + Runnable r = () -> { + }; + Duration duration = Duration.ZERO; + this.executor.scheduleAtFixedRate(r, duration); + + BDDMockito.then(this.delegate).should().scheduleAtFixedRate( + BDDMockito.any(TraceRunnable.class), BDDMockito.eq(duration)); + } + +} diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue410/Issue410Tests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue410/Issue410Tests.java index 1e4c6f445..bfef80152 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue410/Issue410Tests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue410/Issue410Tests.java @@ -20,6 +20,7 @@ import java.lang.invoke.MethodHandles; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.atomic.AtomicReference; import brave.Span; @@ -45,6 +46,7 @@ import org.springframework.core.env.Environment; import org.springframework.scheduling.annotation.Async; import org.springframework.scheduling.annotation.EnableAsync; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; +import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.stereotype.Component; import org.springframework.test.context.junit4.SpringRunner; import org.springframework.web.bind.annotation.RequestMapping; @@ -176,6 +178,58 @@ public class Issue410Tests { then(this.tracer.currentSpan()).isNull(); } + /** + * Related to issue #1232 + */ + @Test + public void should_pass_tracing_info_for_completable_futures_with_threadPoolTaskScheduler() { + Span span = this.tracer.nextSpan().name("foo"); + log.info("Starting test"); + try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) { + String response = this.restTemplate.getForObject( + "http://localhost:" + port() + "/threadPoolTaskScheduler", + String.class); + + then(response).isEqualTo(span.context().traceIdString()); + Awaitility.await().untilAsserted(() -> { + then(this.asyncTask.getSpan().get()).isNotNull(); + then(this.asyncTask.getSpan().get().context().traceId()) + .isEqualTo(span.context().traceId()); + }); + } + finally { + span.finish(); + } + + then(this.tracer.currentSpan()).isNull(); + } + + /** + * Related to issue #1232 + */ + @Test + public void should_pass_tracing_info_for_completable_futures_with_scheduledThreadPoolExecutor() { + Span span = this.tracer.nextSpan().name("foo"); + log.info("Starting test"); + try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) { + String response = this.restTemplate.getForObject( + "http://localhost:" + port() + "/scheduledThreadPoolExecutor", + String.class); + + then(response).isEqualTo(span.context().traceIdString()); + Awaitility.await().untilAsserted(() -> { + then(this.asyncTask.getSpan().get()).isNotNull(); + then(this.asyncTask.getSpan().get().context().traceId()) + .isEqualTo(span.context().traceId()); + }); + } + finally { + span.finish(); + } + + then(this.tracer.currentSpan()).isNull(); + } + private int port() { return this.environment.getProperty("local.server.port", Integer.class); } @@ -203,6 +257,18 @@ class AppConfig { return executor; } + @Bean + public ThreadPoolTaskScheduler threadPoolTaskScheduler() { + ThreadPoolTaskScheduler executor = new ThreadPoolTaskScheduler(); + executor.initialize(); + return executor; + } + + @Bean + public ScheduledThreadPoolExecutor scheduledThreadPoolExecutor() { + return new ScheduledThreadPoolExecutor(10); + } + } @Component @@ -224,6 +290,12 @@ class AsyncTask { @Autowired BeanFactory beanFactory; + @Autowired + ThreadPoolTaskScheduler threadPoolTaskScheduler; + + @Autowired + ScheduledThreadPoolExecutor scheduledThreadPoolExecutor; + private AtomicReference span = new AtomicReference<>(); @Async("poolTaskExecutor") @@ -290,6 +362,26 @@ class AsyncTask { return this.span.get(); } + public Span scheduledThreadPoolExecutor() + throws ExecutionException, InterruptedException { + log.info("This task is running with ScheduledThreadPoolExecutor"); + this.scheduledThreadPoolExecutor.submit(() -> { + log.info("Hello from runnable"); + AsyncTask.this.span.set(AsyncTask.this.tracer.currentSpan()); + }).get(); + return this.span.get(); + } + + public Span threadPoolTaskScheduler() + throws ExecutionException, InterruptedException { + log.info("This task is running with ThreadPoolTaskScheduler"); + this.threadPoolTaskScheduler.submit(() -> { + log.info("Hello from runnable"); + AsyncTask.this.span.set(AsyncTask.this.tracer.currentSpan()); + }).get(); + return this.span.get(); + } + public AtomicReference getSpan() { return this.span; } @@ -335,6 +427,20 @@ class Application { return this.asyncTask.taskScheduler().context().traceIdString(); } + @RequestMapping("/threadPoolTaskScheduler") + public String threadPoolTaskScheduler() + throws ExecutionException, InterruptedException { + log.info("Executing completable via ThreadPoolTaskScheduler"); + return this.asyncTask.threadPoolTaskScheduler().context().traceIdString(); + } + + @RequestMapping("/scheduledThreadPoolExecutor") + public String scheduledThreadPoolExecutor() + throws ExecutionException, InterruptedException { + log.info("Executing completable via ScheduledThreadPoolExecutor"); + return this.asyncTask.scheduledThreadPoolExecutor().context().traceIdString(); + } + /** * Related to issue #445. * @return service bean diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/scheduling/TracingOnScheduledTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/scheduling/TracingOnScheduledTests.java index f58aa58e1..27fec4de7 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/scheduling/TracingOnScheduledTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/scheduling/TracingOnScheduledTests.java @@ -68,6 +68,7 @@ public class TracingOnScheduledTests { public void setup() { this.beanWithScheduledMethod.clear(); this.beanWithScheduledMethodToBeIgnored.clear(); + this.reporter.clear(); } @Test