Ensure span is cleaned up if Feign client throws a non retryable

The retryable case was already covered and tested but if the exception
is not retryable, we need to rethrow and clean up the thread state.

Fixes gh-240

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/hystrix/SleuthHystrixConcurrencyStrategy.java

@@ -33,8 +33,8 @@ import com.netflix.hystrix.strategy.properties.HystrixPropertiesStrategy;
 
 /**
  * A {@link HystrixConcurrencyStrategy} that wraps a {@link Callable} in a
- * {@link Callable} that either starts a new span or continues one
- * if the tracing was already running before the command was executed.
+ * {@link Callable} that either starts a new span or continues one if the tracing was
+ * already running before the command was executed.
  *
  * @author Marcin Grzejszczak
  * @since 1.0.0
@@ -42,7 +42,8 @@ import com.netflix.hystrix.strategy.properties.HystrixPropertiesStrategy;
 public class SleuthHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy {
 
 	private static final String HYSTRIX_COMPONENT = "hystrix";
-	private static final Log log = LogFactory.getLog(SleuthHystrixConcurrencyStrategy.class);
+	private static final Log log = LogFactory
+			.getLog(SleuthHystrixConcurrencyStrategy.class);
 
 	private final Tracer tracer;
 	private final TraceKeys traceKeys;
@@ -53,6 +54,10 @@ public class SleuthHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy
 		this.traceKeys = traceKeys;
 		try {
 			this.delegate = HystrixPlugins.getInstance().getConcurrencyStrategy();
+			if (this.delegate instanceof SleuthHystrixConcurrencyStrategy) {
+				// Welcome to singleton hell...
+				return;
+			}
 			HystrixCommandExecutionHook commandExecutionHook = HystrixPlugins
 					.getInstance().getCommandExecutionHook();
 			HystrixEventNotifier eventNotifier = HystrixPlugins.getInstance()
@@ -65,33 +70,37 @@ public class SleuthHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy
 					propertiesStrategy);
 			HystrixPlugins.reset();
 			HystrixPlugins.getInstance().registerConcurrencyStrategy(this);
-			HystrixPlugins.getInstance().registerCommandExecutionHook(commandExecutionHook);
+			HystrixPlugins.getInstance()
+					.registerCommandExecutionHook(commandExecutionHook);
 			HystrixPlugins.getInstance().registerEventNotifier(eventNotifier);
 			HystrixPlugins.getInstance().registerMetricsPublisher(metricsPublisher);
 			HystrixPlugins.getInstance().registerPropertiesStrategy(propertiesStrategy);
 		}
 		catch (Exception e) {
-			log.error(
-					"Failed to register Sleuth Hystrix Concurrency Strategy", e);
+			log.error("Failed to register Sleuth Hystrix Concurrency Strategy", e);
 		}
 	}
 
 	private void logCurrentStateOfHysrixPlugins(HystrixEventNotifier eventNotifier,
 			HystrixMetricsPublisher metricsPublisher,
 			HystrixPropertiesStrategy propertiesStrategy) {
-		log.debug("Current Hystrix plugins configuration is ["
-				+ "concurrencyStrategy [" + this.delegate + "],"
-				+ "eventNotifier [" + eventNotifier + "],"
-				+ "metricPublisher [" + metricsPublisher + "],"
-				+ "propertiesStrategy [" + propertiesStrategy + "],"
-				+ "]");
+		log.debug("Current Hystrix plugins configuration is [" + "concurrencyStrategy ["
+				+ this.delegate + "]," + "eventNotifier [" + eventNotifier + "],"
+				+ "metricPublisher [" + metricsPublisher + "]," + "propertiesStrategy ["
+				+ propertiesStrategy + "]," + "]");
 		log.debug("Registering Sleuth Hystrix Concurrency Strategy.");
 	}
 
 	@Override
 	public <T> Callable<T> wrapCallable(Callable<T> callable) {
-		Callable<T> wrappedCallable = this.delegate != null ?
-				this.delegate.wrapCallable(callable) : callable;
+		if (callable instanceof HystrixTraceCallable) {
+			return callable;
+		}
+		Callable<T> wrappedCallable = this.delegate != null
+				? this.delegate.wrapCallable(callable) : callable;
+		if (wrappedCallable instanceof HystrixTraceCallable) {
+			return wrappedCallable;
+		}
 		return new HystrixTraceCallable<>(this.tracer, this.traceKeys, wrappedCallable);
 	}
 
@@ -103,7 +112,8 @@ public class SleuthHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy
 		private Callable<S> callable;
 		private Span parent;
 
-		public HystrixTraceCallable(Tracer tracer, TraceKeys traceKeys, Callable<S> callable) {
+		public HystrixTraceCallable(Tracer tracer, TraceKeys traceKeys,
+				Callable<S> callable) {
 			this.tracer = tracer;
 			this.traceKeys = traceKeys;
 			this.callable = callable;
@@ -120,8 +130,10 @@ public class SleuthHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy
 			else {
 				span = this.tracer.createSpan(HYSTRIX_COMPONENT);
 				this.tracer.addTag(Span.SPAN_LOCAL_COMPONENT_TAG_NAME, HYSTRIX_COMPONENT);
-				this.tracer.addTag(this.traceKeys.getAsync().getPrefix() +
-						this.traceKeys.getAsync().getThreadNameKey(), Thread.currentThread().getName());
+				this.tracer.addTag(
+						this.traceKeys.getAsync().getPrefix()
+								+ this.traceKeys.getAsync().getThreadNameKey(),
+						Thread.currentThread().getName());
 				created = true;
 			}
 			try {

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

@@ -24,10 +24,11 @@ import org.springframework.cloud.sleuth.Tracer;
 import feign.Client;
 import feign.Request;
 import feign.Response;
+import feign.RetryableException;
 
 /**
- * A Feign Client that closes a Span if there is no response body.
- * In other cases Span will not get closed cause the Decoder will not get called
+ * A Feign Client that closes a Span if there is no response body. In other cases Span
+ * will get closed because the Decoder will be called
  *
  * @author Marcin Grzejszczak
  *
@@ -49,8 +50,20 @@ final class TraceFeignClient extends FeignEventPublisher implements Client {
 
 	@Override
 	public Response execute(Request request, Request.Options options) throws IOException {
-		Response response = this.delegate.execute(request, options);
-		if (response.body() == null || (response.body() != null
+		Response response = null;
+		try {
+			response = this.delegate.execute(request, options);
+		}
+		catch (RetryableException | IOException e) {
+			// IOException will be wrapped into a RetryableException in the caller
+			throw e;
+		}
+		catch (RuntimeException e) {
+			// Any other exception is going to be propagated so we need to tidy up
+			finish();
+			throw e;
+		}
+		if (response != null && response.body() == null || (response.body() != null
 				&& Objects.equals(response.body().length(), 0))) {
 			finish();
 		}

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

@@ -22,10 +22,10 @@ import feign.RetryableException;
 import feign.Retryer;
 
 /**
- * Execution of this retryer means that an exception occurred while trying to
- * send the request. In that case we need to put information about this span
- * into the {@link FeignRequestContext} in order for the {@link feign.RequestInterceptor}
- * to know that it should be continued or a new one should be created.
+ * Execution of this retryer means that an exception occurred while trying to send the
+ * request. In that case we need to put information about this span into the
+ * {@link FeignRequestContext} in order for the {@link feign.RequestInterceptor} to know
+ * that it should be continued or a new one should be created.
  *
  * @author Marcin Grzejszczak
  *
@@ -34,7 +34,8 @@ import feign.Retryer;
 final class TraceFeignRetryer implements Retryer {
 
 	private final Tracer tracer;
-	private final FeignRequestContext feignRequestContext = FeignRequestContext.getInstance();
+	private final FeignRequestContext feignRequestContext = FeignRequestContext
+			.getInstance();
 	private final Retryer delegate;
 
 	TraceFeignRetryer(Tracer tracer) {

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/util/ExceptionUtils.java

@@ -19,11 +19,11 @@ package org.springframework.cloud.sleuth.util;
 import org.apache.commons.logging.Log;
 
 /**
- * Utility class for logging exceptions. Useful for test purposes -
- * when a warning message should be presented an exception can be thrown.
+ * Utility class for logging exceptions. Useful for test purposes - when a warning message
+ * should be presented an exception can be thrown.
  * <p>
- * The purpose of this class is not to throw exceptions from the user's code
- * when there are some issues with tracing.
+ * The purpose of this class is not to throw exceptions from the user's code when there
+ * are some issues with tracing.
  *
  * @author Spencer Gibb
  * @since 1.0.0
@@ -32,18 +32,27 @@ public final class ExceptionUtils {
 	private static final Log log = org.apache.commons.logging.LogFactory
 			.getLog(ExceptionUtils.class);
 	private static boolean fail = false;
+	private static Exception lastException = null;
 
 	private ExceptionUtils() {
 		throw new IllegalStateException("Utility class can't be instantiated");
 	}
 
 	public static void warn(String msg) {
-		if (fail) {
-			throw new IllegalStateException(msg);
-		}
 		log.warn(msg);
+		if (fail) {
+			IllegalStateException exception = new IllegalStateException(msg);
+			ExceptionUtils.lastException = exception;
+			throw exception;
+		}
 	}
+
+	public static Exception getLastException() {
+		return ExceptionUtils.lastException;
+	}
+
 	public static void setFail(boolean fail) {
 		ExceptionUtils.fail = fail;
+		ExceptionUtils.lastException = null;
 	}
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/AdhocTestSuite.java

@@ -20,8 +20,8 @@ import org.junit.Ignore;
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
 import org.junit.runners.Suite.SuiteClasses;
-import org.springframework.cloud.sleuth.instrument.hystrix.HystrixAnnotationsIntegrationTests;
-import org.springframework.cloud.sleuth.instrument.hystrix.TraceCommandTests;
+import org.springframework.cloud.sleuth.instrument.web.TraceFilterCustomExtractorTests;
+import org.springframework.cloud.sleuth.instrument.web.client.WebClientExceptionTests;
 
 /**
  * A test suite for probing weird ordering problems in the tests.
@@ -29,8 +29,7 @@ import org.springframework.cloud.sleuth.instrument.hystrix.TraceCommandTests;
  * @author Dave Syer
  */
 @RunWith(Suite.class)
-@SuiteClasses({ HystrixAnnotationsIntegrationTests.class,
-		TraceCommandTests.class })
+@SuiteClasses({ TraceFilterCustomExtractorTests.class, WebClientExceptionTests.class })
 @Ignore
 public class AdhocTestSuite {
 

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

@@ -0,0 +1,144 @@
+/*
+ * 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 static junitparams.JUnitParamsRunner.$;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.junit.Assert.assertThat;
+
+import java.io.IOException;
+import java.util.Map;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.boot.test.WebIntegrationTest;
+import org.springframework.cloud.client.discovery.EnableDiscoveryClient;
+import org.springframework.cloud.client.loadbalancer.LoadBalanced;
+import org.springframework.cloud.netflix.feign.EnableFeignClients;
+import org.springframework.cloud.netflix.feign.FeignClient;
+import org.springframework.cloud.netflix.ribbon.RibbonClient;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.Tracer;
+import org.springframework.cloud.sleuth.assertions.SleuthAssertions;
+import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
+import org.springframework.cloud.sleuth.util.ExceptionUtils;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.http.ResponseEntity;
+import org.springframework.test.context.junit4.rules.SpringClassRule;
+import org.springframework.test.context.junit4.rules.SpringMethodRule;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestMethod;
+import org.springframework.web.client.RestTemplate;
+
+import junitparams.JUnitParamsRunner;
+import junitparams.Parameters;
+
+@RunWith(JUnitParamsRunner.class)
+@SpringApplicationConfiguration(classes = {
+		WebClientDiscoveryExceptionTests.TestConfiguration.class })
+@WebIntegrationTest(value = {
+		"spring.application.name=exceptionservice" }, randomPort = true)
+public class WebClientDiscoveryExceptionTests {
+
+	@ClassRule
+	public static final SpringClassRule SCR = new SpringClassRule();
+	@Rule
+	public final SpringMethodRule springMethodRule = new SpringMethodRule();
+
+	@Autowired
+	TestFeignInterfaceWithException testFeignInterfaceWithException;
+	@Autowired
+	@LoadBalanced
+	RestTemplate template;
+	@Autowired
+	Tracer tracer;
+
+	@Before
+	public void open() {
+		TestSpanContextHolder.removeCurrentSpan();
+		ExceptionUtils.setFail(true);
+	}
+
+	@After
+	public void close() {
+		ExceptionUtils.setFail(false);
+		TestSpanContextHolder.removeCurrentSpan();
+	}
+
+	// issue #240
+	@Test
+	@Parameters
+	public void shouldCloseSpanUponException(ResponseEntityProvider provider)
+			throws IOException {
+		Span span = this.tracer.createSpan("new trace");
+
+		try {
+			provider.get(this);
+			Assert.fail("should throw an exception");
+		}
+		catch (RuntimeException e) {
+		}
+
+		assertThat(ExceptionUtils.getLastException(), is(nullValue()));
+
+		SleuthAssertions.then(this.tracer.getCurrentSpan()).isEqualTo(span);
+		this.tracer.close(span);
+	}
+
+	Object[] parametersForShouldCloseSpanUponException() {
+		return $(
+				(ResponseEntityProvider) (tests) -> tests.testFeignInterfaceWithException
+						.shouldFailToConnect(),
+				(ResponseEntityProvider) (tests) -> tests.template
+						.getForEntity("http://exceptionservice/", Map.class));
+	}
+
+	@FeignClient("exceptionservice")
+	public interface TestFeignInterfaceWithException {
+		@RequestMapping(method = RequestMethod.GET, value = "/")
+		ResponseEntity<String> shouldFailToConnect();
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	@EnableDiscoveryClient
+	@EnableFeignClients
+	@RibbonClient("exceptionservice")
+	public static class TestConfiguration {
+
+		@LoadBalanced
+		@Bean
+		public RestTemplate restTemplate() {
+			return new RestTemplate();
+		}
+	}
+
+	@FunctionalInterface
+	interface ResponseEntityProvider {
+		ResponseEntity<?> get(WebClientDiscoveryExceptionTests webClientTests);
+	}
+}

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

@@ -17,6 +17,9 @@
 package org.springframework.cloud.sleuth.instrument.web.client;
 
 import static junitparams.JUnitParamsRunner.$;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.nullValue;
+import static org.junit.Assert.assertThat;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -24,6 +27,7 @@ import java.util.Map;
 
 import org.junit.After;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
@@ -40,6 +44,7 @@ import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.Tracer;
 import org.springframework.cloud.sleuth.assertions.SleuthAssertions;
 import org.springframework.cloud.sleuth.trace.TestSpanContextHolder;
+import org.springframework.cloud.sleuth.util.ExceptionUtils;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.ResponseEntity;
@@ -76,8 +81,15 @@ public class WebClientExceptionTests {
 	@Autowired
 	Tracer tracer;
 
+	@Before
+	public void open() {
+		TestSpanContextHolder.removeCurrentSpan();
+		ExceptionUtils.setFail(true);
+	}
+
 	@After
 	public void close() {
+		ExceptionUtils.setFail(false);
 		TestSpanContextHolder.removeCurrentSpan();
 	}
 
@@ -96,6 +108,8 @@ public class WebClientExceptionTests {
 			SleuthAssertions.then(e).hasRootCauseInstanceOf(IOException.class);
 		}
 
+		assertThat(ExceptionUtils.getLastException(), is(nullValue()));
+
 		SleuthAssertions.then(this.tracer.getCurrentSpan()).isEqualTo(span);
 		this.tracer.close(span);
 	}