Add instrumentation for RetryableFeignBlockingLoadBalancerClient. (#1759)

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

@@ -24,9 +24,11 @@ import org.apache.commons.logging.LogFactory;
 
 import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.BeanFactory;
+import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryFactory;
 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.loadbalancer.RetryableFeignBlockingLoadBalancerClient;
 import org.springframework.cloud.openfeign.ribbon.CachingSpringLoadBalancerFactory;
 import org.springframework.cloud.openfeign.ribbon.LoadBalancerFeignClient;
 import org.springframework.cloud.util.ProxyUtils;
@@ -74,6 +76,8 @@ final class TraceFeignObjectWrapper {
 
 	private Object loadBalancerClient;
 
+	private Object loadBalancerRetryFactory;
+
 	TraceFeignObjectWrapper(BeanFactory beanFactory) {
 		this.beanFactory = beanFactory;
 	}
@@ -91,6 +95,11 @@ final class TraceFeignObjectWrapper {
 					&& !(bean instanceof TraceFeignBlockingLoadBalancerClient)) {
 				return instrumentedFeignLoadBalancerClient(bean);
 			}
+			if (loadBalancerPresent
+					&& bean instanceof RetryableFeignBlockingLoadBalancerClient
+					&& !(bean instanceof TraceRetryableFeignBlockingLoadBalancerClient)) {
+				return instrumentedRetryableFeignLoadBalancerClient(bean);
+			}
 			if (ribbonPresent && bean instanceof TraceFeignBlockingLoadBalancerClient) {
 				return bean;
 			}
@@ -125,6 +134,37 @@ final class TraceFeignObjectWrapper {
 		}
 	}
 
+	private Object instrumentedRetryableFeignLoadBalancerClient(Object bean) {
+		if (AopUtils.getTargetClass(bean)
+				.equals(RetryableFeignBlockingLoadBalancerClient.class)) {
+			RetryableFeignBlockingLoadBalancerClient client = ProxyUtils
+					.getTargetObject(bean);
+			return new TraceRetryableFeignBlockingLoadBalancerClient(
+					(Client) new TraceFeignObjectWrapper(beanFactory)
+							.wrap(client.getDelegate()),
+					(BlockingLoadBalancerClient) loadBalancerClient(),
+					(LoadBalancedRetryFactory) loadBalancerRetryFactory(), beanFactory);
+		}
+		else {
+			RetryableFeignBlockingLoadBalancerClient client = ((RetryableFeignBlockingLoadBalancerClient) bean);
+			try {
+				Field delegate = RetryableFeignBlockingLoadBalancerClient.class
+						.getDeclaredField(DELEGATE);
+				delegate.setAccessible(true);
+				delegate.set(client, new TraceFeignObjectWrapper(beanFactory)
+						.wrap(client.getDelegate()));
+			}
+			catch (NoSuchFieldException | IllegalArgumentException
+					| IllegalAccessException | SecurityException e) {
+				log.warn(EXCEPTION_WARNING, e);
+			}
+			return new TraceRetryableFeignBlockingLoadBalancerClient(client,
+					(BlockingLoadBalancerClient) loadBalancerClient(),
+					(LoadBalancedRetryFactory) loadBalancerRetryFactory(),
+					this.beanFactory);
+		}
+	}
+
 	private Object instrumentedFeignRibbonClient(Object bean) {
 		if (AopUtils.getTargetClass(bean).equals(LoadBalancerFeignClient.class)) {
 			LoadBalancerFeignClient client = ((LoadBalancerFeignClient) bean);
@@ -173,4 +213,12 @@ final class TraceFeignObjectWrapper {
 		return loadBalancerClient;
 	}
 
+	private Object loadBalancerRetryFactory() {
+		if (loadBalancerRetryFactory == null) {
+			loadBalancerRetryFactory = beanFactory
+					.getBean(LoadBalancedRetryFactory.class);
+		}
+		return loadBalancerRetryFactory;
+	}
+
 }

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

@@ -0,0 +1,123 @@
+/*
+ * Copyright 2013-2020 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 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.client.loadbalancer.LoadBalancedRetryFactory;
+import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
+import org.springframework.cloud.openfeign.loadbalancer.RetryableFeignBlockingLoadBalancerClient;
+
+/**
+ * A trace representation of {@link RetryableFeignBlockingLoadBalancerClient}. Needed due
+ * to casts in {@link org.springframework.cloud.openfeign.FeignClientFactoryBean}.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 2.2.0
+ * @see RetryableFeignBlockingLoadBalancerClient
+ */
+class TraceRetryableFeignBlockingLoadBalancerClient
+		extends RetryableFeignBlockingLoadBalancerClient {
+
+	private static final Log LOG = LogFactory
+			.getLog(TraceRetryableFeignBlockingLoadBalancerClient.class);
+
+	private final BeanFactory beanFactory;
+
+	Tracer tracer;
+
+	HttpTracing httpTracing;
+
+	TracingFeignClient tracingFeignClient;
+
+	TraceRetryableFeignBlockingLoadBalancerClient(Client delegate,
+			BlockingLoadBalancerClient loadBalancerClient,
+			LoadBalancedRetryFactory retryFactory, BeanFactory beanFactory) {
+		super(delegate, loadBalancerClient, retryFactory);
+		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");
+				}
+				tracingFeignClient().handleSendAndReceive(fallbackSpan, request, 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;
+	}
+
+}