[#124] Added support for async rest template

fixes #124

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/AbstractTraceHttpRequestInterceptor.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2013-2016 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.web.client;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanAccessor;
+import org.springframework.cloud.sleuth.event.ClientReceivedEvent;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.context.ApplicationEventPublisherAware;
+import org.springframework.http.HttpRequest;
+import org.springframework.util.StringUtils;
+
+/**
+ * Abstraction over classes that interact with Http requests. Allows you
+ * to enrich the request headers with trace related information.
+ *
+ * @author Marcin Grzejszczak
+ */
+public abstract class AbstractTraceHttpRequestInterceptor
+		implements ApplicationEventPublisherAware {
+
+	private ApplicationEventPublisher publisher;
+	private final SpanAccessor accessor;
+
+	public AbstractTraceHttpRequestInterceptor(SpanAccessor accessor) {
+		this.accessor = accessor;
+	}
+
+	@Override
+	public void setApplicationEventPublisher(ApplicationEventPublisher publisher) {
+		this.publisher = publisher;
+	}
+
+	/**
+	 * Adds trace related headers from the span to the request
+	 */
+	public void enrichWithTraceHeaders(HttpRequest request, Span span) {
+		setHeader(request, Span.TRACE_ID_NAME, span.getTraceId());
+		setHeader(request, Span.SPAN_ID_NAME, span.getSpanId());
+		if (!span.isExportable()) {
+			setHeader(request, Span.NOT_SAMPLED_NAME, "true");
+		}
+		setHeader(request, Span.SPAN_NAME_NAME, span.getName().toString());
+		setHeader(request, Span.PARENT_ID_NAME, getParentId(span));
+		setHeader(request, Span.PROCESS_ID_NAME, span.getProcessId());
+	}
+
+	private Long getParentId(Span span) {
+		return !span.getParents().isEmpty() ? span.getParents().get(0) : null;
+	}
+
+	public void setHeader(HttpRequest request, String name, String value) {
+		if (StringUtils.hasText(value) && !request.getHeaders().containsKey(name) && this.accessor.isTracing()) {
+			request.getHeaders().add(name, value);
+		}
+	}
+
+	public void setHeader(HttpRequest request, String name, Long value) {
+		if (value != null) {
+			setHeader(request, name, Span.toHex(value));
+		}
+	}
+
+	/**
+	 * Close the current span and emit the ClientReceivedEvent
+	 */
+	public void close() {
+		if (getCurrentSpan() == null) {
+			return;
+		}
+		publish(new ClientReceivedEvent(this, getCurrentSpan()));
+	}
+
+	protected void publish(ApplicationEvent event) {
+		if (this.publisher != null) {
+			this.publisher.publishEvent(event);
+		}
+	}
+
+	protected Span getCurrentSpan() {
+		return this.accessor.getCurrentSpan();
+	}
+
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncClientHttpRequestFactoryWrapper.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2013-2016 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.web.client;
+
+import java.io.IOException;
+import java.net.URI;
+
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.SpanAccessor;
+import org.springframework.cloud.sleuth.event.ClientSentEvent;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.http.HttpMethod;
+import org.springframework.http.client.AsyncClientHttpRequest;
+import org.springframework.http.client.AsyncClientHttpRequestFactory;
+
+/**
+ * Wrapper that adds trace related headers to the created AsyncClientHttpRequest
+ *
+ * @see org.springframework.web.client.RestTemplate
+ * @see SpanAccessor
+ *
+ * @author Marcin Grzejszczak
+ * @author Spencer Gibb
+ */
+public class TraceAsyncClientHttpRequestFactoryWrapper extends AbstractTraceHttpRequestInterceptor
+		implements AsyncClientHttpRequestFactory {
+
+	private ApplicationEventPublisher publisher;
+	private final AsyncClientHttpRequestFactory delegate;
+
+	public TraceAsyncClientHttpRequestFactoryWrapper(SpanAccessor accessor,
+			AsyncClientHttpRequestFactory delegate) {
+		super(accessor);
+		this.delegate = delegate;
+	}
+
+	@Override
+	public AsyncClientHttpRequest createAsyncRequest(URI uri, HttpMethod httpMethod)
+			throws IOException {
+		AsyncClientHttpRequest request = this.delegate.createAsyncRequest(uri, httpMethod);
+		Span span = getCurrentSpan();
+		if (span == null) {
+			setHeader(request, Span.NOT_SAMPLED_NAME, "true");
+			return request;
+		}
+		enrichWithTraceHeaders(request, span);
+		publish(new ClientSentEvent(this, span));
+		return request;
+	}
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncListenableTaskExecutor.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2013-2016 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.web.client;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.Future;
+
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.core.task.AsyncListenableTaskExecutor;
+import org.springframework.util.concurrent.ListenableFuture;
+
+/**
+ * AsyncListenableTaskExecutor that wraps all Runnable / Callable tasks into
+ * their trace related representation
+ */
+public class TraceAsyncListenableTaskExecutor implements AsyncListenableTaskExecutor {
+
+	private final AsyncListenableTaskExecutor delegate;
+	private final Tracer tracer;
+
+	TraceAsyncListenableTaskExecutor(AsyncListenableTaskExecutor delegate,
+			Tracer tracer) {
+		this.delegate = delegate;
+		this.tracer = tracer;
+	}
+
+	@Override
+	public ListenableFuture<?> submitListenable(Runnable task) {
+		return this.delegate.submitListenable(this.tracer.wrap(task));
+	}
+
+	@Override
+	public <T> ListenableFuture<T> submitListenable(Callable<T> task) {
+		return this.delegate.submitListenable(this.tracer.wrap(task));
+	}
+
+	@Override
+	public void execute(Runnable task, long startTimeout) {
+		this.delegate.execute(this.tracer.wrap(task), startTimeout);
+	}
+
+	@Override
+	public Future<?> submit(Runnable task) {
+		return this.delegate.submit(this.tracer.wrap(task));
+	}
+
+	@Override
+	public <T> Future<T> submit(Callable<T> task) {
+		return this.delegate.submit(this.tracer.wrap(task));
+	}
+
+	@Override
+	public void execute(Runnable task) {
+		this.delegate.execute(this.tracer.wrap(task));
+	}
+
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2015 the original author or authors.
+ * Copyright 2013-2016 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.
@@ -19,16 +19,11 @@ import java.io.IOException;
 
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.SpanAccessor;
-import org.springframework.cloud.sleuth.event.ClientReceivedEvent;
 import org.springframework.cloud.sleuth.event.ClientSentEvent;
-import org.springframework.context.ApplicationEvent;
-import org.springframework.context.ApplicationEventPublisher;
-import org.springframework.context.ApplicationEventPublisherAware;
 import org.springframework.http.HttpRequest;
 import org.springframework.http.client.ClientHttpRequestExecution;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.http.client.ClientHttpResponse;
-import org.springframework.util.StringUtils;
 
 /**
  * Interceptor that verifies whether the trance and span id has been set on the request
@@ -37,23 +32,14 @@ import org.springframework.util.StringUtils;
  * @see org.springframework.web.client.RestTemplate
  * @see SpanAccessor
  *
- * @author Marcin Grzejszczak, 4financeIT
+ * @author Marcin Grzejszczak
  * @author Spencer Gibb
  */
-public class TraceRestTemplateInterceptor
-		implements ClientHttpRequestInterceptor, ApplicationEventPublisherAware {
-
-	private ApplicationEventPublisher publisher;
-
-	private SpanAccessor accessor;
+public class TraceRestTemplateInterceptor extends AbstractTraceHttpRequestInterceptor
+		implements ClientHttpRequestInterceptor {
 
 	public TraceRestTemplateInterceptor(SpanAccessor accessor) {
-		this.accessor = accessor;
-	}
-
-	@Override
-	public void setApplicationEventPublisher(ApplicationEventPublisher publisher) {
-		this.publisher = publisher;
+		super(accessor);
 	}
 
 	@Override
@@ -64,49 +50,9 @@ public class TraceRestTemplateInterceptor
 			setHeader(request, Span.NOT_SAMPLED_NAME, "true");
 			return execution.execute(request, body);
 		}
-		setHeader(request, Span.TRACE_ID_NAME, span.getTraceId());
-		setHeader(request, Span.SPAN_ID_NAME, span.getSpanId());
-		if (!span.isExportable()) {
-			setHeader(request, Span.NOT_SAMPLED_NAME, "true");
-		}
-		setHeader(request, Span.SPAN_NAME_NAME, span.getName().toString());
-		setHeader(request, Span.PARENT_ID_NAME, getParentId(span));
-		setHeader(request, Span.PROCESS_ID_NAME, span.getProcessId());
+		enrichWithTraceHeaders(request, span);
 		publish(new ClientSentEvent(this, span));
 		return new TraceHttpResponse(this, execution.execute(request, body));
 	}
 
-	public void close() {
-		if (getCurrentSpan() == null) {
-			return;
-		}
-		publish(new ClientReceivedEvent(this, getCurrentSpan()));
-	}
-
-	private void publish(ApplicationEvent event) {
-		if (this.publisher != null) {
-			this.publisher.publishEvent(event);
-		}
-	}
-
-	private Long getParentId(Span span) {
-		return !span.getParents().isEmpty() ? span.getParents().get(0) : null;
-	}
-
-	public void setHeader(HttpRequest request, String name, String value) {
-		if (StringUtils.hasText(value) && !request.getHeaders().containsKey(name) && this.accessor.isTracing()) {
-			request.getHeaders().add(name, value);
-		}
-	}
-
-	public void setHeader(HttpRequest request, String name, Long value) {
-		if (value != null) {
-			setHeader(request, name, Span.toHex(value));
-		}
-	}
-
-	private Span getCurrentSpan() {
-		return this.accessor.getCurrentSpan();
-	}
-
 }

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2015 the original author or authors.
+ * Copyright 2013-2016 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.
@@ -29,10 +29,16 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.cloud.sleuth.SpanAccessor;
+import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.task.AsyncListenableTaskExecutor;
+import org.springframework.http.client.AsyncClientHttpRequestFactory;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
+import org.springframework.http.client.SimpleClientHttpRequestFactory;
+import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+import org.springframework.web.client.AsyncRestTemplate;
 import org.springframework.web.client.RestTemplate;
 
 /**
@@ -57,6 +63,29 @@ public class TraceWebClientAutoConfiguration {
 		return new RestTemplate();
 	}
 
+
+	@Bean
+	@ConditionalOnMissingBean
+	public AsyncClientHttpRequestFactory asyncClientHttpRequestFactory(Tracer tracer) {
+		SimpleClientHttpRequestFactory simpleClientHttpRequestFactory = new SimpleClientHttpRequestFactory();
+		simpleClientHttpRequestFactory.setTaskExecutor(asyncListenableTaskExecutor(tracer));
+		return simpleClientHttpRequestFactory;
+	}
+
+	private AsyncListenableTaskExecutor asyncListenableTaskExecutor(Tracer tracer) {
+		ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler();
+		threadPoolTaskScheduler.initialize();
+		return new TraceAsyncListenableTaskExecutor(threadPoolTaskScheduler, tracer);
+	}
+
+	@Bean
+	@ConditionalOnMissingBean
+	public AsyncRestTemplate asyncRestTemplate(SpanAccessor spanAccessor,
+			AsyncClientHttpRequestFactory asyncClientHttpRequestFactory, RestTemplate restTemplate) {
+		return new AsyncRestTemplate(new TraceAsyncClientHttpRequestFactoryWrapper(
+				spanAccessor, asyncClientHttpRequestFactory), restTemplate);
+	}
+
 	@Configuration
 	protected static class TraceInterceptorConfiguration {
 

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java

@@ -1,7 +1,8 @@
 package org.springframework.cloud.sleuth.instrument.web;
 
-import junitparams.JUnitParamsRunner;
-import junitparams.Parameters;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+
 import org.junit.Before;
 import org.junit.ClassRule;
 import org.junit.Rule;
@@ -11,8 +12,8 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.SpringApplicationConfiguration;
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.instrument.DefaultTestAutoConfiguration;
-import org.springframework.cloud.sleuth.instrument.web.common.HttpMockServer;
 import org.springframework.cloud.sleuth.instrument.web.common.AbstractMvcWiremockIntegrationTest;
+import org.springframework.cloud.sleuth.instrument.web.common.HttpMockServer;
 import org.springframework.context.annotation.Import;
 import org.springframework.http.MediaType;
 import org.springframework.scheduling.annotation.EnableAsync;
@@ -23,12 +24,18 @@ import org.springframework.test.web.servlet.request.MockMvcRequestBuilders;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RequestMethod;
 import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.client.AsyncRestTemplate;
 import org.springframework.web.client.RestTemplate;
 import org.springframework.web.context.request.async.WebAsyncTask;
 
-import java.util.concurrent.Callable;
+import junitparams.JUnitParamsRunner;
+import junitparams.Parameters;
 
-import static com.github.tomakehurst.wiremock.client.WireMock.*;
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.get;
+import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.matching;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlMatching;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static junitparams.JUnitParamsRunner.$;
 import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.asyncDispatch;
@@ -54,6 +61,13 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock
 		thenTraceIdHasBeenSetOnARequestHeader();
 	}
 
+	@Test
+	public void should_set_span_data_on_headers_when_sending_a_request_via_async_rest_template() throws Exception {
+		whenARequestIsSentToAAsyncRestTemplateEndpoint();
+
+		thenTraceIdHasBeenSetOnARequestHeader();
+	}
+
 	@Test
 	@Parameters
 	public void should_set_span_data_on_headers_via_aspect_in_asynchronous_call(String url) throws Exception {
@@ -66,6 +80,10 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock
 		return $("/callablePing", "/webAsyncTaskPing");
 	}
 
+	private void whenARequestIsSentToAAsyncRestTemplateEndpoint() throws Exception {
+		this.mockMvc.perform(MockMvcRequestBuilders.get("/asyncRestTemplate").accept(MediaType.TEXT_PLAIN)).andReturn();
+	}
+
 	private void whenARequestIsSentToASyncEndpoint() throws Exception {
 		this.mockMvc.perform(MockMvcRequestBuilders.get("/syncPing").accept(MediaType.TEXT_PLAIN)).andReturn();
 	}
@@ -95,6 +113,13 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock
 
 		@Autowired HttpMockServer httpMockServer;
 		@Autowired RestTemplate restTemplate;
+		@Autowired AsyncRestTemplate asyncRestTemplate;
+
+		@RequestMapping(value = "/asyncRestTemplate", method = RequestMethod.GET, produces = MediaType.TEXT_PLAIN_VALUE)
+		public String asyncRestTemplate()
+				throws ExecutionException, InterruptedException {
+			return callWiremockViaAsyncRestTemplateAndReturnOk();
+		}
 
 		@RequestMapping(value = "/syncPing", method = RequestMethod.GET, produces = MediaType.TEXT_PLAIN_VALUE)
 		public String syncPing() {
@@ -125,5 +150,11 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock
 			this.restTemplate.getForObject("http://localhost:" + this.httpMockServer.port(), String.class);
 			return "OK";
 		}
+
+		private String callWiremockViaAsyncRestTemplateAndReturnOk()
+				throws ExecutionException, InterruptedException {
+			this.asyncRestTemplate.getForEntity("http://localhost:" + this.httpMockServer.port(), String.class).get();
+			return "OK";
+		}
 	}
 }