From 1e9d18226e1484cc99a4bdddfaf3d1263df4df09 Mon Sep 17 00:00:00 2001
From: Adrian Cole <adriancole@users.noreply.github.com>
Date: Mon, 2 Mar 2020 07:53:04 +0800
Subject: [PATCH] Fixes WebClient which started a client span prior to a
 subscription (#1576)

In looking at underlying HttpClient mechanics, I noticed the reactor
call doesn't happen until subscribe. Before this change, we started the
client span at the ExchangeFilterFunction, not at subscribe time.
---
 .../client/HttpClientBeanPostProcessor.java   | 10 ++----
 .../TraceWebClientBeanPostProcessor.java      | 32 ++++++++-----------
 2 files changed, 16 insertions(+), 26 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 7e45a9b2b..888377a63 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
@@ -25,7 +25,6 @@ import brave.Span;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
 import brave.propagation.CurrentTraceContext;
-import brave.propagation.CurrentTraceContext.Scope;
 import brave.propagation.TraceContext;
 import io.netty.bootstrap.Bootstrap;
 import reactor.core.publisher.Mono;
@@ -155,12 +154,9 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 					null);
 			WrappedHttpClientRequest request = new WrappedHttpClientRequest(req);
 
-			// Simplify after openzipkin/brave#1082
-			try (Scope ws = currentTraceContext().maybeScope(parent)) {
-				clientSpan = handler().handleSend(request);
-				parseConnectionAddress(connection, clientSpan);
-				ref.set(clientSpan);
-			}
+			clientSpan = handler().handleSendWithParent(request, parent);
+			parseConnectionAddress(connection, clientSpan);
+			ref.set(clientSpan);
 		}
 
 		static void parseConnectionAddress(Connection connection, Span span) {
diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java
index 86d1c6f65..97faff492 100644
--- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java
+++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientBeanPostProcessor.java
@@ -138,14 +138,7 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 	@Override
 	public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
-		HttpClientRequest wrapper = new HttpClientRequest(request);
-		TraceContext parent = currentTraceContext().get();
-		Span clientSpan = handler().handleSend(wrapper);
-		if (log.isDebugEnabled()) {
-			log.debug("HttpClientHandler::handleSend: " + clientSpan);
-		}
-		return new MonoWebClientTrace(next, wrapper.buildRequest(), this, parent,
-				clientSpan);
+		return new MonoWebClientTrace(next, request, this);
 	}
 
 	CurrentTraceContext currentTraceContext() {
@@ -177,18 +170,14 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 		@Nullable
 		final TraceContext parent;
 
-		private final Span span;
-
 		MonoWebClientTrace(ExchangeFunction next, ClientRequest request,
-				TraceExchangeFilterFunction filterFunction, @Nullable TraceContext parent,
-				Span span) {
+				TraceExchangeFilterFunction filterFunction) {
 			this.next = next;
 			this.request = request;
 			this.handler = filterFunction.handler();
 			this.currentTraceContext = filterFunction.currentTraceContext();
 			this.scopePassingTransformer = filterFunction.scopePassingTransformer;
-			this.parent = parent;
-			this.span = span;
+			this.parent = currentTraceContext.get();
 		}
 
 		@Override
@@ -196,8 +185,14 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 			Context context = subscriber.currentContext();
 
-			this.next.exchange(request).subscribe(new WebClientTracerSubscriber(
-					subscriber, context, parent, span, this));
+			HttpClientRequest wrapper = new HttpClientRequest(request);
+			Span span = handler.handleSendWithParent(wrapper, parent);
+			if (log.isDebugEnabled()) {
+				log.debug("HttpClientHandler::handleSend: " + span);
+			}
+
+			this.next.exchange(wrapper.buildRequest()).subscribe(
+					new WebClientTracerSubscriber(subscriber, context, span, this));
 		}
 
 	}
@@ -224,10 +219,9 @@ final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 		boolean done;
 
 		WebClientTracerSubscriber(CoreSubscriber<? super ClientResponse> actual,
-				Context ctx, @Nullable final TraceContext parent, Span clientSpan,
-				MonoWebClientTrace mono) {
+				Context ctx, Span clientSpan, MonoWebClientTrace mono) {
 			this.actual = actual;
-			this.parent = parent;
+			this.parent = mono.parent;
 			this.clientSpan = clientSpan;
 			this.handler = mono.handler;
 			this.currentTraceContext = mono.currentTraceContext;