[#198] Fixed exception handling for TraceRestTemplateInterceptor

fixes #198

spring-cloud-sleuth-core/pom.xml

@@ -95,6 +95,12 @@
 			<version>0.7.1</version>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>com.squareup.okhttp3</groupId>
+			<artifactId>mockwebserver</artifactId>
+			<version>3.1.2</version>
+			<scope>test</scope>
+		</dependency>
 		<dependency>
 			<groupId>org.assertj</groupId>
 			<artifactId>assertj-core</artifactId>

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

@@ -112,7 +112,7 @@ abstract class AbstractTraceHttpRequestInterceptor
 		}
 	}
 
-	private Span currentSpan() {
+	protected Span currentSpan() {
 		return this.tracer.getCurrentSpan();
 	}
 

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

@@ -16,7 +16,10 @@
 package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.io.IOException;
+import java.lang.invoke.MethodHandles;
 
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.http.HttpRequest;
 import org.springframework.http.client.ClientHttpRequestExecution;
@@ -37,6 +40,8 @@ import org.springframework.http.client.ClientHttpResponse;
 public class TraceRestTemplateInterceptor extends AbstractTraceHttpRequestInterceptor
 		implements ClientHttpRequestInterceptor {
 
+	private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());
+
 	public TraceRestTemplateInterceptor(Tracer tracer) {
 		super(tracer);
 	}
@@ -49,7 +54,17 @@ public class TraceRestTemplateInterceptor extends AbstractTraceHttpRequestInterc
 			return execution.execute(request, body);
 		}
 		publishStartEvent(request);
-		return new TraceHttpResponse(this, execution.execute(request, body));
+		return response(request, body, execution);
+	}
+
+	private ClientHttpResponse response(HttpRequest request, byte[] body,
+			ClientHttpRequestExecution execution) throws IOException {
+		try {
+			return new TraceHttpResponse(this, execution.execute(request, body));
+		} catch (Exception e) {
+			this.tracer.close(currentSpan());
+			throw e;
+		}
 	}
 
 

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptorIntegrationTests.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 java.io.IOException;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Random;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.springframework.cloud.sleuth.DefaultSpanNamer;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.assertions.SleuthAssertions;
+import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
+import org.springframework.cloud.sleuth.trace.DefaultTracer;
+import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
+import org.springframework.context.support.StaticApplicationContext;
+import org.springframework.http.client.ClientHttpRequestFactory;
+import org.springframework.http.client.ClientHttpRequestInterceptor;
+import org.springframework.http.client.HttpComponentsClientHttpRequestFactory;
+import org.springframework.web.client.RestTemplate;
+
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.MockWebServer;
+import okhttp3.mockwebserver.SocketPolicy;
+
+/**
+ * @author Marcin Grzejszczak
+ */
+public class TraceRestTemplateInterceptorIntegrationTests {
+
+	@Rule public final MockWebServer mockWebServer = new MockWebServer();
+
+	private RestTemplate template = new RestTemplate(clientHttpRequestFactory());
+
+	private DefaultTracer tracer;
+
+	private StaticApplicationContext publisher = new StaticApplicationContext();
+
+	@Before
+	public void setup() {
+		this.publisher.refresh();
+		this.tracer = new DefaultTracer(new AlwaysSampler(), new Random(), this.publisher,
+				new DefaultSpanNamer());
+		this.template.setInterceptors(Arrays.<ClientHttpRequestInterceptor>asList(
+				new TraceRestTemplateInterceptor(this.tracer)));
+		TestSpanContextHolder.removeCurrentSpan();
+	}
+
+	@After
+	public void clean() throws IOException {
+		TestSpanContextHolder.removeCurrentSpan();
+	}
+
+	// issue #198
+	@Test
+	public void spanRemovedFromThreadUponException() throws IOException {
+		this.mockWebServer.enqueue(new MockResponse().setSocketPolicy(SocketPolicy.DISCONNECT_AT_START));
+		Span span = this.tracer.startTrace("new trace");
+
+		try {
+			this.template.getForEntity(
+					"http://localhost:" + this.mockWebServer.getPort() + "/exception",
+					Map.class).getBody();
+			Assert.fail("should throw an exception");
+		} catch (RuntimeException e) {
+			SleuthAssertions.then(e).hasRootCauseInstanceOf(IOException.class);
+		}
+
+		SleuthAssertions.then(this.tracer.getCurrentSpan()).isEqualTo(span);
+		this.tracer.close(span);
+	}
+
+	private ClientHttpRequestFactory clientHttpRequestFactory() {
+		HttpComponentsClientHttpRequestFactory factory = new HttpComponentsClientHttpRequestFactory();
+		factory.setReadTimeout(100);
+		factory.setConnectTimeout(100);
+		return factory;
+	}
+
+}

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

@@ -22,10 +22,12 @@ import java.util.Map;
 import java.util.Random;
 
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.cloud.sleuth.DefaultSpanNamer;
 import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.assertions.SleuthAssertions;
 import org.springframework.cloud.sleuth.sampler.AlwaysSampler;
 import org.springframework.cloud.sleuth.trace.DefaultTracer;
 import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
@@ -55,17 +57,17 @@ public class TraceRestTemplateInterceptorTests {
 	private RestTemplate template = new RestTemplate(
 			new MockMvcClientHttpRequestFactory(this.mockMvc));
 
-	private DefaultTracer traces;
+	private DefaultTracer tracer;
 
 	private StaticApplicationContext publisher = new StaticApplicationContext();
 
 	@Before
 	public void setup() {
 		this.publisher.refresh();
-		this.traces = new DefaultTracer(new AlwaysSampler(), new Random(), this.publisher,
+		this.tracer = new DefaultTracer(new AlwaysSampler(), new Random(), this.publisher,
 				new DefaultSpanNamer());
 		this.template.setInterceptors(Arrays.<ClientHttpRequestInterceptor>asList(
-				new TraceRestTemplateInterceptor(this.traces)));
+				new TraceRestTemplateInterceptor(this.tracer)));
 		TestSpanContextHolder.removeCurrentSpan();
 	}
 
@@ -76,7 +78,7 @@ public class TraceRestTemplateInterceptorTests {
 
 	@Test
 	public void headersAddedWhenTracing() {
-		this.traces.continueSpan(Span.builder().traceId(1L).spanId(2L).parent(3L).build());
+		this.tracer.continueSpan(Span.builder().traceId(1L).spanId(2L).parent(3L).build());
 		@SuppressWarnings("unchecked")
 		Map<String, String> headers = this.template.getForEntity("/", Map.class)
 				.getBody();
@@ -87,7 +89,7 @@ public class TraceRestTemplateInterceptorTests {
 
 	@Test
 	public void notSampledHeaderAddedWhenNotExportable() {
-		this.traces.continueSpan(Span.builder().traceId(1L).spanId(2L).exportable(false).build());
+		this.tracer.continueSpan(Span.builder().traceId(1L).spanId(2L).exportable(false).build());
 		@SuppressWarnings("unchecked")
 		Map<String, String> headers = this.template.getForEntity("/", Map.class)
 				.getBody();
@@ -104,8 +106,25 @@ public class TraceRestTemplateInterceptorTests {
 		assertFalse("Wrong headers: " + headers, headers.containsKey(Span.SPAN_ID_NAME));
 	}
 
+	// issue #198
+	@Test
+	public void spanRemovedFromThreadUponException() {
+		Span span = this.tracer.startTrace("new trace");
+
+		try {
+			this.template.getForEntity("/exception", Map.class).getBody();
+			Assert.fail("should throw an exception");
+		} catch (RuntimeException e) {
+			SleuthAssertions.then(e).hasMessage("500 Internal Server Error");
+		}
+
+		SleuthAssertions.then(this.tracer.getCurrentSpan()).isEqualTo(span);
+		this.tracer.close(span);
+	}
+
 	@RestController
 	public static class TestController {
+
 		@RequestMapping("/")
 		public Map<String, String> home(@RequestHeader HttpHeaders headers) {
 			Map<String, String> map = new HashMap<String, String>();
@@ -114,6 +133,11 @@ public class TraceRestTemplateInterceptorTests {
 			return map;
 		}
 
+		@RequestMapping("/exception")
+		public Map<String, String> exception() {
+			throw new RuntimeException("foo");
+		}
+
 		private void addHeaders(Map<String, String> map, HttpHeaders headers,
 				String... names) {
 			if (headers != null) {