Ensures spring context is closed in brave tests (#1572)

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

@@ -182,6 +182,7 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 
 		@Override
 		public void accept(HttpClientResponse response, Connection connection) {
+			// TODO: is there a way to read the request at response time?
 			handle(response.currentContext(), response, null);
 		}
 

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

@@ -220,6 +220,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 		final CurrentTraceContext currentTraceContext;
 
+		// TODO: this isn't implemented correctly. error and success could both be called
 		boolean done;
 
 		WebClientTracerSubscriber(CoreSubscriber<? super ClientResponse> actual,
@@ -251,7 +252,8 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 					try (Scope scope = currentTraceContext.maybeScope(parent)) {
 						subscription.cancel();
 					}
-					finally {
+					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 + "]");
@@ -274,6 +276,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 								.build());
 			}
 			finally {
+				// TODO: is there a way to read the request at response time?
 				handleReceive(response, null);
 			}
 		}

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

@@ -17,26 +17,17 @@
 package org.springframework.cloud.sleuth.instrument.web.client;
 
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
 
 import brave.http.HttpTracing;
 import brave.test.http.ITHttpAsyncClient;
 import io.netty.channel.ChannelOption;
 import io.netty.handler.timeout.ReadTimeoutHandler;
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
-import org.reactivestreams.Subscription;
-import reactor.core.CoreSubscriber;
-import reactor.core.publisher.Hooks;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.Operators;
-import reactor.core.scheduler.Schedulers;
 import reactor.netty.ByteBufFlux;
 import reactor.netty.http.client.HttpClient;
 import reactor.netty.http.client.HttpClientResponse;
-import reactor.util.context.Context;
 import zipkin2.Callback;
 
 import org.springframework.beans.factory.config.BeanPostProcessor;
@@ -46,30 +37,23 @@ import org.springframework.context.annotation.AnnotationConfigApplicationContext
  * This runs Brave's integration tests, ensuring common instrumentation bugs aren't
  * present.
  */
-public class ReactorNettyHttpClientBraveTests extends ITHttpAsyncClient<HttpClient> {
-	@Before
-	@After
-	public void resetHooks() {
-		// There's an assumption some other test is leaking hooks, so we clear them all to
-		// prevent should_not_scope_scalar_subscribe from being interfered with.
-		Hooks.resetOnEachOperator();
-		Hooks.resetOnLastOperator();
-		Schedulers.removeExecutorServiceDecorator("sleuth");
-	}
+// Function of spring context so that shutdown hooks happen!
+public class ReactorNettyHttpClientBraveTests
+		extends ITHttpAsyncClient<AnnotationConfigApplicationContext> {
 
 	/**
 	 * This uses Spring to instrument the {@link HttpClient} using a
 	 * {@link BeanPostProcessor}.
 	 */
 	@Override
-	protected HttpClient newClient(int port) {
+	protected AnnotationConfigApplicationContext newClient(int port) {
 		AnnotationConfigApplicationContext result = new AnnotationConfigApplicationContext();
 		result.registerBean(HttpTracing.class, () -> httpTracing);
 		result.registerBean(HttpClient.class,
-				ReactorNettyHttpClientBraveTests::testHttpClient);
+				() -> testHttpClient().baseUrl("http://127.0.0.1:" + port));
 		result.register(HttpClientBeanPostProcessor.class);
 		result.refresh();
-		return result.getBean(HttpClient.class).baseUrl("http://127.0.0.1:" + port);
+		return result;
 	}
 
 	static HttpClient testHttpClient() {
@@ -82,21 +66,29 @@ public class ReactorNettyHttpClientBraveTests extends ITHttpAsyncClient<HttpClie
 	}
 
 	@Override
-	protected void closeClient(HttpClient client) {
-		// HttpClient is not Closeable
+	protected void closeClient(AnnotationConfigApplicationContext context) {
+		context.close(); // ensures shutdown hooks fire
 	}
 
 	@Override
-	protected void get(HttpClient client, String pathIncludingQuery) {
-		client.get().uri(pathIncludingQuery).response().block();
+	protected void get(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery) {
+		context.getBean(HttpClient.class).get().uri(pathIncludingQuery).response()
+				.block();
 	}
 
 	@Test
-	@Ignore("TODO: consider integrating TracingMapConnect with ScopePassingSpanSubscriber")
+	@Ignore("TODO: NPE reading context: consider integrating TracingMapConnect with ScopePassingSpanSubscriber")
 	@Override
 	public void callbackContextIsFromInvocationTime() {
 	}
 
+	@Test
+	@Ignore("TODO: False negative due to NPE reading context: remove after Brave 5.10")
+	@Override
+	public void asyncRootSpan() {
+	}
+
 	@Test
 	@Ignore("TODO: reactor/reactor-netty#1000")
 	@Override
@@ -128,60 +120,20 @@ public class ReactorNettyHttpClientBraveTests extends ITHttpAsyncClient<HttpClie
 	}
 
 	@Override
-	protected void post(HttpClient client, String pathIncludingQuery, String body) {
-		client.post().send(ByteBufFlux.fromString(Mono.just(body)))
-				.uri(pathIncludingQuery).response().block();
+	protected void post(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery, String body) {
+		context.getBean(HttpClient.class).post()
+				.send(ByteBufFlux.fromString(Mono.just(body))).uri(pathIncludingQuery)
+				.response().block();
 	}
 
 	@Override
-	protected void getAsync(HttpClient client, String path, Callback<Void> callback) {
-		Mono<HttpClientResponse> request = client.get().uri(path).response();
+	protected void getAsync(AnnotationConfigApplicationContext context, String path,
+			Callback<Void> callback) {
+		Mono<HttpClientResponse> request = context.getBean(HttpClient.class).get()
+				.uri(path).response();
 
-		request.subscribe(new CoreSubscriber<HttpClientResponse>() {
-
-			final AtomicReference<Subscription> ref = new AtomicReference<>();
-
-			@Override
-			public void onSubscribe(Subscription s) {
-				if (Operators.validate(ref.getAndSet(s), s)) {
-					s.request(Long.MAX_VALUE);
-				}
-				else {
-					s.cancel();
-				}
-			}
-
-			@Override
-			public void onNext(HttpClientResponse t) {
-				Subscription s = ref.getAndSet(null);
-				if (s != null) {
-					callback.onSuccess(null);
-					s.cancel();
-				}
-				else {
-					Operators.onNextDropped(t, currentContext());
-				}
-			}
-
-			@Override
-			public void onError(Throwable t) {
-				if (ref.getAndSet(null) != null) {
-					callback.onError(t);
-				}
-			}
-
-			@Override
-			public void onComplete() {
-				if (ref.getAndSet(null) != null) {
-					callback.onSuccess(null);
-				}
-			}
-
-			@Override
-			public Context currentContext() {
-				return Context.empty();
-			}
-		});
+		TestCallbackSubscriber.subscribe(request, callback);
 	}
 
 }

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

@@ -0,0 +1,102 @@
+/*
+ * 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;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Mono;
+import reactor.core.publisher.Operators;
+import reactor.util.context.Context;
+import zipkin2.Callback;
+
+/**
+ * {@link #subscribe} is made for reactor-netty and WebFlux client requests used in tests.
+ * This does more assertions than normal, to ensure instrumentation isn't redundantly
+ * signalling, or missing signals.
+ *
+ * <p>
+ * The implementation forwards signals to the supplied {@link Callback}, enforcing
+ * assumptions about a non-empty, {@link Mono} subscription.
+ */
+final class TestCallbackSubscriber<T> implements CoreSubscriber<T> {
+
+	static <T> void subscribe(Mono<T> mono, Callback<Void> callback) {
+		mono.subscribe(new TestCallbackSubscriber<>(callback));
+	}
+
+	final Callback<Void> callback;
+
+	final AtomicReference<Subscription> ref = new AtomicReference<>();
+
+	private TestCallbackSubscriber(Callback<Void> callback) {
+		this.callback = callback;
+	}
+
+	@Override
+	public void onSubscribe(Subscription s) {
+		if (Operators.validate(ref.getAndSet(s), s)) {
+			s.request(Long.MAX_VALUE);
+		}
+		else {
+			// We don't intentionally call subscribe() multiple times in our tests. If we
+			// reach here, possibly instrumentation is redundantly subscribing.
+			callback.onError(new AssertionError("onSubscribe() called twice!"));
+		}
+	}
+
+	@Override
+	public void onNext(T t) {
+		if (ref.getAndSet(null) != null) {
+			callback.onSuccess(null /* because Void */);
+		}
+		else {
+			// This is a Mono, which doesn't signal onNext() twice. If we reach here,
+			// possibly instrumentation is signaling twice.
+			callback.onError(new AssertionError("onNext() called twice!"));
+		}
+	}
+
+	@Override
+	public void onError(Throwable t) {
+		if (ref.getAndSet(null) != null) {
+			callback.onError(t);
+		}
+		else {
+			// We don't expect onError() to signal twice. If we reach here, possibly
+			// instrumentation is signaling twice or onSuccess() threw an exception.
+			callback.onError(new AssertionError("onError() called twice: " + t, t));
+		}
+	}
+
+	@Override
+	public void onComplete() {
+		if (ref.getAndSet(null) != null) {
+			// Tests make a non-empty Mono subscription, which should not signal
+			// onComplete() before onNext(). If we reach here, possibly instrumentation
+			// is not signaling onNext() when it should.
+			callback.onError(new AssertionError("onComplete() called before onNext!"));
+		}
+	}
+
+	@Override
+	public Context currentContext() {
+		return Context.empty();
+	}
+
+}

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

@@ -16,27 +16,15 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client;
 
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
-
 import brave.http.HttpTracing;
 import brave.test.http.ITHttpAsyncClient;
-import io.netty.channel.ChannelOption;
-import io.netty.handler.timeout.ReadTimeoutHandler;
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Ignore;
 import org.junit.Test;
-import org.reactivestreams.Subscription;
-import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.Operators;
 import reactor.netty.http.client.HttpClient;
-import reactor.util.context.Context;
 import zipkin2.Callback;
 
 import org.springframework.beans.factory.config.BeanPostProcessor;
-import org.springframework.cloud.sleuth.instrument.reactor.ScopePassingSpanSubscriberTests;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -49,94 +37,48 @@ 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.
  */
-public class WebClientBraveTests extends ITHttpAsyncClient<WebClient> {
-
-	@Before
-	@After
-	public void resetHooks() {
-		new ScopePassingSpanSubscriberTests().resetHooks();
-	}
+// Function of spring context so that shutdown hooks happen!
+public class WebClientBraveTests
+		extends ITHttpAsyncClient<AnnotationConfigApplicationContext> {
 
 	/**
 	 * This uses Spring to instrument the {@link WebClient} using a
 	 * {@link BeanPostProcessor}.
 	 */
 	@Override
-	protected WebClient newClient(int port) {
+	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.getBean(WebClient.Builder.class).baseUrl("http://127.0.0.1:" + port)
-				.build();
+		return result;
 	}
 
 	@Override
-	protected void closeClient(WebClient client) {
-		// WebClient is not Closeable
+	protected void closeClient(AnnotationConfigApplicationContext context) {
+		context.close(); // ensures shutdown hooks fire
 	}
 
 	@Override
-	protected void get(WebClient client, String pathIncludingQuery) {
-		client.get().uri(pathIncludingQuery).exchange().block();
+	protected void get(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery) {
+		client(context).get().uri(pathIncludingQuery).exchange().block();
 	}
 
 	@Override
-	protected void post(WebClient client, String pathIncludingQuery, String body) {
-		client.post().uri(pathIncludingQuery).body(BodyInserters.fromValue(body))
+	protected void post(AnnotationConfigApplicationContext context,
+			String pathIncludingQuery, String body) {
+		client(context).post().uri(pathIncludingQuery).body(BodyInserters.fromValue(body))
 				.exchange().block();
 	}
 
 	@Override
-	protected void getAsync(WebClient client, String path, Callback<Void> callback) {
-		Mono<ClientResponse> request = client.get().uri(path).exchange();
+	protected void getAsync(AnnotationConfigApplicationContext context, String path,
+			Callback<Void> callback) {
+		Mono<ClientResponse> request = client(context).get().uri(path).exchange();
 
-		request.subscribe(new CoreSubscriber<ClientResponse>() {
-
-			final AtomicReference<Subscription> ref = new AtomicReference<>();
-
-			@Override
-			public void onSubscribe(Subscription s) {
-				if (Operators.validate(ref.getAndSet(s), s)) {
-					s.request(Long.MAX_VALUE);
-				}
-				else {
-					s.cancel();
-				}
-			}
-
-			@Override
-			public void onNext(ClientResponse t) {
-				Subscription s = ref.getAndSet(null);
-				if (s != null) {
-					callback.onSuccess(null);
-					s.cancel();
-				}
-				else {
-					Operators.onNextDropped(t, currentContext());
-				}
-			}
-
-			@Override
-			public void onError(Throwable t) {
-				if (ref.getAndSet(null) != null) {
-					callback.onError(t);
-				}
-			}
-
-			@Override
-			public void onComplete() {
-				if (ref.getAndSet(null) != null) {
-					callback.onSuccess(null);
-				}
-			}
-
-			@Override
-			public Context currentContext() {
-				return Context.empty();
-			}
-		});
+		TestCallbackSubscriber.subscribe(request, callback);
 	}
 
 	@Test
@@ -151,6 +93,11 @@ public class WebClientBraveTests extends ITHttpAsyncClient<WebClient> {
 	public void reportsServerAddress() {
 	}
 
+	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.
@@ -160,13 +107,7 @@ public class WebClientBraveTests extends ITHttpAsyncClient<WebClient> {
 
 		@Bean
 		HttpClient httpClient() {
-			// TODO: ReactorNettyHttpClientBraveTests.testHttpClient() #1554
-			return HttpClient.create()
-					.tcpConfiguration(tcpClient -> tcpClient
-							.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 1000)
-							.doOnConnected(conn -> conn.addHandler(
-									new ReadTimeoutHandler(1, TimeUnit.SECONDS))))
-					.followRedirect(true);
+			return ReactorNettyHttpClientBraveTests.testHttpClient();
 		}
 
 		@Bean