From 964afe10f757cac1ad65a57f996379c9fe627fc2 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin@grzejszczak.pl>
Date: Fri, 16 Dec 2016 13:59:03 +0100
Subject: [PATCH] Premature async rest template (#479)

without this change the async rest template provides wrong value of the span duration
with this change the span is closed via a callback

fixes #475
---
 .../web/client/TraceAsyncRestTemplate.java    | 67 ++++++++++++++-----
 ...stTemplateTraceAspectIntegrationTests.java | 16 ++---
 2 files changed, 60 insertions(+), 23 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java
index 4db2d1567..8543f0849 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java
@@ -20,11 +20,13 @@ import java.net.URI;
 
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.util.ExceptionUtils;
 import org.springframework.core.task.AsyncListenableTaskExecutor;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.AsyncClientHttpRequestFactory;
 import org.springframework.http.client.ClientHttpRequestFactory;
 import org.springframework.util.concurrent.ListenableFuture;
+import org.springframework.util.concurrent.ListenableFutureCallback;
 import org.springframework.web.client.AsyncRequestCallback;
 import org.springframework.web.client.AsyncRestTemplate;
 import org.springframework.web.client.ResponseExtractor;
@@ -75,27 +77,62 @@ public class TraceAsyncRestTemplate extends AsyncRestTemplate {
 	protected <T> ListenableFuture<T> doExecute(URI url, HttpMethod method,
 			AsyncRequestCallback requestCallback, ResponseExtractor<T> responseExtractor)
 			throws RestClientException {
-		try {
-			return super.doExecute(url, method, requestCallback, responseExtractor);
-		} finally {
+		ListenableFuture<T> future = super.doExecute(url, method, requestCallback, responseExtractor);
+		Span span = this.tracer.getCurrentSpan();
+		future.addCallback(new TraceListenableFutureCallback<>(this.tracer, span));
+		// potential race can happen here
+		if (span != null && span.equals(this.tracer.getCurrentSpan())) {
+			this.tracer.detach(span);
+		}
+		return future;
+	}
+
+	private static class TraceListenableFutureCallback<T> implements ListenableFutureCallback<T> {
+
+		private final Tracer tracer;
+		private final Span parent;
+
+		private TraceListenableFutureCallback(Tracer tracer, Span parent) {
+			this.tracer = tracer;
+			this.parent = parent;
+		}
+
+		@Override
+		public void onFailure(Throwable ex) {
+			continueSpan();
+			this.tracer.addTag(Span.SPAN_ERROR_TAG_NAME, ExceptionUtils.getExceptionMessage(ex));
 			finish();
 		}
-	}
 
-	private void finish() {
-		if (!isTracing()) {
-			return;
+		@Override
+		public void onSuccess(T result) {
+			continueSpan();
+			finish();
+		}
+
+		private void continueSpan() {
+			this.tracer.continueSpan(this.parent);
+		}
+
+		private void finish() {
+			if (!isTracing()) {
+				return;
+			}
+			currentSpan().logEvent(Span.CLIENT_RECV);
+			this.tracer.close(currentSpan());
+		}
+
+		private Span currentSpan() {
+			return this.tracer.getCurrentSpan();
+		}
+
+		private boolean isTracing() {
+			return this.tracer.isTracing();
 		}
-		currentSpan().logEvent(Span.CLIENT_RECV);
-		this.tracer.close(this.currentSpan());
 	}
 
-	private Span currentSpan() {
-		return this.tracer.getCurrentSpan();
-	}
 
-	private boolean isTracing() {
-		return this.tracer.isTracing();
-	}
+
+
 
 }
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java
index 80fb7f54b..f99213850 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java
@@ -49,11 +49,9 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.
 @DirtiesContext
 public class RestTemplateTraceAspectIntegrationTests {
 
-	@Autowired
-	private WebApplicationContext context;
-
-	@Autowired
-	private AspectTestingController controller;
+	@Autowired WebApplicationContext context;
+	@Autowired AspectTestingController controller;
+	@Autowired Tracer tracer;
 
 	private MockMvc mockMvc;
 
@@ -61,12 +59,14 @@ public class RestTemplateTraceAspectIntegrationTests {
 	public void init() {
 		this.mockMvc = MockMvcBuilders.webAppContextSetup(this.context).build();
 		this.controller.reset();
-		TestSpanContextHolder.removeCurrentSpan();
+		ExceptionUtils.setFail(true);
 	}
 
+	@Before
 	@After
-	public void cleanup() {
-		TestSpanContextHolder.removeCurrentSpan();
+	public void verify() {
+		then(this.tracer.getCurrentSpan()).isNull();
+		then(ExceptionUtils.getLastException()).isNull();
 	}
 
 	@Test