From b4ecb6f9576c141a099f2f71d5f472685d1ae4b2 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <mgrzejszczak@pivotal.io>
Date: Tue, 21 Nov 2017 19:07:06 +0100
Subject: [PATCH] Ensure we're not tracing the traced feign clients

fixes #791
---
 .../web/client/feign/TraceFeignAspect.java    | 16 ++--
 .../client/feign/TraceFeignAspectTests.java   | 77 +++++++++++++++++++
 2 files changed, 88 insertions(+), 5 deletions(-)
 create mode 100644 spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspectTests.java

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspect.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspect.java
index 0486d4575..829decdbb 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspect.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspect.java
@@ -16,6 +16,8 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client.feign;
 
+import java.io.IOException;
+
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
@@ -41,13 +43,17 @@ class TraceFeignAspect {
 
 	@Around("execution (* feign.Client.*(..)) && !within(is(FinalType))")
 	public Object feignClientWasCalled(final ProceedingJoinPoint pjp) throws Throwable {
-		Object[] args = pjp.getArgs();
-		Request request = (Request) args[0];
-		Request.Options options = (Request.Options) args[1];
 		Object bean = pjp.getTarget();
-		if (!(bean instanceof TraceFeignClient)) {
-			return new TraceFeignClient(this.beanFactory, (Client) bean).execute(request, options);
+		if (!(bean instanceof TraceFeignClient) && !(bean instanceof TraceLoadBalancerFeignClient)) {
+			return executeTraceFeignClient(bean, pjp);
 		}
 		return pjp.proceed();
 	}
+
+	Object executeTraceFeignClient(Object bean, ProceedingJoinPoint pjp) throws IOException {
+		Object[] args = pjp.getArgs();
+		Request request = (Request) args[0];
+		Request.Options options = (Request.Options) args[1];
+		return new TraceFeignClient(this.beanFactory, (Client) bean).execute(request, options);
+	}
 }
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspectTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspectTests.java
new file mode 100644
index 000000000..98d231c74
--- /dev/null
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/feign/TraceFeignAspectTests.java
@@ -0,0 +1,77 @@
+package org.springframework.cloud.sleuth.instrument.web.client.feign;
+
+import java.io.IOException;
+import java.nio.charset.Charset;
+import java.util.HashMap;
+
+import feign.Client;
+import feign.Request;
+import org.aspectj.lang.ProceedingJoinPoint;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import org.springframework.beans.factory.BeanFactory;
+
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+/**
+ * @author Marcin Grzejszczak
+ */
+@RunWith(MockitoJUnitRunner.class)
+public class TraceFeignAspectTests {
+	
+	@Mock BeanFactory beanFactory;
+	@Mock Client client;
+	@Mock ProceedingJoinPoint pjp;
+	@Mock TraceLoadBalancerFeignClient traceLoadBalancerFeignClient;
+	TraceFeignAspect traceFeignAspect;
+	
+	@Before
+	public void setup() {
+		stubPjp();
+		this.traceFeignAspect = new TraceFeignAspect(this.beanFactory) {
+			@Override Object executeTraceFeignClient(Object bean, ProceedingJoinPoint pjp) throws IOException {
+				return null;
+			}
+		};
+	}
+
+	private void stubPjp() {
+		Request request = Request.create("foo", "bar", new HashMap<>(), new byte[] {}, Charset
+				.defaultCharset());
+		Request.Options options = new Request.Options();
+		given(this.pjp.getArgs()).willReturn(new Object[] {request, options} );
+	}
+
+	@Test 
+	public void should_wrap_feign_client_in_trace_representation() throws Throwable {
+		given(this.pjp.getTarget()).willReturn(this.client);
+		
+		this.traceFeignAspect.feignClientWasCalled(this.pjp);
+
+		verify(this.pjp, never()).proceed();
+	}
+	
+	@Test 
+	public void should_not_wrap_traced_feign_client_in_trace_representation() throws Throwable {
+		given(this.pjp.getTarget()).willReturn(new TraceFeignClient(this.beanFactory, this.client));
+
+		this.traceFeignAspect.feignClientWasCalled(this.pjp);
+
+		verify(this.pjp).proceed();
+	}
+	
+	@Test 
+	public void should_not_wrap_traced_load_balancer_feign_client_in_trace_representation() throws Throwable {
+		given(this.pjp.getTarget()).willReturn(this.traceLoadBalancerFeignClient);
+
+		this.traceFeignAspect.feignClientWasCalled(this.pjp);
+
+		verify(this.pjp).proceed();
+	}
+
+}
\ No newline at end of file