Reverts impl to the refactored one

    Revert "Not using httpStatus() method for custom status codes"

    This reverts commit b59277f

    Revert "Came back to previous impl for WebClient instrumentation; fixes gh-1442"

    This reverts commit bee61a9.

Author: marcingrzejszczak

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

@@ -19,20 +19,29 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 import brave.Span;
 import brave.Tracer;
+import brave.Tracing;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
 import brave.propagation.Propagation;
 import brave.propagation.TraceContext;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Mono;
+import reactor.util.annotation.Nullable;
+import reactor.util.context.Context;
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.cloud.sleuth.instrument.reactor.ReactorSleuth;
+import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.web.client.RestClientException;
 import org.springframework.web.reactive.function.client.ClientRequest;
 import org.springframework.web.reactive.function.client.ClientResponse;
@@ -90,9 +99,6 @@ private Consumer<List<ExchangeFilterFunction>> addTraceExchangeFilterFunctionIfN
 final class TraceExchangeFilterFunction implements ExchangeFilterFunction {
 
 	private static final Log log = LogFactory.getLog(TraceExchangeFilterFunction.class);
-
-	private static final String CLIENT_SPAN_KEY = "sleuth.webclient.clientSpan";
-
 	static final Propagation.Setter<ClientRequest.Builder, String> SETTER = new Propagation.Setter<ClientRequest.Builder, String>() {
 		@Override
 		public void put(ClientRequest.Builder carrier, String key, String value) {
@@ -111,12 +117,14 @@ public String toString() {
 		}
 	};
 
-	public static ExchangeFilterFunction create(BeanFactory beanFactory) {
-		return new TraceExchangeFilterFunction(beanFactory);
-	}
+	private static final String CLIENT_SPAN_KEY = "sleuth.webclient.clientSpan";
+
+	private static final String CANCELLED_SUBSCRIPTION_ERROR = "CANCELLED";
 
 	final BeanFactory beanFactory;
 
+	final Function<? super Publisher<DataBuffer>, ? extends Publisher<DataBuffer>> scopePassingTransformer;
+
 	Tracer tracer;
 
 	HttpTracing httpTracing;
@@ -127,86 +135,27 @@ public static ExchangeFilterFunction create(BeanFactory beanFactory) {
 
 	TraceExchangeFilterFunction(BeanFactory beanFactory) {
 		this.beanFactory = beanFactory;
+		this.scopePassingTransformer = ReactorSleuth
+				.scopePassingSpanOperator(beanFactory);
+	}
+
+	public static ExchangeFilterFunction create(BeanFactory beanFactory) {
+		return new TraceExchangeFilterFunction(beanFactory);
 	}
 
 	@Override
 	public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
-		final ClientRequest.Builder builder = ClientRequest.from(request);
-		Mono<ClientResponse> exchange = Mono.defer(() -> next.exchange(builder.build()))
-				.cast(Object.class).onErrorResume(Mono::just)
-				.zipWith(Mono.subscriberContext()).flatMap(anyAndContext -> {
-					if (log.isDebugEnabled()) {
-						log.debug("Wrapping the context [" + anyAndContext + "]");
-					}
-					Object any = anyAndContext.getT1();
-					Span clientSpan = anyAndContext.getT2().get(CLIENT_SPAN_KEY);
-					Mono<ClientResponse> continuation;
-					final Tracer.SpanInScope ws = tracer().withSpanInScope(clientSpan);
-					if (any instanceof Throwable) {
-						continuation = Mono.error((Throwable) any);
-					}
-					else {
-						continuation = Mono.just((ClientResponse) any);
-					}
-					return continuation
-							.doAfterSuccessOrError((clientResponse, throwable1) -> {
-								Throwable throwable = throwable1;
-								if (clientResponse == null
-										|| clientResponse.statusCode() == null) {
-									if (log.isDebugEnabled()) {
-										log.debug(
-												"No response was returned. Will close the span ["
-														+ clientSpan + "]");
-									}
-									handleReceive(clientSpan, ws, clientResponse,
-											throwable);
-									return;
-								}
-								boolean error = clientResponse.statusCode()
-										.is4xxClientError()
-										|| clientResponse.statusCode().is5xxServerError();
-								if (error) {
-									if (log.isDebugEnabled()) {
-										log.debug(
-												"Non positive status code was returned from the call. Will close the span ["
-														+ clientSpan + "]");
-									}
-									throwable = new RestClientException(
-											"Status code of the response is ["
-													+ clientResponse.statusCode().value()
-													+ "] and the reason is ["
-													+ clientResponse.statusCode()
-															.getReasonPhrase()
-													+ "]");
-								}
-								handleReceive(clientSpan, ws, clientResponse, throwable);
-							});
-				}).subscriberContext(c -> {
-					if (log.isDebugEnabled()) {
-						log.debug("Instrumenting WebClient call");
-					}
-					Span parent = c.getOrDefault(Span.class, null);
-					Span clientSpan = handler().handleSend(injector(), builder, request,
-							tracer().nextSpan());
-					if (log.isDebugEnabled()) {
-						log.debug("Handled send of " + clientSpan);
-					}
-					if (parent == null) {
-						c = c.put(Span.class, clientSpan);
-						if (log.isDebugEnabled()) {
-							log.debug("Reactor Context got injected with the client span "
-									+ clientSpan);
-						}
-					}
-					return c.put(CLIENT_SPAN_KEY, clientSpan);
-				});
-		return exchange;
-	}
+		ClientRequest.Builder builder = ClientRequest.from(request);
+		if (log.isDebugEnabled()) {
+			log.debug("Instrumenting WebClient call");
+		}
+		Span span = handler().handleSend(injector(), builder, request,
+				tracer().nextSpan());
+		if (log.isDebugEnabled()) {
+			log.debug("Handled send of " + span);
+		}
 
-	private void handleReceive(Span clientSpan, Tracer.SpanInScope ws,
-			ClientResponse clientResponse, Throwable throwable) {
-		handler().handleReceive(clientResponse, throwable, clientSpan);
-		ws.close();
+		return new MonoWebClientTrace(next, builder.build(), this, span);
 	}
 
 	@SuppressWarnings("unchecked")
@@ -241,6 +190,190 @@ TraceContext.Injector<ClientRequest.Builder> injector() {
 		return this.injector;
 	}
 
+	private static final class MonoWebClientTrace extends Mono<ClientResponse> {
+
+		final ExchangeFunction next;
+
+		final ClientRequest request;
+
+		final Tracer tracer;
+
+		final HttpClientHandler<ClientRequest, ClientResponse> handler;
+
+		final TraceContext.Injector<ClientRequest.Builder> injector;
+
+		final Tracing tracing;
+
+		final Function<? super Publisher<DataBuffer>, ? extends Publisher<DataBuffer>> scopePassingTransformer;
+
+		private final Span span;
+
+		MonoWebClientTrace(ExchangeFunction next, ClientRequest request,
+				TraceExchangeFilterFunction parent, Span span) {
+			this.next = next;
+			this.request = request;
+			this.tracer = parent.tracer();
+			this.handler = parent.handler();
+			this.injector = parent.injector();
+			this.tracing = parent.httpTracing().tracing();
+			this.scopePassingTransformer = parent.scopePassingTransformer;
+			this.span = span;
+		}
+
+		@Override
+		public void subscribe(CoreSubscriber<? super ClientResponse> subscriber) {
+
+			Context context = subscriber.currentContext();
+
+			this.next.exchange(request).subscribe(
+					new WebClientTracerSubscriber(subscriber, context, span, this));
+		}
+
+		static final class WebClientTracerSubscriber
+				implements CoreSubscriber<ClientResponse> {
+
+			final CoreSubscriber<? super ClientResponse> actual;
+
+			final Context context;
+
+			final Span span;
+
+			final Tracer.SpanInScope ws;
+
+			final HttpClientHandler<ClientRequest, ClientResponse> handler;
+
+			final Function<? super Publisher<DataBuffer>, ? extends Publisher<DataBuffer>> scopePassingTransformer;
+
+			final Tracing tracing;
+
+			boolean done;
+
+			WebClientTracerSubscriber(CoreSubscriber<? super ClientResponse> actual,
+					Context context, Span span, MonoWebClientTrace parent) {
+				this.actual = actual;
+				this.span = span;
+				this.handler = parent.handler;
+				this.tracing = parent.tracing;
+				this.scopePassingTransformer = parent.scopePassingTransformer;
+
+				if (!context.hasKey(Span.class)) {
+					context = context.put(Span.class, span);
+					if (log.isDebugEnabled()) {
+						log.debug("Reactor Context got injected with the client span "
+								+ span);
+					}
+				}
+
+				this.context = context.put(CLIENT_SPAN_KEY, span);
+				this.ws = parent.tracer.withSpanInScope(span);
+
+			}
+
+			@Override
+			public void onSubscribe(Subscription subscription) {
+				this.actual.onSubscribe(new Subscription() {
+					@Override
+					public void request(long n) {
+						subscription.request(n);
+					}
+
+					@Override
+					public void cancel() {
+						terminateSpanOnCancel();
+						subscription.cancel();
+					}
+				});
+			}
+
+			@Override
+			public void onNext(ClientResponse response) {
+				this.done = true;
+				try {
+					// decorate response body
+					this.actual
+							.onNext(ClientResponse.from(response)
+									.body(response.bodyToFlux(DataBuffer.class)
+											.transform(this.scopePassingTransformer))
+									.build());
+				}
+				finally {
+					terminateSpan(response, null);
+				}
+			}
+
+			@Override
+			public void onError(Throwable t) {
+				try {
+					this.actual.onError(t);
+				}
+				finally {
+					terminateSpan(null, t);
+				}
+			}
+
+			@Override
+			public void onComplete() {
+				try {
+					this.actual.onComplete();
+				}
+				finally {
+					if (!this.done) {
+						terminateSpan(null, null);
+					}
+				}
+			}
+
+			@Override
+			public Context currentContext() {
+				return this.context;
+			}
+
+			void handleReceive(Span clientSpan, Tracer.SpanInScope ws,
+					ClientResponse clientResponse, Throwable throwable) {
+				this.handler.handleReceive(clientResponse, throwable, clientSpan);
+				ws.close();
+			}
+
+			void terminateSpanOnCancel() {
+				if (log.isDebugEnabled()) {
+					log.debug("Subscription was cancelled. Will close the span ["
+							+ this.span + "]");
+				}
+
+				this.span.tag("error", CANCELLED_SUBSCRIPTION_ERROR);
+				handleReceive(this.span, this.ws, null, null);
+			}
+
+			void terminateSpan(@Nullable ClientResponse clientResponse,
+					@Nullable Throwable throwable) {
+				if (clientResponse == null || clientResponse.statusCode() == null) {
+					if (log.isDebugEnabled()) {
+						log.debug("No response was returned. Will close the span ["
+								+ this.span + "]");
+					}
+					handleReceive(this.span, this.ws, clientResponse, throwable);
+					return;
+				}
+				boolean error = clientResponse.statusCode().is4xxClientError()
+						|| clientResponse.statusCode().is5xxServerError();
+				if (error) {
+					if (log.isDebugEnabled()) {
+						log.debug(
+								"Non positive status code was returned from the call. Will close the span ["
+										+ this.span + "]");
+					}
+					throwable = new RestClientException("Status code of the response is ["
+							+ clientResponse.statusCode().value()
+							+ "] and the reason is ["
+							+ clientResponse.statusCode().getReasonPhrase() + "]");
+				}
+				handleReceive(this.span, this.ws, clientResponse, throwable);
+			}
+
+		}
+
+	}
+
 	static final class HttpAdapter
 			extends brave.http.HttpClientAdapter<ClientRequest, ClientResponse> {