From 8fbe7648a1dc3ffdd9340db2d79a27e70543dba6 Mon Sep 17 00:00:00 2001
From: Marcin Grzejszczak <marcin.grzejszczak@gmail.com>
Date: Tue, 23 Jun 2020 09:32:52 +0200
Subject: [PATCH] Updated changes following Reactor changes

---
 .../client/HttpClientBeanPostProcessor.java   | 19 +++++++++++--------
 .../HttpClientBeanPostProcessorTest.java      |  4 ++--
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
index 83bbcf3dc..9c4221c6d 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessor.java
@@ -17,17 +17,17 @@
 package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
-import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import brave.Span;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
 import brave.propagation.TraceContext;
-import io.netty.bootstrap.Bootstrap;
 import reactor.core.publisher.Mono;
 import reactor.netty.Connection;
 import reactor.netty.http.client.HttpClient;
@@ -82,8 +82,8 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 
 	}
 
-	static class TracingMapConnect implements
-			BiFunction<Mono<? extends Connection>, Bootstrap, Mono<? extends Connection>> {
+	static class TracingMapConnect
+			implements Function<Mono<? extends Connection>, Mono<? extends Connection>> {
 
 		static final Exception CANCELLED_ERROR = new CancellationException("CANCELLED") {
 			@Override
@@ -99,8 +99,7 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 		}
 
 		@Override
-		public Mono<? extends Connection> apply(Mono<? extends Connection> mono,
-				Bootstrap bootstrap) {
+		public Mono<? extends Connection> apply(Mono<? extends Connection> mono) {
 			// This function is invoked once per-request. We keep a reference to the
 			// pending client span here, so that only one signal completes the span.
 			PendingSpan pendingSpan = new PendingSpan();
@@ -173,8 +172,12 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			if (span.isNoop()) {
 				return;
 			}
-			InetSocketAddress socketAddress = connection.address();
-			span.remoteIpAndPort(socketAddress.getHostString(), socketAddress.getPort());
+			SocketAddress socketAddress = connection.address();
+			if (socketAddress instanceof InetSocketAddress) {
+				InetSocketAddress inetSocketAddress = (InetSocketAddress) socketAddress;
+				span.remoteIpAndPort(inetSocketAddress.getHostString(),
+						inetSocketAddress.getPort());
+			}
 		}
 
 	}
diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessorTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessorTest.java
index d3645be79..944795599 100644
--- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessorTest.java
+++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/HttpClientBeanPostProcessorTest.java
@@ -71,7 +71,7 @@ public class HttpClientBeanPostProcessorTest {
 				});
 
 		// Wrap and run the assertions
-		tracingMapConnect.apply(original, bootstrap).log().subscribe();
+		tracingMapConnect.apply(original).log().subscribe();
 	}
 
 	@Test
@@ -90,7 +90,7 @@ public class HttpClientBeanPostProcessorTest {
 				});
 
 		// Wrap and run the assertions
-		tracingMapConnect.apply(original, bootstrap).log().subscribe();
+		tracingMapConnect.apply(original).log().subscribe();
 	}
 
 }