Adds base integration test for reactor clients and adds cancelation tests (#1578)

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

@@ -17,14 +17,15 @@
 package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.net.InetSocketAddress;
+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.Supplier;
 
 import brave.Span;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
-import brave.propagation.CurrentTraceContext;
 import brave.propagation.TraceContext;
 import io.netty.bootstrap.Bootstrap;
 import reactor.core.publisher.Mono;
@@ -63,48 +64,60 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 					.doOnResponse(new TracingDoOnResponse(httpTracing))
 					.doOnRequestError(new TracingDoOnErrorRequest(httpTracing))
 					.doOnRequest(new TracingDoOnRequest(httpTracing))
-					.mapConnect(new TracingMapConnect(httpTracing));
+					.mapConnect(new TracingMapConnect(() -> {
+						HttpTracing ref = httpTracing.get();
+						return ref != null ? ref.tracing().currentTraceContext().get()
+								: null;
+					}));
 		}
 		return bean;
 	}
 
-	/** current client span, cleared on completion. */
-	private static final class CurrentClientSpan extends AtomicReference<Span> {
+	/** The current client span, cleared on completion for any reason. */
+	static final class PendingSpan extends AtomicReference<Span> {
 
 	}
 
-	private static class TracingMapConnect implements
+	static class TracingMapConnect implements
 			BiFunction<Mono<? extends Connection>, Bootstrap, Mono<? extends Connection>> {
 
-		final LazyBean<HttpTracing> httpTracing;
+		static final Exception CANCELLED_ERROR = new CancellationException("CANCELLED") {
+			@Override
+			public Throwable fillInStackTrace() {
+				return this; // stack trace doesn't add value here
+			}
+		};
 
-		CurrentTraceContext currentTraceContext;
+		final Supplier<TraceContext> currentTraceContext;
 
-		TracingMapConnect(LazyBean<HttpTracing> httpTracing) {
-			this.httpTracing = httpTracing;
+		TracingMapConnect(Supplier<TraceContext> currentTraceContext) {
+			this.currentTraceContext = currentTraceContext;
 		}
 
 		@Override
 		public Mono<? extends Connection> apply(Mono<? extends Connection> mono,
 				Bootstrap bootstrap) {
+			// 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();
 			return mono.subscriberContext(context -> {
-				TraceContext invocationContext = currentTraceContext().get();
+				TraceContext invocationContext = currentTraceContext.get();
 				if (invocationContext != null) {
 					// Read in this processor and also in ScopePassingSpanSubscriber
 					context = context.put(TraceContext.class, invocationContext);
 				}
-				return context.put(CurrentClientSpan.class, new CurrentClientSpan());
+				return context.put(PendingSpan.class, pendingSpan);
+			}).doOnCancel(() -> {
+				// Check to see if Subscription.cancel() happened before another signal,
+				// like onComplete() completed the span (clearing the reference).
+				Span span = pendingSpan.getAndSet(null);
+				if (span != null) {
+					span.error(CANCELLED_ERROR);
+					span.finish();
+				}
 			});
 		}
 
-		CurrentTraceContext currentTraceContext() {
-			if (this.currentTraceContext == null) {
-				this.currentTraceContext = this.httpTracing.get().tracing()
-						.currentTraceContext();
-			}
-			return this.currentTraceContext;
-		}
-
 	}
 
 	private static class TracingDoOnRequest
@@ -125,28 +138,24 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			return this.handler;
 		}
 
-		CurrentTraceContext currentTraceContext() {
-			return httpTracing.get().tracing().currentTraceContext();
-		}
-
 		@Override
 		public void accept(HttpClientRequest req, Connection connection) {
-			CurrentClientSpan ref = req.currentContext()
-					.getOrDefault(CurrentClientSpan.class, null);
-			if (ref == null) { // Somehow TracingMapConnect was not invoked.. skip out
-				return;
+			PendingSpan pendingSpan = req.currentContext().getOrDefault(PendingSpan.class,
+					null);
+			if (pendingSpan == null) {
+				return; // Somehow TracingMapConnect was not invoked.. skip out
 			}
 
 			// This might be re-entrant on auto-redirect or connection retry:
 			// See reactor/reactor-netty#1000 for follow-ups.
-			Span clientSpan = ref.getAndSet(null);
-			if (clientSpan != null) {
+			Span span = pendingSpan.getAndSet(null);
+			if (span != null) {
 				// Retry from a connect fail wouldn't have parsed the request, leading to
 				// an empty span with no data if we finished it. An auto-redirect would
 				// have parsed the request, but we have no idea which status code it
 				// finished with. Since we can't see the preceding request state, we
 				// abandon its span in favor of the next.
-				clientSpan.abandon();
+				span.abandon();
 			}
 
 			// Start a new client span with the appropriate parent
@@ -154,9 +163,9 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 					null);
 			HttpClientRequestWrapper request = new HttpClientRequestWrapper(req);
 
-			clientSpan = handler().handleSendWithParent(request, parent);
-			parseConnectionAddress(connection, clientSpan);
-			ref.set(clientSpan);
+			span = handler().handleSendWithParent(request, parent);
+			parseConnectionAddress(connection, span);
+			pendingSpan.set(span);
 		}
 
 		static void parseConnectionAddress(Connection connection, Span span) {
@@ -231,18 +240,18 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 
 		void handle(Context context, @Nullable HttpClientResponse resp,
 				@Nullable Throwable error) {
-			CurrentClientSpan ref = context.getOrDefault(CurrentClientSpan.class, null);
-			if (ref == null) { // Somehow TracingMapConnect was not invoked.. skip out
-				return;
+			PendingSpan pendingSpan = context.getOrDefault(PendingSpan.class, null);
+			if (pendingSpan == null) {
+				return; // Somehow TracingMapConnect was not invoked.. skip out
 			}
 
-			Span clientSpan = ref.getAndSet(null);
-			if (clientSpan == null) {
+			Span span = pendingSpan.getAndSet(null);
+			if (span == null) {
 				return; // Unexpected. In the handle method, without a span to finish!
 			}
 			HttpClientResponseWrapper response = resp != null
 					? new HttpClientResponseWrapper(resp) : null;
-			handler().handleReceive(response, error, clientSpan);
+			handler().handleReceive(response, error, span);
 		}
 
 	}

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

@@ -18,6 +18,8 @@ package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.util.List;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -112,13 +114,6 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 	private static final Log log = LogFactory.getLog(TraceExchangeFilterFunction.class);
 
-	static final Exception CANCELLED_ERROR = new CancellationException("CANCELLED") {
-		@Override
-		public Throwable fillInStackTrace() {
-			return this; // stack trace doesn't add value here
-		}
-	};
-
 	final LazyBean<HttpTracing> httpTracing;
 
 	final Function<? super Publisher<DataBuffer>, ? extends Publisher<DataBuffer>> scopePassingTransformer;
@@ -193,13 +188,16 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 				log.debug("HttpClientHandler::handleSend: " + span);
 			}
 
+			// NOTE: We are starting the client span for the request here, but it could be
+			// canceled prior to actually being invoked. TraceWebClientSubscription will
+			// abandon this span, if cancel() happens before request().
 			this.next.exchange(wrapper.buildRequest()).subscribe(
-					new WebClientTracerSubscriber(subscriber, context, span, this));
+					new TraceWebClientSubscriber(subscriber, context, span, this));
 		}
 
 	}
 
-	private static final class WebClientTracerSubscriber
+	static final class TraceWebClientSubscriber extends AtomicReference<Span>
 			implements CoreSubscriber<ClientResponse> {
 
 		final CoreSubscriber<? super ClientResponse> actual;
@@ -209,61 +207,33 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 		@Nullable
 		final TraceContext parent;
 
-		final Span clientSpan;
-
 		final HttpClientHandler<HttpClientRequest, HttpClientResponse> handler;
 
 		final Function<? super Publisher<DataBuffer>, ? extends Publisher<DataBuffer>> scopePassingTransformer;
 
 		final CurrentTraceContext currentTraceContext;
 
-		// TODO: this isn't implemented correctly. error and success could both be called
-		boolean done;
-
-		WebClientTracerSubscriber(CoreSubscriber<? super ClientResponse> actual,
+		TraceWebClientSubscriber(CoreSubscriber<? super ClientResponse> actual,
 				Context ctx, Span clientSpan, MonoWebClientTrace mono) {
 			this.actual = actual;
 			this.parent = mono.parent;
-			this.clientSpan = clientSpan;
 			this.handler = mono.handler;
 			this.currentTraceContext = mono.currentTraceContext;
 			this.scopePassingTransformer = mono.scopePassingTransformer;
 			this.context = parent != null
 					&& !parent.equals(ctx.getOrDefault(TraceContext.class, null))
 							? ctx.put(TraceContext.class, parent) : ctx;
+			set(clientSpan);
 		}
 
 		@Override
 		public void onSubscribe(Subscription subscription) {
-			this.actual.onSubscribe(new Subscription() {
-				@Override
-				public void request(long n) {
-					try (Scope scope = currentTraceContext.maybeScope(parent)) {
-						subscription.request(n);
-					}
-				}
-
-				@Override
-				public void cancel() {
-					try (Scope scope = currentTraceContext.maybeScope(parent)) {
-						subscription.cancel();
-					}
-					finally { // TODO: this is probably incorrect as cancel happens
-								// routinely in unary subscription.
-						if (log.isDebugEnabled()) {
-							log.debug("Subscription was cancelled. Will close the span ["
-									+ clientSpan + "]");
-						}
-						handleReceive(null, CANCELLED_ERROR);
-					}
-				}
-			});
+			this.actual.onSubscribe(new TraceWebClientSubscription(subscription, this));
 		}
 
 		@Override
 		public void onNext(ClientResponse response) {
 			try (Scope scope = currentTraceContext.maybeScope(parent)) {
-				this.done = true;
 				// decorate response body
 				this.actual
 						.onNext(ClientResponse.from(response)
@@ -272,8 +242,12 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 								.build());
 			}
 			finally {
-				// TODO: is there a way to read the request at response time?
-				handleReceive(response, null);
+				Span span = getAndSet(null);
+				if (span != null) {
+					// TODO: is there a way to read the request at response time?
+					this.handler.handleReceive(new ClientResponseWrapper(response), null,
+							span);
+				}
 			}
 		}
 
@@ -283,7 +257,11 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 				this.actual.onError(t);
 			}
 			finally {
-				handleReceive(null, t);
+				Span span = getAndSet(null);
+				if (span != null) {
+					span.error(t);
+					span.finish();
+				}
 			}
 		}
 
@@ -293,13 +271,14 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 				this.actual.onComplete();
 			}
 			finally {
-				// TODO: onComplete should be after onNext. Why are we handling this?
-				if (!this.done) { // unknown state
+				Span span = getAndSet(null);
+				if (span != null) {
+					// TODO: backfill empty test:
+					// https://github.com/spring-cloud/spring-cloud-sleuth/issues/1570
 					if (log.isDebugEnabled()) {
-						log.debug("Reached OnComplete without finishing ["
-								+ this.clientSpan + "]");
+						log.debug("Reached OnComplete without finishing [" + span + "]");
 					}
-					this.clientSpan.abandon();
+					span.abandon();
 				}
 			}
 		}
@@ -309,10 +288,57 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 			return this.context;
 		}
 
-		void handleReceive(@Nullable ClientResponse res, @Nullable Throwable error) {
-			ClientResponseWrapper response = res != null ? new ClientResponseWrapper(res)
-					: null;
-			this.handler.handleReceive(response, error, clientSpan);
+	}
+
+	static class TraceWebClientSubscription extends AtomicBoolean
+			implements Subscription {
+
+		static final Exception CANCELLED_ERROR = new CancellationException("CANCELLED") {
+			@Override
+			public Throwable fillInStackTrace() {
+				return this; // stack trace doesn't add value here
+			}
+		};
+
+		final AtomicReference<Span> pendingSpan;
+
+		final Subscription delegate;
+
+		TraceWebClientSubscription(Subscription delegate,
+				AtomicReference<Span> pendingSpan) {
+			this.delegate = delegate;
+			this.pendingSpan = pendingSpan;
+		}
+
+		@Override
+		public void request(long n) {
+			if (compareAndSet(false, true)) {
+				delegate.request(n); // Not scoping to save overhead
+			}
+		}
+
+		@Override
+		public void cancel() {
+			delegate.cancel(); // Not scoping to save overhead
+
+			// Check to see if Subscription.cancel() happened after request(),
+			// but before another signal (like onComplete) completed the span.
+			Span span = pendingSpan.getAndSet(null);
+			if (span != null) {
+				if (log.isDebugEnabled()) {
+					log.debug(
+							"Subscription was cancelled. TraceWebClientBeanPostProcessor Will close the span ["
+									+ span + "]");
+				}
+
+				if (!get()) { // Subscription.request() not called: Abandon the span.
+					span.abandon();
+				}
+				else { // Request was canceled in-flight
+					span.error(CANCELLED_ERROR);
+					span.finish();
+				}
+			}
 		}
 
 	}

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

@@ -16,15 +16,23 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client;
 
+import java.util.concurrent.atomic.AtomicReference;
+
+import brave.Span;
 import org.assertj.core.api.BDDAssertions;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
+import org.reactivestreams.Subscription;
 
+import org.springframework.cloud.sleuth.instrument.web.client.TraceExchangeFilterFunction.TraceWebClientSubscription;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.web.reactive.function.client.WebClient;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.verify;
+
 /**
  * @author Marcin Grzejszczak
  */
@@ -34,6 +42,12 @@ public class TraceWebClientBeanPostProcessorTest {
 	@Mock
 	ConfigurableApplicationContext springContext;
 
+	@Mock
+	Subscription subscription;
+
+	@Mock
+	Span span;
+
 	@Test
 	public void should_add_filter_only_once_to_web_client() {
 		TraceWebClientBeanPostProcessor processor = new TraceWebClientBeanPostProcessor(
@@ -68,4 +82,31 @@ public class TraceWebClientBeanPostProcessorTest {
 		});
 	}
 
+	@Test
+	public void should_close_span_on_cancel() {
+		TraceWebClientSubscription traceSubscription = new TraceWebClientSubscription(
+				subscription, new AtomicReference<>(span));
+
+		traceSubscription.request(1);
+		traceSubscription.cancel();
+
+		verify(span).error(TraceWebClientSubscription.CANCELLED_ERROR);
+		verify(span).finish();
+
+		// Check that the ref is clear following span completion
+		assertThat(traceSubscription.pendingSpan.get()).isNull();
+	}
+
+	@Test
+	public void should_not_crash_on_cancel_when_span_clear() {
+		TraceWebClientSubscription traceSubscription = new TraceWebClientSubscription(
+				subscription, new AtomicReference<>());
+
+		traceSubscription.request(1);
+		traceSubscription.cancel();
+
+		verify(subscription).request(1);
+		verify(subscription).cancel();
+	}
+
 }

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/integration/WebClientTests.java

@@ -372,8 +372,12 @@ public class WebClientTests {
 				.contains("CLIENT");
 	}
 
+	/**
+	 * Cancel before {@link Subscription#request(long)} means a network request was never
+	 * sent
+	 */
 	@Test
-	public void shouldTagOnCancel() {
+	public void shouldNotTagOnCancel() {
 		this.webClient.get().uri("http://localhost:" + this.port + "/doNotSkip")
 				.retrieve().bodyToMono(String.class)
 				.subscribe(new BaseSubscriber<String>() {
@@ -383,8 +387,7 @@ public class WebClientTests {
 					}
 				});
 
-		then(this.reporter.getSpans()).isNotEmpty();
-		then(this.reporter.getSpans().get(0).tags()).containsEntry("error", "CANCELLED");
+		then(this.reporter.getSpans()).isEmpty();
 	}
 
 	@Test

tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/ITSpringConfiguredReactorClient.java

@@ -0,0 +1,161 @@
+/*
+ * 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;
+
+import java.net.URI;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import brave.http.HttpTracing;
+import brave.test.http.ITHttpAsyncClient;
+import io.netty.channel.ChannelOption;
+import io.netty.handler.timeout.ReadTimeoutHandler;
+import okhttp3.mockwebserver.Dispatcher;
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.RecordedRequest;
+import org.junit.Test;
+import org.reactivestreams.Subscription;
+import reactor.core.publisher.BaseSubscriber;
+import reactor.core.publisher.Mono;
+import reactor.netty.http.client.HttpClient;
+import zipkin2.Callback;
+
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * This runs Brave's integration tests, ensuring common instrumentation bugs aren't
+ * present.
+ */
+// Function of spring context so that shutdown hooks happen!
+abstract class ITSpringConfiguredReactorClient
+		extends ITHttpAsyncClient<AnnotationConfigApplicationContext> {
+
+	final Class<?>[] componentClasses;
+
+	/**
+	 * @param componentClasses configure instrumentation given {@linkplain URI baseUrl},
+	 * {@link HttpClient} and {@link HttpTracing} bindings exist.
+	 */
+	ITSpringConfiguredReactorClient(Class<?>... componentClasses) {
+		this.componentClasses = componentClasses;
+	}
+
+	@Override
+	final protected AnnotationConfigApplicationContext newClient(int port) {
+		AnnotationConfigApplicationContext result = new AnnotationConfigApplicationContext();
+		URI baseUrl = URI.create("http://127.0.0.1:" + server.getPort());
+		result.registerBean(HttpTracing.class, () -> httpTracing);
+		result.registerBean(HttpClient.class, () -> testHttpClient(baseUrl));
+		result.registerBean(URI.class, () -> baseUrl);
+		result.register(componentClasses);
+		result.refresh();
+		return result;
+	}
+
+	static HttpClient testHttpClient(URI baseUrl) {
+		return HttpClient.create().baseUrl(baseUrl.toString())
+				.tcpConfiguration(tcpClient -> tcpClient
+						.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 1000)
+						.doOnConnected(conn -> conn
+								.addHandler(new ReadTimeoutHandler(1, TimeUnit.SECONDS))))
+				.followRedirect(true);
+	}
+
+	@Override
+	final protected void closeClient(AnnotationConfigApplicationContext context) {
+		context.close(); // ensures shutdown hooks fire
+	}
+
+	@Override
+	final protected void get(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery) {
+		getMono(context, pathIncludingQuery).block();
+	}
+
+	@Override
+	final protected void post(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery, String body) {
+		postMono(context, pathIncludingQuery, body).block();
+	}
+
+	@Override
+	final protected void getAsync(AnnotationConfigApplicationContext context, String path,
+			Callback<Integer> callback) {
+		TestHttpCallbackSubscriber.subscribe(getMono(context, path), callback);
+	}
+
+	/** Returns a {@link Mono} of the HTTP status code from the given "POST" request. */
+	abstract Mono<Integer> postMono(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery, String body);
+
+	/** Returns a {@link Mono} of the HTTP status code. */
+	abstract Mono<Integer> getMono(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery);
+
+	/**
+	 * This assumes that implementations do not issue an HTTP request until
+	 * {@link Subscription#request(long)} is called. Since a client span is only for
+	 * remote operations, we should not create one when we know a network request won't
+	 * happen. In this case, we ensure a canceled subscription doesn't end up traced.
+	 */
+	@Test
+	public void cancelledSubscription_doesntTrace() throws Exception {
+		CountDownLatch latch = new CountDownLatch(1);
+
+		BaseSubscriber<Integer> subscriber = new BaseSubscriber<Integer>() {
+			@Override
+			protected void hookOnSubscribe(Subscription subscription) {
+				subscription.cancel();
+				latch.countDown();
+			}
+		};
+
+		getMono(client, "/foo").subscribe(subscriber);
+
+		latch.await();
+
+		assertThat(server.getRequestCount()).isZero();
+		// post-conditions will prove no span was created
+	}
+
+	@Test
+	public void cancelInFlight() throws Exception {
+		BaseSubscriber<Integer> subscriber = new BaseSubscriber<Integer>() {
+		};
+
+		CountDownLatch latch = new CountDownLatch(1);
+
+		server.setDispatcher(new Dispatcher() {
+			@Override
+			public MockResponse dispatch(RecordedRequest request) {
+				subscriber.cancel();
+				latch.countDown();
+				return new MockResponse();
+			}
+		});
+
+		getMono(client, "/foo").subscribe(subscriber);
+
+		latch.await();
+
+		assertThat(server.getRequestCount()).isOne();
+		assertThat(takeSpan().tags()).containsKey("error");
+	}
+
+}

tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/ReactorNettyHttpClientBraveTests.java

@@ -16,65 +16,23 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client;
 
-import java.util.concurrent.TimeUnit;
-
-import brave.http.HttpTracing;
-import brave.test.http.ITHttpAsyncClient;
-import io.netty.channel.ChannelOption;
-import io.netty.handler.timeout.ReadTimeoutHandler;
 import org.junit.Ignore;
 import org.junit.Test;
 import reactor.core.publisher.Mono;
 import reactor.netty.ByteBufFlux;
 import reactor.netty.http.client.HttpClient;
-import reactor.netty.http.client.HttpClientResponse;
-import zipkin2.Callback;
 
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 
-/**
- * This runs Brave's integration tests, ensuring common instrumentation bugs aren't
- * present.
- */
-// Function of spring context so that shutdown hooks happen!
-public class ReactorNettyHttpClientBraveTests
-		extends ITHttpAsyncClient<AnnotationConfigApplicationContext> {
+public class ReactorNettyHttpClientBraveTests extends ITSpringConfiguredReactorClient {
 
 	/**
 	 * This uses Spring to instrument the {@link HttpClient} using a
 	 * {@link BeanPostProcessor}.
 	 */
-	@Override
-	protected AnnotationConfigApplicationContext newClient(int port) {
-		AnnotationConfigApplicationContext result = new AnnotationConfigApplicationContext();
-		result.registerBean(HttpTracing.class, () -> httpTracing);
-		result.registerBean(HttpClient.class,
-				() -> testHttpClient().baseUrl("http://127.0.0.1:" + port));
-		result.register(HttpClientBeanPostProcessor.class);
-		result.refresh();
-		return result;
-	}
-
-	static HttpClient testHttpClient() {
-		return HttpClient.create()
-				.tcpConfiguration(tcpClient -> tcpClient
-						.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 1000)
-						.doOnConnected(conn -> conn
-								.addHandler(new ReadTimeoutHandler(1, TimeUnit.SECONDS))))
-				.followRedirect(true);
-	}
-
-	@Override
-	protected void closeClient(AnnotationConfigApplicationContext context) {
-		context.close(); // ensures shutdown hooks fire
-	}
-
-	@Override
-	protected void get(AnnotationConfigApplicationContext context,
-			String pathIncludingQuery) {
-		context.getBean(HttpClient.class).get().uri(pathIncludingQuery).response()
-				.block();
+	public ReactorNettyHttpClientBraveTests() {
+		super(HttpClientBeanPostProcessor.class);
 	}
 
 	@Test
@@ -127,20 +85,18 @@ public class ReactorNettyHttpClientBraveTests
 	}
 
 	@Override
-	protected void post(AnnotationConfigApplicationContext context,
+	Mono<Integer> postMono(AnnotationConfigApplicationContext context,
 			String pathIncludingQuery, String body) {
-		context.getBean(HttpClient.class).post()
+		return context.getBean(HttpClient.class).post()
 				.send(ByteBufFlux.fromString(Mono.just(body))).uri(pathIncludingQuery)
-				.response().block();
+				.response().map(r -> r.status().code());
 	}
 
 	@Override
-	protected void getAsync(AnnotationConfigApplicationContext context, String path,
-			Callback<Integer> callback) {
-		Mono<HttpClientResponse> request = context.getBean(HttpClient.class).get()
-				.uri(path).response();
-
-		TestHttpCallbackSubscriber.subscribe(request, r -> r.status().code(), callback);
+	Mono<Integer> getMono(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery) {
+		return context.getBean(HttpClient.class).get().uri(pathIncludingQuery).response()
+				.map(r -> r.status().code());
 	}
 
 }

tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/TestHttpCallbackSubscriber.java

@@ -17,7 +17,6 @@
 package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Function;
 
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
@@ -35,22 +34,17 @@ import zipkin2.Callback;
  * The implementation forwards signals to the supplied {@link Callback}, enforcing
  * assumptions about a non-empty, {@link Mono} subscription.
  */
-final class TestHttpCallbackSubscriber<T> implements CoreSubscriber<T> {
+final class TestHttpCallbackSubscriber implements CoreSubscriber<Integer> {
 
-	static <T> void subscribe(Mono<T> mono, Function<T, Integer> statusCodeFunction,
-			Callback<Integer> callback) {
-		mono.subscribe(new TestHttpCallbackSubscriber<>(statusCodeFunction, callback));
+	static void subscribe(Mono<Integer> mono, Callback<Integer> callback) {
+		mono.subscribe(new TestHttpCallbackSubscriber(callback));
 	}
 
-	final Function<T, Integer> statusCodeFunction;
-
 	final Callback<Integer> callback;
 
 	final AtomicReference<Subscription> ref = new AtomicReference<>();
 
-	private TestHttpCallbackSubscriber(Function<T, Integer> statusCodeFunction,
-			Callback<Integer> callback) {
-		this.statusCodeFunction = statusCodeFunction;
+	private TestHttpCallbackSubscriber(Callback<Integer> callback) {
 		this.callback = callback;
 	}
 
@@ -67,9 +61,9 @@ final class TestHttpCallbackSubscriber<T> implements CoreSubscriber<T> {
 	}
 
 	@Override
-	public void onNext(T t) {
+	public void onNext(Integer t) {
 		if (ref.getAndSet(null) != null) {
-			callback.onSuccess(statusCodeFunction.apply(t));
+			callback.onSuccess(t);
 		}
 		else {
 			// This is a Mono, which doesn't signal onNext() twice. If we reach here,

tests/spring-cloud-sleuth-instrumentation-reactor-tests/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/WebClientBraveTests.java

@@ -16,71 +16,54 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client;
 
-import brave.http.HttpTracing;
-import brave.test.http.ITHttpAsyncClient;
+import java.net.URI;
+
 import org.junit.Ignore;
 import org.junit.Test;
 import reactor.core.publisher.Mono;
 import reactor.netty.http.client.HttpClient;
-import zipkin2.Callback;
 
 import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.boot.autoconfigure.web.reactive.function.client.WebClientAutoConfiguration;
+import org.springframework.boot.web.reactive.function.client.WebClientCustomizer;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.http.client.reactive.ClientHttpConnector;
+import org.springframework.core.annotation.Order;
 import org.springframework.http.client.reactive.ReactorClientHttpConnector;
 import org.springframework.web.reactive.function.BodyInserters;
 import org.springframework.web.reactive.function.client.ClientResponse;
 import org.springframework.web.reactive.function.client.WebClient;
 
 /**
- * This runs Brave's integration tests without underlying instrumentation, which would
- * happen when a 3rd party client like Jetty is in use.
+ * This runs Brave's integration tests without underlying instrumentation, which is
+ * default in Spring Boot due to static instantiation in
+ * {@link org.springframework.boot.autoconfigure.web.reactive.function.client.ClientHttpConnectorConfiguration}.
  */
-// Function of spring context so that shutdown hooks happen!
-public class WebClientBraveTests
-		extends ITHttpAsyncClient<AnnotationConfigApplicationContext> {
+public class WebClientBraveTests extends ITSpringConfiguredReactorClient {
 
 	/**
 	 * This uses Spring to instrument the {@link WebClient} using a
 	 * {@link BeanPostProcessor}.
 	 */
-	@Override
-	protected AnnotationConfigApplicationContext newClient(int port) {
-		AnnotationConfigApplicationContext result = new AnnotationConfigApplicationContext();
-		result.registerBean(HttpTracing.class, () -> httpTracing);
-		result.register(WebClientBuilderConfiguration.class);
-		result.register(TraceWebClientBeanPostProcessor.class);
-		result.refresh();
-		return result;
+	public WebClientBraveTests() {
+		super(WebClientConfiguration.class, WebClientAutoConfiguration.class,
+				TraceWebClientBeanPostProcessor.class);
 	}
 
 	@Override
-	protected void closeClient(AnnotationConfigApplicationContext context) {
-		context.close(); // ensures shutdown hooks fire
-	}
-
-	@Override
-	protected void get(AnnotationConfigApplicationContext context,
-			String pathIncludingQuery) {
-		client(context).get().uri(pathIncludingQuery).exchange().block();
-	}
-
-	@Override
-	protected void post(AnnotationConfigApplicationContext context,
+	Mono<Integer> postMono(AnnotationConfigApplicationContext context,
 			String pathIncludingQuery, String body) {
-		client(context).post().uri(pathIncludingQuery).body(BodyInserters.fromValue(body))
-				.exchange().block();
+		return context.getBean(WebClient.Builder.class).build().post()
+				.uri(pathIncludingQuery).body(BodyInserters.fromValue(body)).exchange()
+				.map(ClientResponse::rawStatusCode);
 	}
 
 	@Override
-	protected void getAsync(AnnotationConfigApplicationContext context, String path,
-			Callback<Integer> callback) {
-		Mono<ClientResponse> request = client(context).get().uri(path).exchange();
-
-		TestHttpCallbackSubscriber.subscribe(request, ClientResponse::rawStatusCode,
-				callback);
+	Mono<Integer> getMono(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery) {
+		return context.getBean(WebClient.Builder.class).build().get()
+				.uri(pathIncludingQuery).exchange().map(ClientResponse::rawStatusCode);
 	}
 
 	@Test
@@ -101,31 +84,19 @@ public class WebClientBraveTests
 	public void readsRequestAtResponseTime() {
 	}
 
-	WebClient client(AnnotationConfigApplicationContext context) {
-		return context.getBean(WebClient.Builder.class)
-				.baseUrl("http://127.0.0.1:" + server.getPort()).build();
-	}
-
-	/**
-	 * This fakes auto-configuration which wouldn't configure reactor's trace
-	 * instrumentation.
-	 */
 	@Configuration
-	static class WebClientBuilderConfiguration {
+	static class WebClientConfiguration {
 
+		/**
+		 * Normally, the HTTP connector would be statically initialized. This ensures the
+		 * {@link HttpClient} is configured for the mock endpoint.
+		 */
 		@Bean
-		HttpClient httpClient() {
-			return ReactorNettyHttpClientBraveTests.testHttpClient();
-		}
-
-		@Bean
-		ClientHttpConnector clientHttpConnector(HttpClient httpClient) {
-			return new ReactorClientHttpConnector(httpClient);
-		}
-
-		@Bean
-		WebClient.Builder webClientBuilder(ClientHttpConnector clientHttpConnector) {
-			return WebClient.builder().clientConnector(clientHttpConnector);
+		@Order(0)
+		public WebClientCustomizer clientConnectorCustomizer(HttpClient httpClient,
+				URI baseUrl) {
+			return (builder) -> builder.baseUrl(baseUrl.toString())
+					.clientConnector(new ReactorClientHttpConnector(httpClient));
 		}
 
 	}