Fix ThreadPoolTaskScheduler proxy mechanism (#1447)

This commit is contained in:
Nick
2019-09-19 10:49:46 +03:00
committed by Marcin Grzejszczak
parent 2e6ae5d314
commit a5d41840b3
2 changed files with 54 additions and 8 deletions

View File

@@ -195,8 +195,13 @@ class ExecutorBeanPostProcessor implements BeanPostProcessor {
Object createAsyncTaskExecutorProxy(Object bean, boolean cglibProxy,
AsyncTaskExecutor executor) {
return getProxiedObject(bean, cglibProxy, executor,
() -> new LazyTraceAsyncTaskExecutor(this.beanFactory, executor));
return getProxiedObject(bean, cglibProxy, executor, () -> {
if (bean instanceof ThreadPoolTaskScheduler) {
return new LazyTraceThreadPoolTaskScheduler(this.beanFactory,
(ThreadPoolTaskScheduler) executor);
}
return new LazyTraceAsyncTaskExecutor(this.beanFactory, executor);
});
}
private Object getProxiedObject(Object bean, boolean cglibProxy, Executor executor,

View File

@@ -17,6 +17,7 @@
package org.springframework.cloud.sleuth.instrument.async.issues.issue410;
import java.lang.invoke.MethodHandles;
import java.util.Date;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
@@ -184,12 +185,35 @@ public class Issue410Tests {
* Related to issue #1232
*/
@Test
public void should_pass_tracing_info_for_completable_futures_with_threadPoolTaskScheduler() {
public void should_pass_tracing_info_for_submitted_tasks_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",
"http://localhost:" + port() + "/threadPoolTaskScheduler_submit",
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();
}
@Test
public void should_pass_tracing_info_for_scheduled_tasks_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_schedule",
String.class);
then(response).isEqualTo(span.context().traceIdString());
@@ -379,7 +403,7 @@ class AsyncTask {
return this.span.get();
}
public Span threadPoolTaskScheduler()
public Span threadPoolTaskSchedulerSubmit()
throws ExecutionException, InterruptedException {
log.info("This task is running with ThreadPoolTaskScheduler");
this.threadPoolTaskScheduler.submit(() -> {
@@ -389,6 +413,16 @@ class AsyncTask {
return this.span.get();
}
public Span threadPoolTaskSchedulerSchedule()
throws ExecutionException, InterruptedException {
log.info("This task is running with ThreadPoolTaskScheduler");
this.threadPoolTaskScheduler.schedule(() -> {
log.info("Hello from runnable");
AsyncTask.this.span.set(AsyncTask.this.tracer.currentSpan());
}, new Date()).get();
return this.span.get();
}
public AtomicReference<Span> getSpan() {
return this.span;
}
@@ -434,11 +468,18 @@ class Application {
return this.asyncTask.taskScheduler().context().traceIdString();
}
@RequestMapping("/threadPoolTaskScheduler")
public String threadPoolTaskScheduler()
@RequestMapping("/threadPoolTaskScheduler_submit")
public String threadPoolTaskSchedulerSubmit()
throws ExecutionException, InterruptedException {
log.info("Executing completable via ThreadPoolTaskScheduler");
return this.asyncTask.threadPoolTaskScheduler().context().traceIdString();
return this.asyncTask.threadPoolTaskSchedulerSubmit().context().traceIdString();
}
@RequestMapping("/threadPoolTaskScheduler_schedule")
public String threadPoolTaskSchedulerSchedule()
throws ExecutionException, InterruptedException {
log.info("Executing completable via ThreadPoolTaskScheduler");
return this.asyncTask.threadPoolTaskSchedulerSchedule().context().traceIdString();
}
@RequestMapping("/scheduledThreadPoolExecutor")