Merge pull request #78 from spring-cloud/issues_#39_hystrix_concurrency_strategy

Issues #39 hystrix concurrency strategy

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

@@ -0,0 +1,19 @@
+package org.springframework.cloud.sleuth.instrument.hystrix;
+
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+import com.netflix.hystrix.HystrixCommand;
+
+@Configuration
+@ConditionalOnClass(HystrixCommand.class)
+@ConditionalOnProperty(value = "spring.sleuth.hystrix.strategy.enabled", matchIfMissing = true)
+public class SleuthHystrixAutoConfiguration {
+
+	@Bean SleuthHystrixConcurrencyStrategy sleuthHystrixConcurrencyStrategy(TraceManager traceManager) {
+		return new SleuthHystrixConcurrencyStrategy(traceManager);
+	}
+}

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

@@ -0,0 +1,24 @@
+package org.springframework.cloud.sleuth.instrument.hystrix;
+
+import java.util.concurrent.Callable;
+
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.cloud.sleuth.instrument.TraceCallable;
+
+import com.netflix.hystrix.strategy.HystrixPlugins;
+import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
+
+public class SleuthHystrixConcurrencyStrategy extends HystrixConcurrencyStrategy {
+
+	private final TraceManager traceManager;
+
+	public SleuthHystrixConcurrencyStrategy(TraceManager traceManager) {
+		this.traceManager = traceManager;
+		HystrixPlugins.getInstance().registerConcurrencyStrategy(this);
+	}
+
+	@Override
+	public <T> Callable<T> wrapCallable(Callable<T> callable) {
+		return new TraceCallable<>(traceManager, callable);
+	}
+}

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

@@ -26,6 +26,7 @@ import java.util.Map;
 
 import org.springframework.beans.factory.ObjectFactory;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
@@ -40,6 +41,8 @@ import org.springframework.cloud.sleuth.TraceAccessor;
 import org.springframework.cloud.sleuth.TraceManager;
 import org.springframework.cloud.sleuth.event.ClientReceivedEvent;
 import org.springframework.cloud.sleuth.event.ClientSentEvent;
+import org.springframework.cloud.sleuth.instrument.hystrix.SleuthHystrixConcurrencyStrategy;
+import org.springframework.cloud.sleuth.instrument.hystrix.SleuthHystrixAutoConfiguration;
 import org.springframework.context.ApplicationEvent;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.context.annotation.Bean;
@@ -67,6 +70,7 @@ import feign.hystrix.HystrixFeign;
 @ConditionalOnProperty(value = "spring.sleuth.feign.enabled", matchIfMissing = true)
 @ConditionalOnClass(Client.class)
 @AutoConfigureBefore(FeignAutoConfiguration.class)
+@AutoConfigureAfter(SleuthHystrixAutoConfiguration.class)
 public class TraceFeignClientAutoConfiguration {
 
 	@Autowired
@@ -81,6 +85,7 @@ public class TraceFeignClientAutoConfiguration {
 	@Bean
 	@Scope("prototype")
 	@ConditionalOnClass(HystrixCommand.class)
+	@ConditionalOnMissingBean(SleuthHystrixConcurrencyStrategy.class)
 	@ConditionalOnProperty(name = "feign.hystrix.enabled", matchIfMissing = true)
 	public Feign.Builder feignHystrixBuilder(TraceManager traceManager) {
 		return HystrixFeign.builder()

spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories

@@ -5,6 +5,7 @@ org.springframework.cloud.sleuth.log.SleuthLogAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.integration.TraceSpringIntegrationAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.async.AsyncCustomAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.async.AsyncDefaultAutoConfiguration,\
+org.springframework.cloud.sleuth.instrument.hystrix.SleuthHystrixAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.scheduling.TraceSchedulingAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.web.TraceWebAutoConfiguration,\
 org.springframework.cloud.sleuth.instrument.web.client.TraceWebClientAutoConfiguration,\

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

@@ -0,0 +1,16 @@
+package org.springframework.cloud.sleuth.assertions;
+
+import org.assertj.core.api.BDDAssertions;
+import org.springframework.cloud.sleuth.Span;
+
+public class SleuthAssertions extends BDDAssertions {
+
+	public static SpanAssert then(Span actual) {
+		return assertThat(actual);
+	}
+
+	public static SpanAssert assertThat(Span actual) {
+		return new SpanAssert(actual);
+	}
+
+}

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

@@ -0,0 +1,33 @@
+package org.springframework.cloud.sleuth.assertions;
+
+import java.util.Objects;
+
+import org.assertj.core.api.AbstractAssert;
+import org.springframework.cloud.sleuth.Span;
+
+public class SpanAssert extends AbstractAssert<SpanAssert, Span> {
+
+  public SpanAssert(Span actual) {
+    super(actual, SpanAssert.class);
+  }
+
+  public static SpanAssert then(Span actual) {
+    return new SpanAssert(actual);
+  }
+
+  public SpanAssert hasTraceIdEqualTo(String traceId) {
+    isNotNull();
+    if (!Objects.equals(actual.getTraceId(), traceId)) {
+      failWithMessage("Expected span's traceId to be <%s> but was <%s>", traceId, actual.getTraceId());
+    }
+    return this;
+  }
+
+  public SpanAssert hasNameNotEqualTo(String name) {
+    isNotNull();
+    if (Objects.equals(actual.getName(), name)) {
+      failWithMessage("Expected span's name not to be <%s> but was <%s>", name, actual.getName());
+    }
+    return this;
+  }
+}

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

@@ -0,0 +1,106 @@
+package org.springframework.cloud.sleuth.instrument.hystrix;
+
+import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.After;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.SpringApplicationConfiguration;
+import org.springframework.cloud.netflix.hystrix.EnableHystrix;
+import org.springframework.cloud.sleuth.Span;
+import org.springframework.cloud.sleuth.TraceManager;
+import org.springframework.cloud.sleuth.instrument.DefaultTestAutoConfiguration;
+import org.springframework.cloud.sleuth.trace.TraceContextHolder;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import com.jayway.awaitility.Awaitility;
+import com.netflix.hystrix.contrib.javanica.annotation.HystrixCommand;
+
+@RunWith(SpringJUnit4ClassRunner.class)
+@SpringApplicationConfiguration(classes = {
+		SpanPassingForHystrixViaAnnotationsITest.TestConfig.class })
+public class SpanPassingForHystrixViaAnnotationsITest {
+
+	@Autowired HystrixCommandInvocationSpanCatcher hystrixCommandInvocationSpanCatcher;
+	@Autowired TraceManager traceManager;
+
+	@Test
+	public void should_set_span_on_an_hystrix_command_annotated_method() {
+		Span span = givenASpanInCurrentThread();
+
+		whenHystrixCommandAnnotatedMethodGetsExecuted();
+
+		thenTraceIdIsPassedFromTheCurrentThreadToTheHystrixOne(span);
+	}
+
+	private Span givenASpanInCurrentThread() {
+		Span span = traceManager.startSpan("existing").getSpan();
+		traceManager.continueSpan(span);
+		return span;
+	}
+
+	private void whenHystrixCommandAnnotatedMethodGetsExecuted() {
+		hystrixCommandInvocationSpanCatcher.invokeLogicWrappedInHystrixCommand();
+	}
+
+	private void thenTraceIdIsPassedFromTheCurrentThreadToTheHystrixOne(final Span span) {
+		Awaitility.await().until(new Runnable() {
+			@Override
+			public void run() {
+				then(span)
+						.hasTraceIdEqualTo(hystrixCommandInvocationSpanCatcher.getTraceId())
+						.hasNameNotEqualTo(hystrixCommandInvocationSpanCatcher.getSpanName());
+			}
+		});
+	}
+
+	@After
+	public void cleanTrace() {
+		TraceContextHolder.removeCurrentTrace();
+	}
+
+	@DefaultTestAutoConfiguration
+	@EnableHystrix
+	@Configuration
+	static class TestConfig {
+
+		@Bean HystrixCommandInvocationSpanCatcher spanCatcher() {
+			return new HystrixCommandInvocationSpanCatcher();
+		}
+
+	}
+
+	static class HystrixCommandInvocationSpanCatcher {
+
+		AtomicReference<Span> spanCaughtFromHystrixThread;
+
+		@HystrixCommand
+		public void invokeLogicWrappedInHystrixCommand() {
+			spanCaughtFromHystrixThread = new AtomicReference<>(TraceContextHolder.getCurrentSpan());
+		}
+
+		public String getTraceId() {
+			if (spanCaughtFromHystrixThread == null ||
+					spanCaughtFromHystrixThread.get() == null ||
+					(spanCaughtFromHystrixThread.get() != null &&
+							spanCaughtFromHystrixThread.get().getTraceId() == null)) {
+				return null;
+			}
+			return spanCaughtFromHystrixThread.get().getTraceId();
+		}
+
+		public String getSpanName() {
+			if (spanCaughtFromHystrixThread == null ||
+					(spanCaughtFromHystrixThread.get() != null &&
+							spanCaughtFromHystrixThread.get().getName() == null)) {
+				return null;
+			}
+			return spanCaughtFromHystrixThread.get().getName();
+		}
+	}
+}

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

@@ -1,7 +1,7 @@
 
 package org.springframework.cloud.sleuth.instrument.web;
 
-import static org.assertj.core.api.BDDAssertions.then;
+import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
 
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -16,7 +16,6 @@ import org.springframework.cloud.sleuth.instrument.DefaultTestAutoConfiguration;
 import org.springframework.cloud.sleuth.trace.TraceContextHolder;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.EnableAspectJAutoProxy;
 import org.springframework.scheduling.annotation.Async;
 import org.springframework.scheduling.annotation.EnableAsync;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
@@ -25,43 +24,38 @@ import com.jayway.awaitility.Awaitility;
 
 @RunWith(SpringJUnit4ClassRunner.class)
 @SpringApplicationConfiguration(classes = {
-		TraceAsyncITest.CorrelationIdAsyncSpecConfiguration.class })
+		TraceAsyncITest.TraceAsyncITestConfiguration.class })
 public class TraceAsyncITest {
 
-	@Autowired
-	AsyncClass asyncClass;
-	@Autowired
-	AsyncDelegation asyncDelegation;
-	@Autowired
-	TraceManager traceManager;
+	@Autowired ClassPerformingAsyncLogic classPerformingAsyncLogic;
+	@Autowired TraceManager traceManager;
 
 	@Test
 	public void should_set_span_on_an_async_annotated_method() {
-		final Span span = givenASpanInCurrentThread();
+		Span span = givenASpanInCurrentThread();
 
 		whenAsyncProcessingTakesPlace();
 
-		thenSpanPutInTheAsyncThreadIsSameAs(span);
+		thenTraceIdIsPassedFromTheCurrentThreadToTheAsyncOne(span);
 	}
 
 	private Span givenASpanInCurrentThread() {
-		Span span = this.traceManager.startSpan("existing").getSpan();
-		this.traceManager.continueSpan(span);
+		Span span = traceManager.startSpan("existing").getSpan();
+		traceManager.continueSpan(span);
 		return span;
 	}
 
 	private void whenAsyncProcessingTakesPlace() {
-		this.asyncDelegation.doSthThatDelegatesToAsync();
+		classPerformingAsyncLogic.invokeAsynchronousLogic();
 	}
 
-	private void thenSpanPutInTheAsyncThreadIsSameAs(final Span span) {
+	private void thenTraceIdIsPassedFromTheCurrentThreadToTheAsyncOne(final Span span) {
 		Awaitility.await().until(new Runnable() {
 			@Override
 			public void run() {
-				then(span.getTraceId()).isNotNull()
-						.isEqualTo(TraceAsyncITest.this.asyncClass.getTraceId());
-				then(span.getName())
-						.isNotEqualTo(TraceAsyncITest.this.asyncClass.getSpanName());
+				then(span)
+						.hasTraceIdEqualTo(classPerformingAsyncLogic.getTraceId())
+						.hasNameNotEqualTo(classPerformingAsyncLogic.getSpanName());
 			}
 		});
 	}
@@ -73,57 +67,39 @@ public class TraceAsyncITest {
 
 	@DefaultTestAutoConfiguration
 	@EnableAsync
-	@EnableAspectJAutoProxy(proxyTargetClass = true)
 	@Configuration
-	public static class CorrelationIdAsyncSpecConfiguration {
+	static class TraceAsyncITestConfiguration {
 
 		@Bean
-		AsyncClass asyncClass() {
-			return new AsyncClass();
+		ClassPerformingAsyncLogic asyncClass() {
+			return new ClassPerformingAsyncLogic();
 		}
 
-		@Bean
-		AsyncDelegation asyncDelegation() {
-			return new AsyncDelegation(asyncClass());
-		}
 	}
 
-	public static class AsyncDelegation {
-
-		private final AsyncClass asyncClass;
-
-		public AsyncDelegation(AsyncClass asyncClass) {
-			this.asyncClass = asyncClass;
-		}
-
-		public void doSthThatDelegatesToAsync() {
-			this.asyncClass.doSth();
-		}
-	}
-
-	public static class AsyncClass {
+	static class ClassPerformingAsyncLogic {
 
 		AtomicReference<Span> span;
 
 		@Async
-		public void doSth() {
-			this.span = new AtomicReference<>(TraceContextHolder.getCurrentSpan());
+		public void invokeAsynchronousLogic() {
+			span = new AtomicReference<>(TraceContextHolder.getCurrentSpan());
 		}
 
 		public String getTraceId() {
-			if (this.span == null || (this.span.get() != null
-					&& this.span.get().getTraceId() == null)) {
+			if (span == null || (span.get() != null
+					&& span.get().getTraceId() == null)) {
 				return null;
 			}
-			return this.span.get().getTraceId();
+			return span.get().getTraceId();
 		}
 
 		public String getSpanName() {
-			if (this.span == null
-					|| (this.span.get() != null && this.span.get().getName() == null)) {
+			if (span == null
+					|| (span.get() != null && span.get().getName() == null)) {
 				return null;
 			}
-			return this.span.get().getName();
+			return span.get().getName();
 		}
 	}
 }