Fix ThreadPoolTaskScheduler proxy mechanism (#1447)

Author: nkavtur

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/async/ExecutorBeanPostProcessor.java

@@ -185,8 +185,13 @@ Object createExecutorServiceProxy(Object bean, boolean cglibProxy,
 
 	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,

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue410/Issue410Tests.java

@@ -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;
@@ -182,12 +183,35 @@ public void should_pass_tracing_info_for_completable_futures_with_task_scheduler
 	 * 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());
@@ -372,7 +396,7 @@ public Span scheduledThreadPoolExecutor()
 		return this.span.get();
 	}
 
-	public Span threadPoolTaskScheduler()
+	public Span threadPoolTaskSchedulerSubmit()
 			throws ExecutionException, InterruptedException {
 		log.info("This task is running with ThreadPoolTaskScheduler");
 		this.threadPoolTaskScheduler.submit(() -> {
@@ -382,6 +406,16 @@ public Span threadPoolTaskScheduler()
 		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;
 	}
@@ -427,11 +461,18 @@ public String taskScheduler() throws ExecutionException, InterruptedException {
 		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.threadPoolTaskSchedulerSubmit().context().traceIdString();
+	}
+
+	@RequestMapping("/threadPoolTaskScheduler_schedule")
+	public String threadPoolTaskSchedulerSchedule()
 			throws ExecutionException, InterruptedException {
 		log.info("Executing completable via ThreadPoolTaskScheduler");
-		return this.asyncTask.threadPoolTaskScheduler().context().traceIdString();
+		return this.asyncTask.threadPoolTaskSchedulerSchedule().context().traceIdString();
 	}
 
 	@RequestMapping("/scheduledThreadPoolExecutor")