Finish async customizer

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/LazyTraceAsyncCustomizer.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2015 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.scheduling;
+
+import java.util.concurrent.Executor;
+
+import lombok.RequiredArgsConstructor;
+
+import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.cloud.sleuth.Trace;
+import org.springframework.scheduling.annotation.AsyncConfigurer;
+import org.springframework.scheduling.annotation.AsyncConfigurerSupport;
+
+/**
+ * @author Dave Syer
+ *
+ */
+@RequiredArgsConstructor
+public class LazyTraceAsyncCustomizer extends AsyncConfigurerSupport {
+
+	private Trace trace;
+	private final BeanFactory beanFactory;
+	private final AsyncConfigurer delegate;
+
+	@Override
+	public Executor getAsyncExecutor() {
+		if (this.trace == null) {
+			this.trace = this.beanFactory.getBean(Trace.class);
+		}
+		return new TraceExecutor(this.trace, this.delegate.getAsyncExecutor());
+	}
+
+	@Override
+	public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
+		return this.delegate.getAsyncUncaughtExceptionHandler();
+	}
+
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/scheduling/TraceSchedulingAutoConfiguration.java

@@ -7,13 +7,19 @@ package org.springframework.cloud.sleuth.instrument.scheduling;
 import java.util.concurrent.Executor;
 
 import org.aspectj.lang.ProceedingJoinPoint;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.cloud.sleuth.Trace;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.EnableAspectJAutoProxy;
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
+import org.springframework.scheduling.annotation.AsyncConfigurer;
 import org.springframework.scheduling.annotation.AsyncConfigurerSupport;
 import org.springframework.scheduling.annotation.EnableAsync;
 
@@ -37,12 +43,12 @@ public class TraceSchedulingAutoConfiguration {
 
 	@EnableAsync
 	@Configuration
-	protected static class AsyncConfiguration extends AsyncConfigurerSupport {
+	@ConditionalOnMissingBean(AsyncConfigurer.class)
+	protected static class AsyncDefaultConfiguration extends AsyncConfigurerSupport {
 
 		@Autowired
 		private Trace trace;
 
-		// TODO: look for an existing AsyncConfigurer and steal its Executor
 		@Override
 		public Executor getAsyncExecutor() {
 			return new TraceExecutor(this.trace, new SimpleAsyncTaskExecutor());
@@ -50,4 +56,29 @@ public class TraceSchedulingAutoConfiguration {
 
 	}
 
+	@Configuration
+	@ConditionalOnBean(AsyncConfigurer.class)
+	protected static class AsyncCustomConfiguration implements BeanPostProcessor {
+
+		@Autowired
+		private BeanFactory beanFactory;
+
+		@Override
+		public Object postProcessBeforeInitialization(Object bean, String beanName)
+				throws BeansException {
+			return bean;
+		}
+
+		@Override
+		public Object postProcessAfterInitialization(Object bean, String beanName)
+				throws BeansException {
+			if (bean instanceof AsyncConfigurer) {
+				AsyncConfigurer configurer = (AsyncConfigurer) bean;
+				return new LazyTraceAsyncCustomizer(this.beanFactory, configurer);
+			}
+			return bean;
+		}
+
+	}
+
 }