Gh 1528 instrument feign sc loadbalancer client (#1531)

* Add optional SC LoadBalancer dependency. Create TraceFeignBlockingLoadBalancerClient.

* Instrument Spring Cloud LoadBalancer in TraceFeignObjectWrapper.

* Do not use class in field in order to avoid NoClassDefFound.

fixes #1528

spring-cloud-sleuth-core/pom.xml

@@ -156,6 +156,11 @@
 			<artifactId>spring-cloud-context</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework.cloud</groupId>
+			<artifactId>spring-cloud-starter-loadbalancer</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>com.netflix.hystrix</groupId>
 			<artifactId>hystrix-core</artifactId>

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

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2013-2019 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
+ *
+ *      https://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.feign;
+
+import java.io.IOException;
+import java.util.HashMap;
+
+import brave.Span;
+import brave.Tracer;
+import brave.http.HttpTracing;
+import com.netflix.client.ClientException;
+import feign.Client;
+import feign.Request;
+import feign.Response;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
+import org.springframework.cloud.openfeign.loadbalancer.FeignBlockingLoadBalancerClient;
+
+/**
+ * A trace representation of {@link FeignBlockingLoadBalancerClient}. Needed due to casts
+ * in {@link org.springframework.cloud.openfeign.FeignClientFactoryBean}. Based on
+ * {@link TraceLoadBalancerFeignClient}.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 2.2.0
+ * @see TraceLoadBalancerFeignClient
+ * @see FeignBlockingLoadBalancerClient
+ */
+public class TraceFeignBlockingLoadBalancerClient
+		extends FeignBlockingLoadBalancerClient {
+
+	private static final Log LOG = LogFactory
+			.getLog(TraceFeignBlockingLoadBalancerClient.class);
+
+	private final BeanFactory beanFactory;
+
+	Tracer tracer;
+
+	HttpTracing httpTracing;
+
+	TracingFeignClient tracingFeignClient;
+
+	TraceFeignBlockingLoadBalancerClient(Client delegate,
+			BlockingLoadBalancerClient loadBalancerClient, BeanFactory beanFactory) {
+		super(delegate, loadBalancerClient);
+		this.beanFactory = beanFactory;
+	}
+
+	@Override
+	public Response execute(Request request, Request.Options options) throws IOException {
+		if (LOG.isDebugEnabled()) {
+			LOG.debug("Before send");
+		}
+		Response response = null;
+		Span fallbackSpan = tracer().nextSpan().start();
+		try {
+			response = super.execute(request, options);
+			if (LOG.isDebugEnabled()) {
+				LOG.debug("After receive");
+			}
+			return response;
+		}
+		catch (Exception e) {
+			if (LOG.isDebugEnabled()) {
+				LOG.debug("Exception thrown", e);
+			}
+			if (e instanceof IOException || e.getCause() != null
+					&& e.getCause() instanceof ClientException
+					&& ((ClientException) e.getCause())
+							.getErrorType() == ClientException.ErrorType.GENERAL) {
+				if (LOG.isDebugEnabled()) {
+					LOG.debug(
+							"General exception was thrown, so most likely the traced client wasn't called. Falling back to a manual span");
+				}
+				fallbackSpan = tracingFeignClient().handleSend(
+						new HashMap<>(request.headers()), request, fallbackSpan);
+				tracingFeignClient().handleReceive(fallbackSpan, response, e);
+			}
+			throw e;
+		}
+		finally {
+			fallbackSpan.abandon();
+		}
+	}
+
+	private Tracer tracer() {
+		if (tracer == null) {
+			tracer = beanFactory.getBean(Tracer.class);
+		}
+		return tracer;
+	}
+
+	private HttpTracing httpTracing() {
+		if (httpTracing == null) {
+			httpTracing = beanFactory.getBean(HttpTracing.class);
+		}
+		return httpTracing;
+	}
+
+	private TracingFeignClient tracingFeignClient() {
+		if (tracingFeignClient == null) {
+			tracingFeignClient = (TracingFeignClient) TracingFeignClient
+					.create(httpTracing(), getDelegate());
+		}
+		return tracingFeignClient;
+	}
+
+}

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

@@ -24,7 +24,9 @@ import org.apache.commons.logging.LogFactory;
 
 import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.BeanFactory;
+import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
 import org.springframework.cloud.netflix.ribbon.SpringClientFactory;
+import org.springframework.cloud.openfeign.loadbalancer.FeignBlockingLoadBalancerClient;
 import org.springframework.cloud.openfeign.ribbon.CachingSpringLoadBalancerFactory;
 import org.springframework.cloud.openfeign.ribbon.LoadBalancerFeignClient;
 import org.springframework.util.ClassUtils;
@@ -33,14 +35,21 @@ import org.springframework.util.ClassUtils;
  * Class that wraps Feign related classes into their Trace representative.
  *
  * @author Marcin Grzejszczak
+ * @author Olga Maciaszek-Sharma
  * @since 1.0.1
  */
 final class TraceFeignObjectWrapper {
 
-	private static final Log log = LogFactory.getLog(TraceFeignObjectWrapper.class);
+	public static final String EXCEPTION_WARNING = "Exception occurred while trying to access the delegate's field. Will fallback to default instrumentation mechanism, which means that the delegate might not be instrumented";
 
 	private static final boolean ribbonPresent;
 
+	private static final Log log = LogFactory.getLog(TraceFeignObjectWrapper.class);
+
+	private static final boolean loadBalancerPresent;
+
+	private static final String DELEGATE = "delegate";
+
 	static {
 		ribbonPresent = ClassUtils.isPresent(
 				"org.springframework.cloud.openfeign.ribbon.LoadBalancerFeignClient",
@@ -48,6 +57,12 @@ final class TraceFeignObjectWrapper {
 				&& ClassUtils.isPresent(
 						"org.springframework.cloud.netflix.ribbon.SpringClientFactory",
 						null);
+		loadBalancerPresent = ClassUtils.isPresent(
+				"org.springframework.cloud.openfeign.loadbalancer.FeignBlockingLoadBalancerClient",
+				null)
+				&& ClassUtils.isPresent(
+						"org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient",
+						null);
 	}
 
 	private final BeanFactory beanFactory;
@@ -56,6 +71,8 @@ final class TraceFeignObjectWrapper {
 
 	private Object springClientFactory;
 
+	private Object loadBalancerClient;
+
 	TraceFeignObjectWrapper(BeanFactory beanFactory) {
 		this.beanFactory = beanFactory;
 	}
@@ -64,34 +81,16 @@ final class TraceFeignObjectWrapper {
 		if (bean instanceof Client && !(bean instanceof TracingFeignClient)) {
 			if (ribbonPresent && bean instanceof LoadBalancerFeignClient
 					&& !(bean instanceof TraceLoadBalancerFeignClient)) {
-				if (AopUtils.getTargetClass(bean).equals(LoadBalancerFeignClient.class)) {
-					LoadBalancerFeignClient client = ((LoadBalancerFeignClient) bean);
-					return new TraceLoadBalancerFeignClient(
-							(Client) new TraceFeignObjectWrapper(this.beanFactory)
-									.wrap(client.getDelegate()),
-							factory(), (SpringClientFactory) clientFactory(),
-							this.beanFactory);
-				}
-				else {
-					LoadBalancerFeignClient client = ((LoadBalancerFeignClient) bean);
-					try {
-						Field delegate = LoadBalancerFeignClient.class
-								.getDeclaredField("delegate");
-						delegate.setAccessible(true);
-						delegate.set(client, new TraceFeignObjectWrapper(this.beanFactory)
-								.wrap(client.getDelegate()));
-					}
-					catch (NoSuchFieldException | IllegalArgumentException
-							| IllegalAccessException | SecurityException e) {
-						log.warn(
-								"Exception occurred while trying to access the delegate's field. Will fallback to default instrumentation mechanism, which means that the delegate might not be instrumented",
-								e);
-					}
-					return new TraceLoadBalancerFeignClient(client, factory(),
-							(SpringClientFactory) clientFactory(), this.beanFactory);
-				}
+				return instrumentedFeignRibbonClient(bean);
 			}
-			else if (ribbonPresent && bean instanceof TraceLoadBalancerFeignClient) {
+			if (ribbonPresent && bean instanceof TraceLoadBalancerFeignClient) {
+				return bean;
+			}
+			if (loadBalancerPresent && bean instanceof FeignBlockingLoadBalancerClient
+					&& !(bean instanceof TraceFeignBlockingLoadBalancerClient)) {
+				return instrumentedFeignLoadBalancerClient(bean);
+			}
+			if (ribbonPresent && bean instanceof TraceFeignBlockingLoadBalancerClient) {
 				return bean;
 			}
 			return new LazyTracingFeignClient(this.beanFactory, (Client) bean);
@@ -99,6 +98,57 @@ final class TraceFeignObjectWrapper {
 		return bean;
 	}
 
+	private Object instrumentedFeignLoadBalancerClient(Object bean) {
+		if (AopUtils.getTargetClass(bean).equals(FeignBlockingLoadBalancerClient.class)) {
+			FeignBlockingLoadBalancerClient client = ((FeignBlockingLoadBalancerClient) bean);
+			return new TraceFeignBlockingLoadBalancerClient(
+					(Client) new TraceFeignObjectWrapper(this.beanFactory)
+							.wrap(client.getDelegate()),
+					(BlockingLoadBalancerClient) loadBalancerClient(), this.beanFactory);
+		}
+		else {
+			FeignBlockingLoadBalancerClient client = ((FeignBlockingLoadBalancerClient) bean);
+			try {
+				Field delegate = FeignBlockingLoadBalancerClient.class
+						.getDeclaredField(DELEGATE);
+				delegate.setAccessible(true);
+				delegate.set(client, new TraceFeignObjectWrapper(this.beanFactory)
+						.wrap(client.getDelegate()));
+			}
+			catch (NoSuchFieldException | IllegalArgumentException
+					| IllegalAccessException | SecurityException e) {
+				log.warn(EXCEPTION_WARNING, e);
+			}
+			return new TraceFeignBlockingLoadBalancerClient(client,
+					(BlockingLoadBalancerClient) loadBalancerClient(), this.beanFactory);
+		}
+	}
+
+	private Object instrumentedFeignRibbonClient(Object bean) {
+		if (AopUtils.getTargetClass(bean).equals(LoadBalancerFeignClient.class)) {
+			LoadBalancerFeignClient client = ((LoadBalancerFeignClient) bean);
+			return new TraceLoadBalancerFeignClient(
+					(Client) new TraceFeignObjectWrapper(this.beanFactory)
+							.wrap(client.getDelegate()),
+					factory(), (SpringClientFactory) clientFactory(), this.beanFactory);
+		}
+		else {
+			LoadBalancerFeignClient client = ((LoadBalancerFeignClient) bean);
+			try {
+				Field delegate = LoadBalancerFeignClient.class.getDeclaredField(DELEGATE);
+				delegate.setAccessible(true);
+				delegate.set(client, new TraceFeignObjectWrapper(this.beanFactory)
+						.wrap(client.getDelegate()));
+			}
+			catch (NoSuchFieldException | IllegalArgumentException
+					| IllegalAccessException | SecurityException e) {
+				log.warn(EXCEPTION_WARNING, e);
+			}
+			return new TraceLoadBalancerFeignClient(client, factory(),
+					(SpringClientFactory) clientFactory(), this.beanFactory);
+		}
+	}
+
 	private CachingSpringLoadBalancerFactory factory() {
 		if (this.cachingSpringLoadBalancerFactory == null) {
 			this.cachingSpringLoadBalancerFactory = this.beanFactory
@@ -115,4 +165,11 @@ final class TraceFeignObjectWrapper {
 		return this.springClientFactory;
 	}
 
+	private Object loadBalancerClient() {
+		if (loadBalancerClient == null) {
+			loadBalancerClient = beanFactory.getBean(BlockingLoadBalancerClient.class);
+		}
+		return loadBalancerClient;
+	}
+
 }

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

@@ -16,8 +16,6 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client.feign;
 
-import brave.Tracing;
-import brave.http.HttpTracing;
 import feign.Client;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -26,9 +24,13 @@ import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
 import org.springframework.beans.factory.BeanFactory;
+import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
+import org.springframework.cloud.openfeign.loadbalancer.FeignBlockingLoadBalancerClient;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.BDDAssertions.then;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /**
  * @author Marcin Grzejszczak
@@ -36,10 +38,6 @@ import static org.mockito.Mockito.mock;
 @RunWith(MockitoJUnitRunner.class)
 public class TracingFeignObjectWrapperTests {
 
-	Tracing tracing = Tracing.newBuilder().build();
-
-	HttpTracing httpTracing = HttpTracing.create(this.tracing);
-
 	@Mock
 	BeanFactory beanFactory;
 
@@ -47,16 +45,57 @@ public class TracingFeignObjectWrapperTests {
 	TraceFeignObjectWrapper traceFeignObjectWrapper;
 
 	@Test
-	public void should_wrap_a_client_into_lazy_trace_client() throws Exception {
+	public void should_wrap_a_client_into_lazy_trace_client() {
 		then(this.traceFeignObjectWrapper.wrap(mock(Client.class)))
 				.isExactlyInstanceOf(LazyTracingFeignClient.class);
 	}
 
 	@Test
-	public void should_not_wrap_a_bean_that_is_not_feign_related() throws Exception {
+	public void should_not_wrap_a_bean_that_is_not_feign_related() {
 		String notFeignRelatedObject = "object";
 		then(this.traceFeignObjectWrapper.wrap(notFeignRelatedObject))
 				.isSameAs(notFeignRelatedObject);
 	}
 
+	// gh-1528
+	@Test
+	public void should_wrap_feign_loadbalancer_client() {
+		Client delegate = mock(Client.class);
+		BlockingLoadBalancerClient loadBalancerClient = mock(
+				BlockingLoadBalancerClient.class);
+		when(beanFactory.getBean(BlockingLoadBalancerClient.class))
+				.thenReturn(loadBalancerClient);
+
+		Object wrapped = traceFeignObjectWrapper
+				.wrap(new FeignBlockingLoadBalancerClient(delegate, loadBalancerClient));
+
+		assertThat(wrapped).isInstanceOf(TraceFeignBlockingLoadBalancerClient.class);
+	}
+
+	// gh-1528, gh-1125
+	@Test
+	public void should_wrap_subclass_of_feign_loadbalancer_client() {
+		Client delegate = mock(Client.class);
+		BlockingLoadBalancerClient loadBalancerClient = mock(
+				BlockingLoadBalancerClient.class);
+		when(beanFactory.getBean(BlockingLoadBalancerClient.class))
+				.thenReturn(loadBalancerClient);
+
+		Object wrapped = traceFeignObjectWrapper.wrap(
+				new TestFeignBlockingLoadBalancerClient(delegate, loadBalancerClient));
+
+		assertThat(wrapped).isInstanceOf(TraceFeignBlockingLoadBalancerClient.class);
+
+	}
+
+	static class TestFeignBlockingLoadBalancerClient
+			extends FeignBlockingLoadBalancerClient {
+
+		TestFeignBlockingLoadBalancerClient(Client delegate,
+				BlockingLoadBalancerClient loadBalancerClient) {
+			super(delegate, loadBalancerClient);
+		}
+
+	}
+
 }

tests/spring-cloud-sleuth-instrumentation-feign-tests/src/test/java/org/springframework/cloud/sleuth/instrument/feign/issues/issue1125delegates/ManuallyCreatedDelegateLoadBalancerFeignClientTests.java

@@ -61,7 +61,7 @@ import static org.assertj.core.api.BDDAssertions.then;
 @RunWith(SpringRunner.class)
 @SpringBootTest(classes = Application.class,
 		webEnvironment = SpringBootTest.WebEnvironment.NONE,
-		properties = {"feign.hystrix.enabled=false"})
+		properties = { "feign.hystrix.enabled=false" })
 @DirtiesContext(classMode = ClassMode.AFTER_EACH_TEST_METHOD)
 public class ManuallyCreatedDelegateLoadBalancerFeignClientTests {
 
@@ -138,8 +138,8 @@ class Application {
 	public MyNameRemote myNameRemote(Client client, Decoder decoder, Encoder encoder,
 			Contract contract) {
 		return Feign.builder().client(client).encoder(encoder).decoder(decoder)
-				.contract(contract).target(new HardCodedTarget<>(
-						MyNameRemote.class, "foo", "https://non.existing.url"));
+				.contract(contract).target(new HardCodedTarget<>(MyNameRemote.class,
+						"foo", "https://non.existing.url"));
 	}
 
 	@Bean