Fixes propagation break and error handling in reactor-netty HttpClient (#1552)

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

@@ -31,11 +31,13 @@ import reactor.netty.Connection;
 import reactor.netty.http.client.HttpClient;
 import reactor.netty.http.client.HttpClientRequest;
 import reactor.netty.http.client.HttpClientResponse;
+import reactor.util.context.Context;
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.cloud.sleuth.internal.LazyBean;
 import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.lang.Nullable;
 
 class HttpClientBeanPostProcessor implements BeanPostProcessor {
 
@@ -51,11 +53,15 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 		LazyBean<HttpTracing> httpTracing = LazyBean.create(this.springContext,
 				HttpTracing.class);
 		if (bean instanceof HttpClient) {
-			return ((HttpClient) bean).mapConnect(new TracingMapConnect(httpTracing))
-					.doOnRequest(TracingDoOnRequest.create(httpTracing))
-					.doOnRequestError(TracingDoOnErrorRequest.create(httpTracing))
-					.doOnResponse(TracingDoOnResponse.create(httpTracing))
-					.doOnResponseError(TracingDoOnErrorResponse.create(httpTracing));
+			// This adds handlers to manage the span lifecycle. All require explicit
+			// propagation of the current span as a reactor context property.
+			// This done in mapConnect, added last so that it is setup first.
+			// https://projectreactor.io/docs/core/release/reference/#_simple_context_examples
+			return ((HttpClient) bean).doOnRequest(new TracingDoOnRequest(httpTracing))
+					.doOnRequestError(new TracingDoOnErrorRequest(httpTracing))
+					.doOnResponse(new TracingDoOnResponse(httpTracing))
+					.doOnResponseError(new TracingDoOnErrorResponse(httpTracing))
+					.mapConnect(new TracingMapConnect(httpTracing));
 		}
 		return bean;
 	}
@@ -74,11 +80,12 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 		@Override
 		public Mono<? extends Connection> apply(Mono<? extends Connection> mono,
 				Bootstrap bootstrap) {
+			// This is read in this class and also inside ScopePassingSpanSubscriber
 			return mono.subscriberContext(context -> context.put(AtomicReference.class,
 					new AtomicReference<>(tracer().currentSpan())));
 		}
 
-		private Tracer tracer() {
+		Tracer tracer() {
 			if (this.tracer == null) {
 				this.tracer = this.httpTracing.get().tracing().tracer();
 			}
@@ -100,10 +107,6 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			this.httpTracing = httpTracing;
 		}
 
-		static TracingDoOnRequest create(LazyBean<HttpTracing> httpTracing) {
-			return new TracingDoOnRequest(httpTracing);
-		}
-
 		List<String> propagationKeys() {
 			if (this.propagationKeys == null) {
 				this.propagationKeys = httpTracing.get().tracing().propagation().keys();
@@ -128,12 +131,21 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 					return;
 				}
 			}
-			AtomicReference<Span> reference = req.currentContext()
-					.getOrDefault(AtomicReference.class, new AtomicReference<>());
+
+			// Look for a parent propagated by TracingMapConnect
+			AtomicReference<Span> ref = req.currentContext()
+					.getOrDefault(AtomicReference.class, null);
+			Span parent = ref != null ? ref.get() : null;
+
+			// Start a new client span with the appropriate parent
 			WrappedHttpClientRequest request = new WrappedHttpClientRequest(req);
-			Span span = reference.get() == null ? handler().handleSend(request)
-					: handler().handleSend(request, reference.get());
-			reference.set(span);
+			Span clientSpan = parent != null ? handler().handleSend(request, parent)
+					: handler().handleSend(request);
+
+			// Swap the ref with the client span, so that other hooks can see it
+			if (ref != null) {
+				ref.set(clientSpan);
+			}
 		}
 
 	}
@@ -145,13 +157,9 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			super(httpTracing);
 		}
 
-		static TracingDoOnResponse create(LazyBean<HttpTracing> httpTracing) {
-			return new TracingDoOnResponse(httpTracing);
-		}
-
 		@Override
-		public void accept(HttpClientResponse httpClientResponse, Connection connection) {
-			handle(httpClientResponse, null);
+		public void accept(HttpClientResponse response, Connection connection) {
+			handle(response.currentContext(), response, null);
 		}
 
 	}
@@ -163,13 +171,10 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			super(httpTracing);
 		}
 
-		static TracingDoOnErrorRequest create(LazyBean<HttpTracing> httpTracing) {
-			return new TracingDoOnErrorRequest(httpTracing);
-		}
-
 		@Override
-		public void accept(HttpClientRequest request, Throwable throwable) {
-			handle(null, throwable);
+		public void accept(HttpClientRequest request, Throwable error) {
+			// TODO: the current context here does not have the AtomicReference<Span>
+			handle(request.currentContext(), null, error);
 		}
 
 	}
@@ -181,13 +186,9 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			super(httpTracing);
 		}
 
-		static TracingDoOnErrorResponse create(LazyBean<HttpTracing> httpTracing) {
-			return new TracingDoOnErrorResponse(httpTracing);
-		}
-
 		@Override
-		public void accept(HttpClientResponse httpClientResponse, Throwable throwable) {
-			handle(httpClientResponse, throwable);
+		public void accept(HttpClientResponse response, Throwable error) {
+			handle(response.currentContext(), response, error);
 		}
 
 	}
@@ -209,18 +210,16 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			return this.handler;
 		}
 
-		protected void handle(HttpClientResponse httpClientResponse,
-				Throwable throwable) {
-			if (httpClientResponse == null) {
-				return;
+		void handle(Context context, @Nullable HttpClientResponse resp,
+				@Nullable Throwable error) {
+			AtomicReference<Span> ref = context.getOrDefault(AtomicReference.class, null);
+			Span span = ref != null ? ref.get() : null;
+			if (span == null) {
+				return; // Unexpected. In the handle method, without a span to finish!
 			}
-			AtomicReference reference = httpClientResponse.currentContext()
-					.getOrDefault(AtomicReference.class, null);
-			if (reference == null || reference.get() == null) {
-				return;
-			}
-			handler().handleReceive(new WrappedHttpClientResponse(httpClientResponse),
-					throwable, (Span) reference.get());
+			WrappedHttpClientResponse response = resp != null
+					? new WrappedHttpClientResponse(resp) : null;
+			handler().handleReceive(response, error, span);
 		}
 
 	}

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

@@ -30,6 +30,7 @@ import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.netty.DisposableServer;
 import reactor.netty.http.client.HttpClient;
+import reactor.netty.http.client.PrematureCloseException;
 import reactor.netty.http.server.HttpServer;
 import zipkin2.Span;
 import zipkin2.reporter.Reporter;
@@ -43,6 +44,7 @@ import org.springframework.test.context.junit4.SpringRunner;
 import org.springframework.web.reactive.function.client.WebClient;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
  * This tests {@link HttpClient} instrumentation performed by
@@ -93,6 +95,23 @@ public class ReactorNettyHttpClientSpringBootTests {
 				.isEqualTo(clientSpan.traceId() + "-" + clientSpan.id() + "-1");
 	}
 
+	@Test
+	public void shouldTagOnRequestError() throws InterruptedException {
+		disposableServer = HttpServer.create().port(0).handle((req, resp) -> {
+			throw new RuntimeException("test");
+		}).bindNow();
+
+		Mono<String> request = httpClient.port(disposableServer.port()).get().uri("/")
+				.responseContent().aggregate().asString();
+
+		assertThatThrownBy(request::block)
+				.hasCauseInstanceOf(PrematureCloseException.class);
+
+		Span clientSpan = takeClientSpan();
+
+		assertThat(clientSpan.tags()).containsKey("error");
+	}
+
 	/** Call this to block until a span was reported */
 	Span takeClientSpan() throws InterruptedException {
 		Span result = spans.poll(1, TimeUnit.SECONDS);