Fixed webclient sender not working properly; fixes gh-2126

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/zipkin2/ZipkinAutoConfiguration.java

@@ -16,9 +16,9 @@
 
 package org.springframework.cloud.sleuth.autoconfig.zipkin2;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
@@ -108,25 +108,29 @@ public class ZipkinAutoConfiguration {
 	}
 
 	/** Limits {@link Sender#check()} to {@code deadlineMillis}. */
-	static CheckResult checkResult(ExecutorService zipkinExecutor, Sender sender, long deadlineMillis) {
-		Future<CheckResult> future = zipkinExecutor.submit(sender::check);
-		try {
-			return future.get(deadlineMillis, TimeUnit.MILLISECONDS);
-		}
-		catch (TimeoutException e) {
-			return CheckResult.failed(new TimeoutException("Timed out after " + deadlineMillis + "ms"));
-		}
-		catch (Exception e) {
-			return CheckResult.failed(e);
-		}
+	static CompletableFuture<CheckResult> checkResult(ExecutorService zipkinExecutor, Sender sender,
+			long deadlineMillis) {
+		return CompletableFuture.supplyAsync(sender::check, zipkinExecutor).whenComplete((checkResult, throwable) -> {
+			Throwable exception = throwable instanceof TimeoutException
+					? new TimeoutException("Timed out after " + deadlineMillis + "ms") : throwable;
+			CheckResult result;
+			if (checkResult != null && checkResult.error() != null
+					&& checkResult.error().getCause() instanceof TimeoutException) {
+				exception = new TimeoutException("Timed out after " + deadlineMillis + "ms");
+				result = CheckResult.failed(exception);
+			}
+			else {
+				result = checkResult == null ? CheckResult.failed(exception) : checkResult;
+			}
+			logCheckResult(sender, result);
+		});
 	}
 
 	@Bean(REPORTER_BEAN_NAME)
 	@ConditionalOnMissingBean(name = REPORTER_BEAN_NAME)
 	Reporter<Span> reporter(ReporterMetrics reporterMetrics, ZipkinProperties zipkin,
 			@Qualifier(SENDER_BEAN_NAME) Sender sender) {
-		CheckResult checkResult = checkResult(zipkinExecutor, sender, 1_000L);
-		logCheckResult(sender, checkResult);
+		checkResult(zipkinExecutor, sender, zipkin.getCheckTimeout());
 
 		// Note: AsyncReporter supports memory bounds
 		AsyncReporter<Span> asyncReporter = AsyncReporter.builder(sender).queuedMaxSpans(zipkin.getQueuedMaxSpans())
@@ -152,7 +156,7 @@ public class ZipkinAutoConfiguration {
 		return asyncReporter;
 	}
 
-	private void logCheckResult(Sender sender, CheckResult checkResult) {
+	private static void logCheckResult(Sender sender, CheckResult checkResult) {
 		if (log.isDebugEnabled() && checkResult != null && checkResult.ok()) {
 			log.debug("Check result of the [" + sender.toString() + "] is [" + checkResult + "]");
 		}

spring-cloud-sleuth-autoconfigure/src/main/java/org/springframework/cloud/sleuth/autoconfig/zipkin2/ZipkinHttpSenderConfiguration.java

@@ -85,7 +85,7 @@ class ZipkinHttpSenderConfiguration {
 		Sender webClientSender(ZipkinProperties zipkin, ZipkinWebClientBuilderProvider zipkinWebClientBuilderProvider) {
 			WebClient.Builder webClientBuilder = zipkinWebClientBuilderProvider.zipkinWebClientBuilder();
 			return new WebClientSender(webClientBuilder.build(), zipkin.getBaseUrl(), zipkin.getApiPath(),
-					zipkin.getEncoder());
+					zipkin.getEncoder(), zipkin.getCheckTimeout());
 		}
 
 		@Bean

spring-cloud-sleuth-autoconfigure/src/test/java/org/springframework/cloud/sleuth/autoconfig/zipkin2/ZipkinAutoConfigurationTests.java

@@ -0,0 +1,149 @@
+/*
+ * Copyright 2013-2021 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.autoconfig.zipkin2;
+
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Supplier;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import zipkin2.Call;
+import zipkin2.CheckResult;
+import zipkin2.codec.Encoding;
+import zipkin2.reporter.Sender;
+
+import org.springframework.boot.test.system.CapturedOutput;
+import org.springframework.boot.test.system.OutputCaptureExtension;
+
+import static org.assertj.core.api.BDDAssertions.then;
+import static org.awaitility.Awaitility.await;
+
+@ExtendWith(OutputCaptureExtension.class)
+class ZipkinAutoConfigurationTests {
+
+	ExecutorService service = Executors.newSingleThreadExecutor();
+
+	@AfterEach
+	void clean() {
+		this.service.shutdown();
+	}
+
+	@Test
+	void shouldReturnShortTimeoutExceptionWhenRootCauseTimeout(CapturedOutput capture) {
+		ZipkinAutoConfiguration.checkResult(service,
+				new CustomSender(() -> CheckResult.failed(new RuntimeException(new TimeoutException("boom")))), 1L);
+
+		await().untilAsserted(() -> then(capture.toString()).doesNotContain("CheckResult{ok=true, error=null}")
+				.contains("CheckResult{ok=false, error=java.util.concurrent.TimeoutException: Timed out after 1ms}"));
+	}
+
+	@Test
+	void shouldReturnOriginalExceptionWhenRootCauseNotTimeout(CapturedOutput capture) {
+		ZipkinAutoConfiguration.checkResult(service,
+				new CustomSender(() -> CheckResult.failed(new RuntimeException(new RuntimeException("boom")))), 1L);
+
+		await().atMost(1, TimeUnit.SECONDS).untilAsserted(
+				() -> then(capture.toString()).doesNotContain("CheckResult{ok=true, error=null}").contains(
+						"CheckResult{ok=false, error=java.lang.RuntimeException: java.lang.RuntimeException: boom}"));
+	}
+
+	@Test
+	void shouldReturnCheckResultWhenNoExceptionPresent(CapturedOutput capture) {
+		ZipkinAutoConfiguration.checkResult(service, new CustomSender(() -> CheckResult.OK), 1L);
+
+		await().atMost(1, TimeUnit.SECONDS)
+				.untilAsserted(() -> then(capture.toString()).contains("CheckResult{ok=true, error=null}"));
+	}
+
+	@Test
+	void shouldReturnExceptionWhenNoCheckResultPresent(CapturedOutput capture) {
+		ZipkinAutoConfiguration.checkResult(service, new ExceptionThrowingSender(), 1L);
+
+		await().atMost(1, TimeUnit.SECONDS).untilAsserted(() -> then(capture.toString()).contains(
+				"CheckResult{ok=false, error=java.util.concurrent.CompletionException: java.lang.RuntimeException: boom}"));
+	}
+
+	private static final class ExceptionThrowingSender extends Sender {
+
+		@Override
+		public Encoding encoding() {
+			return null;
+		}
+
+		@Override
+		public int messageMaxBytes() {
+			return 0;
+		}
+
+		@Override
+		public int messageSizeInBytes(List<byte[]> list) {
+			return 0;
+		}
+
+		@Override
+		public Call<Void> sendSpans(List<byte[]> list) {
+			return null;
+		}
+
+		@Override
+		public CheckResult check() {
+			throw new RuntimeException("boom");
+		}
+
+	}
+
+	private static final class CustomSender extends Sender {
+
+		private final Supplier<CheckResult> doSth;
+
+		private CustomSender(Supplier<CheckResult> doSth) {
+			this.doSth = doSth;
+		}
+
+		@Override
+		public Encoding encoding() {
+			return null;
+		}
+
+		@Override
+		public int messageMaxBytes() {
+			return 0;
+		}
+
+		@Override
+		public int messageSizeInBytes(List<byte[]> list) {
+			return 0;
+		}
+
+		@Override
+		public Call<Void> sendSpans(List<byte[]> list) {
+			return null;
+		}
+
+		@Override
+		public CheckResult check() {
+			return this.doSth.get();
+		}
+
+	}
+
+}

spring-cloud-sleuth-zipkin/src/main/java/org/springframework/cloud/sleuth/zipkin2/WebClientSender.java

@@ -17,6 +17,7 @@
 package org.springframework.cloud.sleuth.zipkin2;
 
 import java.net.URI;
+import java.time.Duration;
 
 import zipkin2.Span;
 import zipkin2.codec.BytesEncoder;
@@ -32,13 +33,40 @@ import org.springframework.web.reactive.function.client.WebClient;
  */
 public class WebClientSender extends HttpSender {
 
+	private static final long DEFAULT_CHECK_TIMEOUT = 1_000L;
+
+	/**
+	 * Use
+	 * {@link WebClientSender#WebClientSender(WebClient, String, String, BytesEncoder, long)}.
+	 * @param webClient web client
+	 * @param baseUrl base url
+	 * @param apiPath api path
+	 * @param encoder encoder
+	 * @deprecated use
+	 * {@link WebClientSender#WebClientSender(WebClient, String, String, BytesEncoder, long)}
+	 */
+	@Deprecated
 	public WebClientSender(WebClient webClient, String baseUrl, String apiPath, BytesEncoder<Span> encoder) {
-		super((url, mediaType, bytes) -> post(url, mediaType, bytes, webClient), baseUrl, apiPath, encoder);
+		this(webClient, baseUrl, apiPath, encoder, DEFAULT_CHECK_TIMEOUT);
 	}
 
-	private static void post(String url, MediaType mediaType, byte[] json, WebClient webClient) {
+	/**
+	 * Creates a new instance of {@link WebClientSender}.
+	 * @param webClient web client
+	 * @param baseUrl base url
+	 * @param apiPath api path
+	 * @param encoder encoder
+	 * @param checkTimeout check timeout
+	 */
+	public WebClientSender(WebClient webClient, String baseUrl, String apiPath, BytesEncoder<Span> encoder,
+			long checkTimeout) {
+		super((url, mediaType, bytes) -> post(url, mediaType, bytes, webClient, checkTimeout), baseUrl, apiPath,
+				encoder);
+	}
+
+	private static void post(String url, MediaType mediaType, byte[] json, WebClient webClient, long checkTimeout) {
 		webClient.post().uri(URI.create(url)).accept(mediaType).contentType(mediaType).bodyValue(json).retrieve()
-				.toBodilessEntity().subscribe();
+				.toBodilessEntity().timeout(Duration.ofMillis(checkTimeout)).block();
 	}
 
 	@Override

spring-cloud-sleuth-zipkin/src/main/java/org/springframework/cloud/sleuth/zipkin2/ZipkinProperties.java

@@ -55,6 +55,11 @@ public class ZipkinProperties {
 	 */
 	private boolean enabled = true;
 
+	/**
+	 * Timeout in millis for the check for Zipkin availability.
+	 */
+	private int checkTimeout = 1_000;
+
 	/**
 	 * Timeout in seconds before pending spans will be sent in batches to Zipkin.
 	 */
@@ -128,6 +133,14 @@ public class ZipkinProperties {
 		this.messageTimeout = messageTimeout;
 	}
 
+	public int getCheckTimeout() {
+		return this.checkTimeout;
+	}
+
+	public void setCheckTimeout(int checkTimeout) {
+		this.checkTimeout = checkTimeout;
+	}
+
 	public Compression getCompression() {
 		return this.compression;
 	}

spring-cloud-sleuth-zipkin/src/main/java/org/springframework/cloud/sleuth/zipkin2/ZipkinRestTemplateWrapper.java

@@ -42,8 +42,6 @@ public class ZipkinRestTemplateWrapper extends RestTemplate {
 
 	private static final Log log = LogFactory.getLog(ZipkinRestTemplateWrapper.class);
 
-	private static final int DEFAULT_TIMEOUT = 500;
-
 	private final ZipkinProperties zipkinProperties;
 
 	private final ZipkinUrlExtractor extractor;
@@ -51,13 +49,13 @@ public class ZipkinRestTemplateWrapper extends RestTemplate {
 	public ZipkinRestTemplateWrapper(ZipkinProperties zipkinProperties, ZipkinUrlExtractor extractor) {
 		this.zipkinProperties = zipkinProperties;
 		this.extractor = extractor;
-		setRequestFactory(clientHttpRequestFactory());
+		setRequestFactory(clientHttpRequestFactory(zipkinProperties));
 	}
 
-	private ClientHttpRequestFactory clientHttpRequestFactory() {
+	private ClientHttpRequestFactory clientHttpRequestFactory(ZipkinProperties zipkinProperties) {
 		SimpleClientHttpRequestFactory factory = new SimpleClientHttpRequestFactory();
-		factory.setReadTimeout(DEFAULT_TIMEOUT);
-		factory.setConnectTimeout(DEFAULT_TIMEOUT);
+		factory.setReadTimeout(zipkinProperties.getCheckTimeout());
+		factory.setConnectTimeout(zipkinProperties.getCheckTimeout());
 		return factory;
 	}
 

spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/AbstractSenderTest.java

@@ -17,15 +17,20 @@
 package org.springframework.cloud.sleuth.zipkin2;
 
 import java.io.IOException;
+import java.net.InetAddress;
+import java.net.ServerSocket;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 
+import javax.net.ServerSocketFactory;
+
 import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
 import okhttp3.mockwebserver.RecordedRequest;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import zipkin2.Call;
+import zipkin2.CheckResult;
 import zipkin2.Endpoint;
 import zipkin2.Span;
 import zipkin2.codec.Encoding;
@@ -123,6 +128,62 @@ abstract class AbstractSenderTest {
 		assertThat(this.sender).hasToString(expectedToString());
 	}
 
+	@Test
+	public void testWhereServerDown() throws IOException {
+		this.server.shutdown();
+
+		final Sender sender = jsonSender();
+		CheckResult checkResult = sender.check();
+		assertThat(checkResult.ok()).isFalse();
+		assertThat(checkResult.error()).hasMessageContaining("Connection refused");
+	}
+
+	@Test
+	public void testWhereServerSlow() throws IOException {
+		MockWebServer server = new MockWebServer();
+		ServerSocketFactory socketFactory = ServerSocketFactory.getDefault();
+		ServerSocketFactory factory = new ServerSocketFactory() {
+			@Override
+			public ServerSocket createServerSocket(int port) throws IOException {
+				slow();
+				return socketFactory.createServerSocket(port);
+			}
+
+			private void slow() {
+				try {
+					Thread.sleep(2_000);
+				}
+				catch (InterruptedException e) {
+					throw new RuntimeException(e);
+				}
+			}
+
+			@Override
+			public ServerSocket createServerSocket(int port, int backlog) throws IOException {
+				slow();
+				return socketFactory.createServerSocket(port, backlog);
+			}
+
+			@Override
+			public ServerSocket createServerSocket(int port, int backlog, InetAddress ifAddress) throws IOException {
+				slow();
+				return socketFactory.createServerSocket(port, backlog, ifAddress);
+			}
+		};
+		server.setServerSocketFactory(factory);
+		MockResponse mockResponse = new MockResponse();
+		mockResponse.setBodyDelay(1, TimeUnit.SECONDS);
+		server.enqueue(mockResponse);
+		final Sender sender = jsonSender();
+
+		CheckResult checkResult = sender.check();
+
+		assertThat(checkResult.ok()).isFalse();
+		assertThat(checkResult.error()).hasMessageContaining("TimeoutException");
+
+		server.shutdown();
+	}
+
 	Call<Void> send(Span... spans) {
 		SpanBytesEncoder bytesEncoder = this.sender.encoding() == Encoding.JSON ? SpanBytesEncoder.JSON_V2
 				: SpanBytesEncoder.PROTO3;

spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/RestTemplateSenderTest.java

@@ -18,6 +18,8 @@ package org.springframework.cloud.sleuth.zipkin2;
 
 import zipkin2.reporter.Sender;
 
+import org.springframework.http.client.ClientHttpRequestFactory;
+import org.springframework.http.client.SimpleClientHttpRequestFactory;
 import org.springframework.web.client.RestTemplate;
 
 import static zipkin2.codec.SpanBytesEncoder.JSON_V2;
@@ -25,19 +27,22 @@ import static zipkin2.codec.SpanBytesEncoder.PROTO3;
 
 class RestTemplateSenderTest extends AbstractSenderTest {
 
+	public static final int DEFAULT_CHECK_TIMEOUT = 400;
+
 	@Override
 	Sender jsonSender() {
-		return new RestTemplateSender(new RestTemplate(), this.endpoint, null, JSON_V2);
+		return new RestTemplateSender(new RestTemplate(clientHttpRequestFactory()), this.endpoint, null, JSON_V2);
 	}
 
 	@Override
 	Sender jsonSender(String mockedApiPath) {
-		return new RestTemplateSender(new RestTemplate(), this.endpoint, mockedApiPath, JSON_V2);
+		return new RestTemplateSender(new RestTemplate(clientHttpRequestFactory()), this.endpoint, mockedApiPath,
+				JSON_V2);
 	}
 
 	@Override
 	Sender protoSender() {
-		return new RestTemplateSender(new RestTemplate(), this.endpoint, "", PROTO3);
+		return new RestTemplateSender(new RestTemplate(clientHttpRequestFactory()), this.endpoint, "", PROTO3);
 	}
 
 	@Override
@@ -53,4 +58,11 @@ class RestTemplateSenderTest extends AbstractSenderTest {
 		return "RestTemplateSender{" + this.endpoint + mockedApiPath + "}";
 	}
 
+	private ClientHttpRequestFactory clientHttpRequestFactory() {
+		SimpleClientHttpRequestFactory factory = new SimpleClientHttpRequestFactory();
+		factory.setReadTimeout(DEFAULT_CHECK_TIMEOUT);
+		factory.setConnectTimeout(DEFAULT_CHECK_TIMEOUT);
+		return factory;
+	}
+
 }

spring-cloud-sleuth-zipkin/src/test/java/org/springframework/cloud/sleuth/zipkin2/WebClientSenderTests.java

@@ -26,22 +26,24 @@ import static zipkin2.codec.SpanBytesEncoder.PROTO3;
 
 class WebClientSenderTests extends AbstractSenderTest {
 
+	public static final int DEFAULT_CHECK_TIMEOUT = 400;
+
 	@Override
 	Sender jsonSender() {
 		return new WebClientSender(WebClient.builder().clientConnector(new ReactorClientHttpConnector()).build(),
-				this.endpoint, null, JSON_V2);
+				this.endpoint, null, JSON_V2, DEFAULT_CHECK_TIMEOUT);
 	}
 
 	@Override
 	Sender jsonSender(String mockedApiPath) {
 		return new WebClientSender(WebClient.builder().clientConnector(new ReactorClientHttpConnector()).build(),
-				this.endpoint, mockedApiPath, JSON_V2);
+				this.endpoint, mockedApiPath, JSON_V2, DEFAULT_CHECK_TIMEOUT);
 	}
 
 	@Override
 	Sender protoSender() {
 		return new WebClientSender(WebClient.builder().clientConnector(new ReactorClientHttpConnector()).build(),
-				this.endpoint, "", PROTO3);
+				this.endpoint, "", PROTO3, DEFAULT_CHECK_TIMEOUT);
 	}
 
 	@Override

tests/common/src/main/java/org/springframework/cloud/sleuth/autoconfig/zipkin2/ZipkinAutoConfigurationTests.java

@@ -18,10 +18,10 @@ package org.springframework.cloud.sleuth.autoconfig.zipkin2;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 
 import io.micrometer.core.instrument.MeterRegistry;
 import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
@@ -227,71 +227,20 @@ public abstract class ZipkinAutoConfigurationTests {
 	}
 
 	@Test
-	public void checkResult_onTime() {
+	public void checkResult_onTime() throws ExecutionException, InterruptedException {
 		Sender sender = mock(Sender.class);
 		when(sender.check()).thenReturn(CheckResult.OK);
 
-		assertThat(ZipkinAutoConfiguration.checkResult(zipkinExecutor, sender, 200).ok()).isTrue();
+		assertThat(ZipkinAutoConfiguration.checkResult(zipkinExecutor, sender, 200).get().ok()).isTrue();
 	}
 
 	@Test
-	public void checkResult_onTime_notOk() {
+	public void checkResult_onTime_notOk() throws ExecutionException, InterruptedException {
 		Sender sender = mock(Sender.class);
 		RuntimeException exception = new RuntimeException("dead");
 		when(sender.check()).thenReturn(CheckResult.failed(exception));
 
-		assertThat(ZipkinAutoConfiguration.checkResult(zipkinExecutor, sender, 200).error()).isSameAs(exception);
-	}
-
-	/** Bug in {@link Sender} as it shouldn't throw */
-	@Test
-	public void checkResult_thrown() {
-		Sender sender = mock(Sender.class);
-		RuntimeException exception = new RuntimeException("dead");
-		when(sender.check()).thenThrow(exception);
-
-		assertThat(ZipkinAutoConfiguration.checkResult(zipkinExecutor, sender, 200).error()).hasCause(exception);
-	}
-
-	@Test
-	public void checkResult_slow() {
-		assertThat(ZipkinAutoConfiguration.checkResult(zipkinExecutor, new Sender() {
-			@Override
-			public CheckResult check() {
-				try {
-					Thread.sleep(500L);
-				}
-				catch (InterruptedException e) {
-					throw new AssertionError(e);
-				}
-				return CheckResult.OK;
-			}
-
-			@Override
-			public Encoding encoding() {
-				return Encoding.JSON;
-			}
-
-			@Override
-			public int messageMaxBytes() {
-				return 0;
-			}
-
-			@Override
-			public int messageSizeInBytes(List<byte[]> list) {
-				return 0;
-			}
-
-			@Override
-			public Call<Void> sendSpans(List<byte[]> list) {
-				return Call.create(null);
-			}
-
-			@Override
-			public String toString() {
-				return "FakeSender{}";
-			}
-		}, 200).error()).isInstanceOf(TimeoutException.class).hasMessage("Timed out after 200ms");
+		assertThat(ZipkinAutoConfiguration.checkResult(zipkinExecutor, sender, 200).get().error()).isSameAs(exception);
 	}
 
 	@Configuration(proxyBeanMethods = false)