Adds support for Hystrix fallback methods; fixes gh-1084

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

@@ -16,6 +16,8 @@
 
 package org.springframework.cloud.sleuth.instrument.hystrix;
 
+import java.util.concurrent.atomic.AtomicReference;
+
 import brave.Span;
 import brave.Tracer;
 import com.netflix.hystrix.HystrixCommand;
@@ -36,30 +38,53 @@ public abstract class TraceCommand<R> extends HystrixCommand<R> {
 	private static final String COMMAND_KEY = "commandKey";
 	private static final String COMMAND_GROUP_KEY = "commandGroup";
 	private static final String THREAD_POOL_KEY = "threadPoolKey";
+	private static final String FALLBACK_METHOD_NAME_KEY = "fallbackMethodName";
 
 	private final Tracer tracer;
-	private final Span span;
+	private final AtomicReference<Span> span;
 
 	protected TraceCommand(Tracer tracer, Setter setter) {
 		super(setter);
 		this.tracer = tracer;
-		this.span = this.tracer.nextSpan();
+		this.span = new AtomicReference<>(this.tracer.nextSpan());
 	}
 
 	@Override
 	protected R run() throws Exception {
 		String commandKeyName = getCommandKey().name();
-		Span span = this.span.name(commandKeyName);
+		Span span = this.span.get().name(commandKeyName);
 		span.tag(COMMAND_KEY, commandKeyName);
 		span.tag(COMMAND_GROUP_KEY, getCommandGroup().name());
 		span.tag(THREAD_POOL_KEY, getThreadPoolKey().name());
+		Throwable throwable = null;
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span.start())) {
 			return doRun();
-		}
-		finally {
-			span.finish();
+		} catch (Throwable t) {
+			throwable = t;
+			throw t;
+		} finally {
+			if (throwable == null) {
+				span.finish();
+				this.span.set(null);
+			}
+			// else there will be fallback
 		}
 	}
 
 	public abstract R doRun() throws Exception;
+
+	@Override protected R getFallback() {
+		Span span = this.span.get();
+		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
+			span.tag(FALLBACK_METHOD_NAME_KEY, getFallbackMethodName());
+			return doGetFallback();
+		} finally {
+			span.finish();
+			this.span.set(null);
+		}
+	}
+
+	public R doGetFallback() {
+		return super.getFallback();
+	}
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/hystrix/TraceCommandTests.java

@@ -17,16 +17,19 @@
 package org.springframework.cloud.sleuth.instrument.hystrix;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 import brave.Span;
 import brave.Tracer;
 import brave.Tracing;
 import brave.propagation.StrictCurrentTraceContext;
+import brave.sampler.Sampler;
 import com.netflix.hystrix.HystrixCommand;
 import com.netflix.hystrix.HystrixCommandKey;
 import com.netflix.hystrix.HystrixCommandProperties;
 import com.netflix.hystrix.HystrixThreadPoolProperties;
 import com.netflix.hystrix.strategy.HystrixPlugins;
+import org.assertj.core.api.BDDAssertions;
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.cloud.sleuth.util.ArrayListSpanReporter;
@@ -41,6 +44,7 @@ public class TraceCommandTests {
 	Tracing tracing = Tracing.newBuilder()
 			.currentTraceContext(new StrictCurrentTraceContext())
 			.spanReporter(this.reporter)
+			.sampler(Sampler.ALWAYS_SAMPLE)
 			.build();
 	Tracer tracer = this.tracing.tracer();
 
@@ -120,6 +124,42 @@ public class TraceCommandTests {
 		then(resultFromHystrixCommand).isEqualTo(resultFromTraceCommand);
 	}
 
+	@Test
+	public void should_pass_tracing_information_when_using_Hystrix_commands_with_fallback() {
+		Tracer tracer = this.tracer;
+		AtomicReference<Span> spanBeforeThrowingException = new AtomicReference<>();
+		HystrixCommand.Setter setter = withGroupKey(asKey("group"))
+				.andCommandKey(HystrixCommandKey.Factory.asKey("command"));
+		TraceCommand<Span> traceCommand = new TraceCommand<Span>(tracer, setter) {
+			@Override
+			public Span doRun() throws Exception {
+				spanBeforeThrowingException.set(tracer.currentSpan());
+				throw new FooException();
+			}
+
+			@Override public Span doGetFallback() {
+				return tracer.currentSpan();
+			}
+
+			@Override protected String getFallbackMethodName() {
+				return super.getFallbackMethodName() + "_foobar";
+			}
+		};
+
+		Span span = whenCommandIsExecuted(traceCommand);
+
+		BDDAssertions.then(span.context().traceIdString())
+				.isEqualTo(spanBeforeThrowingException.get().context().traceIdString());
+		List<zipkin2.Span> spans = this.reporter.getSpans();
+		then(spans).hasSize(1);
+		then(spans.get(0).traceId()).isEqualTo(span.context().traceIdString());
+		then(spans.get(0).tags())
+				.containsEntry("commandKey", "command")
+				.containsEntry("commandGroup", "group")
+				.containsEntry("threadPoolKey", "group")
+				.containsEntry("fallbackMethodName", "getFallback_foobar");
+	}
+
 	private String someLogic(){
 		return "some logic";
 	}
@@ -146,4 +186,6 @@ public class TraceCommandTests {
 	private Span givenACommandWasExecuted(TraceCommand<Span> command) {
 		return whenCommandIsExecuted(command);
 	}
-}
+}
+
+class FooException extends RuntimeException {}