Moves client instrumentation off deprecated Brave code (#1532)

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

@@ -16,6 +16,7 @@
 
 package org.springframework.cloud.sleuth.instrument.web.client;
 
+import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
@@ -24,10 +25,7 @@ import brave.Span;
 import brave.Tracer;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
-import brave.propagation.Propagation;
-import brave.propagation.TraceContext;
 import io.netty.bootstrap.Bootstrap;
-import io.netty.handler.codec.http.HttpHeaders;
 import reactor.core.publisher.Mono;
 import reactor.netty.Connection;
 import reactor.netty.http.client.HttpClient;
@@ -89,31 +87,13 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 	private static class TracingDoOnRequest
 			implements BiConsumer<HttpClientRequest, Connection> {
 
-		static final Propagation.Setter<HttpHeaders, String> SETTER = new Propagation.Setter<HttpHeaders, String>() {
-			@Override
-			public void put(HttpHeaders carrier, String key, String value) {
-				if (!carrier.contains(key)) {
-					carrier.add(key, value);
-				}
-			}
-
-			@Override
-			public String toString() {
-				return "HttpHeaders::add";
-			}
-		};
-
 		final BeanFactory beanFactory;
 
 		HttpTracing httpTracing;
 
-		Tracer tracer;
+		List<String> propagationKeys;
 
-		HttpClientHandler<HttpClientRequest, HttpClientResponse> handler;
-
-		TraceContext.Injector<HttpHeaders> injector;
-
-		Propagation<String> propagation;
+		HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> handler;
 
 		TracingDoOnRequest(BeanFactory beanFactory) {
 			this.beanFactory = beanFactory;
@@ -130,23 +110,16 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			return this.httpTracing;
 		}
 
-		private Propagation<String> propagation() {
-			if (this.propagation == null) {
-				this.propagation = httpTracing().tracing().propagation();
+		private List<String> propagationKeys() {
+			if (this.propagationKeys == null) {
+				this.propagationKeys = httpTracing().tracing().propagation().keys();
 			}
-			return this.propagation;
+			return this.propagationKeys;
 		}
 
-		private TraceContext.Injector<HttpHeaders> injector() {
-			if (this.injector == null) {
-				this.injector = propagation().injector(SETTER);
-			}
-			return this.injector;
-		}
-
-		private HttpClientHandler<HttpClientRequest, HttpClientResponse> handler() {
+		private HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> handler() {
 			if (this.handler == null) {
-				this.handler = HttpClientHandler.create(httpTracing(), new HttpAdapter());
+				this.handler = HttpClientHandler.create(httpTracing());
 			}
 			return this.handler;
 		}
@@ -154,16 +127,18 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 		@Override
 		public void accept(HttpClientRequest req, Connection connection) {
 			// request already instrumented
-			for (String key : propagation().keys()) {
+			// TODO: consider another, cheaper way, like flagging a context
+			// property. If not, comment why.
+			for (String key : propagationKeys()) {
 				if (req.requestHeaders().contains(key)) {
 					return;
 				}
 			}
-			AtomicReference reference = req.currentContext()
-					.getOrDefault(AtomicReference.class, new AtomicReference());
-			Span span = handler().handleSend(injector(), req.requestHeaders(), req,
-					reference.get() == null ? handler().nextSpan(req)
-							: (Span) reference.get());
+			AtomicReference<Span> reference = req.currentContext()
+					.getOrDefault(AtomicReference.class, new AtomicReference<>());
+			WrappedHttpClientRequest request = new WrappedHttpClientRequest(req);
+			Span span = reference.get() == null ? handler().handleSend(request)
+					: handler().handleSend(request, reference.get());
 			reference.set(span);
 		}
 
@@ -229,7 +204,7 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 
 		HttpTracing httpTracing;
 
-		HttpClientHandler<HttpClientRequest, HttpClientResponse> handler;
+		HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> handler;
 
 		AbstractTracingDoOnHandler(BeanFactory beanFactory) {
 			this.beanFactory = beanFactory;
@@ -242,9 +217,9 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			return this.httpTracing;
 		}
 
-		private HttpClientHandler<HttpClientRequest, HttpClientResponse> handler() {
+		private HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> handler() {
 			if (this.handler == null) {
-				this.handler = HttpClientHandler.create(httpTracing(), new HttpAdapter());
+				this.handler = HttpClientHandler.create(httpTracing());
 			}
 			return this.handler;
 		}
@@ -259,34 +234,68 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor {
 			if (reference == null || reference.get() == null) {
 				return;
 			}
-			handler().handleReceive(httpClientResponse, throwable,
-					(Span) reference.get());
+			handler().handleReceive(new WrappedHttpClientResponse(httpClientResponse),
+					throwable, (Span) reference.get());
 		}
 
 	}
 
-	private static class HttpAdapter
-			extends brave.http.HttpClientAdapter<HttpClientRequest, HttpClientResponse> {
+	static final class WrappedHttpClientRequest extends brave.http.HttpClientRequest {
 
-		@Override
-		public String method(HttpClientRequest request) {
-			return request.method().name();
+		final HttpClientRequest delegate;
+
+		WrappedHttpClientRequest(HttpClientRequest delegate) {
+			this.delegate = delegate;
 		}
 
 		@Override
-		public String url(HttpClientRequest request) {
-			return request.uri();
+		public Object unwrap() {
+			return delegate;
 		}
 
 		@Override
-		public String requestHeader(HttpClientRequest request, String name) {
-			Object result = request.requestHeaders().get(name);
-			return result != null ? result.toString() : "";
+		public String method() {
+			return delegate.method().name();
 		}
 
 		@Override
-		public Integer statusCode(HttpClientResponse response) {
-			return response.status().code();
+		public String path() {
+			return delegate.path();
+		}
+
+		@Override
+		public String url() {
+			return delegate.uri();
+		}
+
+		@Override
+		public String header(String name) {
+			return delegate.requestHeaders().get(name);
+		}
+
+		@Override
+		public void header(String name, String value) {
+			delegate.header(name, value);
+		}
+
+	}
+
+	static final class WrappedHttpClientResponse extends brave.http.HttpClientResponse {
+
+		final HttpClientResponse delegate;
+
+		WrappedHttpClientResponse(HttpClientResponse delegate) {
+			this.delegate = delegate;
+		}
+
+		@Override
+		public Object unwrap() {
+			return delegate;
+		}
+
+		@Override
+		public int statusCode() {
+			return delegate.status().code();
 		}
 
 	}

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

@@ -23,8 +23,7 @@ import brave.Span;
 import brave.Tracer;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
-import brave.propagation.Propagation;
-import brave.propagation.TraceContext;
+import brave.propagation.TraceContext.Extractor;
 import brave.propagation.TraceContextOrSamplingFlags;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -33,7 +32,6 @@ import org.springframework.cloud.gateway.filter.headers.HttpHeadersFilter;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
-import org.springframework.lang.NonNull;
 import org.springframework.web.server.ServerWebExchange;
 
 final class TraceRequestHttpHeadersFilter extends AbstractHttpHeadersFilter {
@@ -54,18 +52,18 @@ final class TraceRequestHttpHeadersFilter extends AbstractHttpHeadersFilter {
 			log.debug("Will instrument the HTTP request headers ["
 					+ exchange.getRequest().getHeaders() + "]");
 		}
-		TraceCarrier carrier = new TraceCarrier(exchange.getRequest(), input);
-		Span currentSpan = currentSpan(carrier);
-		Span span = injectedSpan(carrier, currentSpan);
+		HttpClientRequest request = new HttpClientRequest(exchange.getRequest(), input);
+		Span currentSpan = currentSpan(request);
+		Span span = injectedSpan(request, currentSpan);
 		if (log.isDebugEnabled()) {
 			log.debug(
 					"Client span  " + span + " created for the request. New headers are "
-							+ carrier.filteredHeaders.toSingleValueMap());
+							+ request.filteredHeaders.toSingleValueMap());
 		}
 		exchange.getAttributes().put(SPAN_ATTRIBUTE, span);
 		HttpHeaders headersWithInput = new HttpHeaders();
 		headersWithInput.addAll(input);
-		addHeadersWithInput(carrier.filteredHeaders, headersWithInput);
+		addHeadersWithInput(request.filteredHeaders, headersWithInput);
 		if (headersWithInput.containsKey("b3") || headersWithInput.containsKey("B3")) {
 			headersWithInput.keySet().remove("b3");
 			headersWithInput.keySet().remove("B3");
@@ -73,22 +71,24 @@ final class TraceRequestHttpHeadersFilter extends AbstractHttpHeadersFilter {
 		return headersWithInput;
 	}
 
-	private Span currentSpan(TraceCarrier carrier) {
+	private Span currentSpan(HttpClientRequest request) {
 		Span currentSpan = this.tracer.currentSpan();
 		if (currentSpan != null) {
 			return currentSpan;
 		}
-		TraceContextOrSamplingFlags contextOrFlags = this.extractor.extract(carrier);
+		// Usually, an HTTP client would not attempt to resume a trace from headers, as a
+		// server would always place its span in scope. However, in commit 848442e,
+		// this behavior was added in support of gateway.
+		TraceContextOrSamplingFlags contextOrFlags = this.extractor.extract(request);
 		return this.tracer.nextSpan(contextOrFlags);
 	}
 
-	private Span injectedSpan(TraceCarrier carrier, Span currentSpan) {
+	private Span injectedSpan(HttpClientRequest request, Span currentSpan) {
 		if (currentSpan == null) {
-			return this.handler.handleSend(this.injector, carrier);
+			return this.handler.handleSend(request);
 		}
-		Span clientSpan = this.tracer
-				.nextSpan(TraceContextOrSamplingFlags.create(currentSpan.context()));
-		return this.handler.handleSend(this.injector, carrier, clientSpan);
+		Span clientSpan = this.tracer.newChild(currentSpan.context());
+		return this.handler.handleSend(request, clientSpan);
 	}
 
 	private void addHeadersWithInput(HttpHeaders filteredHeaders,
@@ -107,20 +107,6 @@ final class TraceRequestHttpHeadersFilter extends AbstractHttpHeadersFilter {
 
 }
 
-class TraceCarrier {
-
-	final ServerHttpRequest originalRequest;
-
-	final HttpHeaders filteredHeaders;
-
-	TraceCarrier(@NonNull ServerHttpRequest originalRequest,
-			@NonNull HttpHeaders filteredHeaders) {
-		this.originalRequest = originalRequest;
-		this.filteredHeaders = filteredHeaders;
-	}
-
-}
-
 final class TraceResponseHttpHeadersFilter extends AbstractHttpHeadersFilter {
 
 	private static final Log log = LogFactory
@@ -143,7 +129,8 @@ final class TraceResponseHttpHeadersFilter extends AbstractHttpHeadersFilter {
 		if (log.isDebugEnabled()) {
 			log.debug("Will instrument the response");
 		}
-		this.handler.handleReceive(exchange.getResponse(), null, (Span) storedSpan);
+		HttpClientResponse response = new HttpClientResponse(exchange.getResponse());
+		this.handler.handleReceive(response, null, (Span) storedSpan);
 		if (log.isDebugEnabled()) {
 			log.debug("The response was handled for span " + storedSpan);
 		}
@@ -161,71 +148,82 @@ abstract class AbstractHttpHeadersFilter implements HttpHeadersFilter {
 
 	static final String SPAN_ATTRIBUTE = Span.class.getName();
 
-	private static final Propagation.Setter<TraceCarrier, String> SETTER = new Propagation.Setter<TraceCarrier, String>() {
-		@Override
-		public void put(TraceCarrier carrier, String key, String value) {
-			carrier.filteredHeaders.set(key, value);
-		}
-
-		@Override
-		public String toString() {
-			return "TraceCarrier::httpHeaders::set";
-		}
-	};
-
-	private static final Propagation.Getter<TraceCarrier, String> GETTER = new Propagation.Getter<TraceCarrier, String>() {
-		@Override
-		public String get(TraceCarrier carrier, String key) {
-			return carrier.filteredHeaders.getFirst(key);
-		}
-
-		@Override
-		public String toString() {
-			return "TraceCarrier::httpHeaders::getFirst";
-		}
-	};
-
 	final Tracer tracer;
 
-	final HttpClientHandler<TraceCarrier, ServerHttpResponse> handler;
-
-	final TraceContext.Injector<TraceCarrier> injector;
-
-	final TraceContext.Extractor<TraceCarrier> extractor;
+	final HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> handler;
 
 	final HttpTracing httpTracing;
 
+	final Extractor<HttpClientRequest> extractor;
+
 	AbstractHttpHeadersFilter(HttpTracing httpTracing) {
 		this.tracer = httpTracing.tracing().tracer();
-		this.handler = HttpClientHandler.create(httpTracing, new ServerHttpAdapter());
-		this.injector = httpTracing.tracing().propagation().injector(SETTER);
-		this.extractor = httpTracing.tracing().propagation().extractor(GETTER);
+		this.extractor = httpTracing.tracing().propagation()
+				.extractor(HttpClientRequest::header);
+		this.handler = HttpClientHandler.create(httpTracing);
 		this.httpTracing = httpTracing;
 	}
 
-	private static class ServerHttpAdapter
-			extends brave.http.HttpClientAdapter<TraceCarrier, ServerHttpResponse> {
+	static final class HttpClientRequest extends brave.http.HttpClientRequest {
 
-		@Override
-		public String method(TraceCarrier request) {
-			return request.originalRequest.getMethodValue();
+		final ServerHttpRequest delegate;
+
+		final HttpHeaders filteredHeaders;
+
+		HttpClientRequest(ServerHttpRequest delegate, HttpHeaders filteredHeaders) {
+			this.delegate = delegate;
+			this.filteredHeaders = filteredHeaders;
 		}
 
 		@Override
-		public String url(TraceCarrier request) {
-			return request.originalRequest.getURI().toString();
+		public Object unwrap() {
+			return delegate;
 		}
 
 		@Override
-		public String requestHeader(TraceCarrier request, String name) {
-			Object result = request.filteredHeaders.get(name);
-			return result != null ? result.toString() : "";
+		public String method() {
+			return delegate.getMethodValue();
 		}
 
 		@Override
-		public Integer statusCode(ServerHttpResponse response) {
-			return response.getStatusCode() != null ? response.getStatusCode().value()
-					: null;
+		public String path() {
+			return delegate.getURI().getPath();
+		}
+
+		@Override
+		public String url() {
+			return delegate.getURI().toString();
+		}
+
+		@Override
+		public String header(String name) {
+			return filteredHeaders.getFirst(name);
+		}
+
+		@Override
+		public void header(String name, String value) {
+			filteredHeaders.set(name, value);
+		}
+
+	}
+
+	static final class HttpClientResponse extends brave.http.HttpClientResponse {
+
+		final ServerHttpResponse delegate;
+
+		HttpClientResponse(ServerHttpResponse delegate) {
+			this.delegate = delegate;
+		}
+
+		@Override
+		public Object unwrap() {
+			return delegate;
+		}
+
+		@Override
+		public int statusCode() {
+			return delegate.getStatusCode() != null ? delegate.getStatusCode().value()
+					: 0;
 		}
 
 	}

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

@@ -17,7 +17,6 @@
 package org.springframework.cloud.sleuth.instrument.web.client.feign;
 
 import java.io.IOException;
-import java.util.HashMap;
 
 import brave.Span;
 import brave.Tracer;
@@ -89,9 +88,8 @@ public class TraceFeignBlockingLoadBalancerClient
 					LOG.debug(
 							"General exception was thrown, so most likely the traced client wasn't called. Falling back to a manual span");
 				}
-				fallbackSpan = tracingFeignClient().handleSend(
-						new HashMap<>(request.headers()), request, fallbackSpan);
-				tracingFeignClient().handleReceive(fallbackSpan, response, e);
+				tracingFeignClient().handleSendAndReceive(fallbackSpan, request, response,
+						e);
 			}
 			throw e;
 		}

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

@@ -17,7 +17,6 @@
 package org.springframework.cloud.sleuth.instrument.web.client.feign;
 
 import java.io.IOException;
-import java.util.HashMap;
 
 import brave.Span;
 import brave.Tracer;
@@ -86,9 +85,8 @@ public class TraceLoadBalancerFeignClient extends LoadBalancerFeignClient {
 					log.debug(
 							"General exception was thrown, so most likely the traced client wasn't called. Falling back to a manual span");
 				}
-				fallbackSpan = tracingFeignClient().handleSend(
-						new HashMap<>(request.headers()), request, fallbackSpan);
-				tracingFeignClient().handleReceive(fallbackSpan, response, e);
+				tracingFeignClient().handleSendAndReceive(fallbackSpan, request, response,
+						e);
 			}
 			throw e;
 		}

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

@@ -17,6 +17,7 @@
 package org.springframework.cloud.sleuth.instrument.web.client.feign;
 
 import java.io.IOException;
+import java.net.URI;
 import java.nio.charset.Charset;
 import java.util.Collection;
 import java.util.Collections;
@@ -28,7 +29,6 @@ import brave.Tracer;
 import brave.http.HttpClientHandler;
 import brave.http.HttpTracing;
 import brave.propagation.Propagation;
-import brave.propagation.TraceContext;
 import feign.Client;
 import feign.Request;
 import feign.Response;
@@ -72,14 +72,11 @@ final class TracingFeignClient implements Client {
 
 	final Client delegate;
 
-	final HttpClientHandler<Request, Response> handler;
-
-	final TraceContext.Injector<Map<String, Collection<String>>> injector;
+	final HttpClientHandler<brave.http.HttpClientRequest, brave.http.HttpClientResponse> handler;
 
 	TracingFeignClient(HttpTracing httpTracing, Client delegate) {
 		this.tracer = httpTracing.tracing().tracer();
-		this.handler = HttpClientHandler.create(httpTracing, new HttpAdapter());
-		this.injector = httpTracing.tracing().propagation().injector(SETTER);
+		this.handler = HttpClientHandler.create(httpTracing);
 		this.delegate = delegate;
 	}
 
@@ -88,74 +85,131 @@ final class TracingFeignClient implements Client {
 	}
 
 	@Override
-	public Response execute(Request request, Request.Options options) throws IOException {
-		Map<String, Collection<String>> headers = new LinkedHashMap<>(request.headers());
-		Span span = handleSend(headers, request, null);
+	public Response execute(Request req, Request.Options options) throws IOException {
+		HttpClientRequest request = new HttpClientRequest(req);
+		Span span = this.handler.handleSend(request);
 		if (log.isDebugEnabled()) {
 			log.debug("Handled send of " + span);
 		}
-		Response response = null;
+		HttpClientResponse response = null;
 		Throwable error = null;
 		try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
-			response = this.delegate.execute(modifiedRequest(request, headers), options);
-			return response;
+			Response res = this.delegate.execute(request.build(), options);
+			if (res != null) { // possibly null on bad implementation or mocks
+				response = new HttpClientResponse(res);
+			}
+			return res;
 		}
 		catch (IOException | RuntimeException | Error e) {
 			error = e;
 			throw e;
 		}
 		finally {
-			handleReceive(span, response, error);
+			this.handler.handleReceive(response, error, span);
+
 			if (log.isDebugEnabled()) {
 				log.debug("Handled receive of " + span);
 			}
 		}
 	}
 
-	Span handleSend(Map<String, Collection<String>> headers, Request request,
-			Span clientSpan) {
-		if (clientSpan != null) {
-			return this.handler.handleSend(this.injector, headers, request, clientSpan);
-		}
-		return this.handler.handleSend(this.injector, headers, request);
+	void handleSendAndReceive(Span span, Request request, Response response,
+			Throwable error) {
+		this.handler.handleSend(new HttpClientRequest(request), span);
+		this.handler.handleReceive(
+				response != null ? new HttpClientResponse(response) : null, error, span);
 	}
 
-	void handleReceive(Span span, Response response, Throwable error) {
-		this.handler.handleReceive(response, error, span);
-	}
+	static final class HttpClientRequest extends brave.http.HttpClientRequest {
 
-	private Request modifiedRequest(Request request,
-			Map<String, Collection<String>> headers) {
-		String method = request.method();
-		String url = request.url();
-		byte[] body = request.body();
-		Charset charset = request.charset();
-		return Request.create(method, url, headers, body, charset);
-	}
+		final Request delegate;
 
-	static final class HttpAdapter
-			extends brave.http.HttpClientAdapter<Request, Response> {
+		Map<String, Collection<String>> headers;
 
-		@Override
-		public String method(Request request) {
-			return request.method();
+		HttpClientRequest(Request delegate) {
+			this.delegate = delegate;
 		}
 
 		@Override
-		public String url(Request request) {
-			return request.url();
+		public Object unwrap() {
+			return delegate;
 		}
 
 		@Override
-		public String requestHeader(Request request, String name) {
-			Collection<String> result = request.headers().get(name);
+		public String method() {
+			return delegate.method();
+		}
+
+		@Override
+		public String path() {
+			String url = url();
+			if (url == null) {
+				return null;
+			}
+			return URI.create(url).getPath();
+		}
+
+		@Override
+		public String url() {
+			return delegate.url();
+		}
+
+		@Override
+		public String header(String name) {
+			Collection<String> result = delegate.headers().get(name);
 			return result != null && result.iterator().hasNext()
 					? result.iterator().next() : null;
 		}
 
 		@Override
-		public Integer statusCode(Response response) {
-			return response.status();
+		public void header(String name, String value) {
+			if (headers == null) {
+				headers = new LinkedHashMap<>(delegate.headers());
+			}
+			if (!headers.containsKey(name)) {
+				headers.put(name, Collections.singletonList(value));
+				if (log.isTraceEnabled()) {
+					log.trace(
+							"Added key [" + name + "] and header value [" + value + "]");
+				}
+			}
+			else {
+				// TODO: this is incorrect to ignore as opposed to overwrite!
+				if (log.isTraceEnabled()) {
+					log.trace("Key [" + name + "] already there in the headers");
+				}
+			}
+		}
+
+		Request build() {
+			if (headers == null) {
+				return delegate;
+			}
+			String method = delegate.method();
+			String url = delegate.url();
+			byte[] body = delegate.body();
+			Charset charset = delegate.charset();
+			return Request.create(method, url, headers, body, charset);
+		}
+
+	}
+
+	static final class HttpClientResponse extends brave.http.HttpClientResponse {
+
+		final Response delegate;
+
+		HttpClientResponse(Response delegate) {
+			this.delegate = delegate;
+		}
+
+		@Override
+		public Object unwrap() {
+			return delegate;
+		}
+
+		@Override
+		public int statusCode() {
+			return delegate.status();
 		}
 
 	}