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 052cf5467..a358784da 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 @@ -17,15 +17,16 @@ package org.springframework.cloud.sleuth.instrument.web.client; import java.net.InetSocketAddress; -import java.util.List; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiConsumer; import java.util.function.BiFunction; import brave.Span; -import brave.Tracer; 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; import reactor.netty.Connection; @@ -58,21 +59,27 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor { // 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)) + return ((HttpClient) bean) .doOnResponseError(new TracingDoOnErrorResponse(httpTracing)) + .doOnResponse(new TracingDoOnResponse(httpTracing)) + .doOnRequestError(new TracingDoOnErrorRequest(httpTracing)) + .doOnRequest(new TracingDoOnRequest(httpTracing)) .mapConnect(new TracingMapConnect(httpTracing)); } return bean; } + /** current client span, cleared on completion. */ + private static final class CurrentClientSpan extends AtomicReference { + + } + private static class TracingMapConnect implements BiFunction, Bootstrap, Mono> { final LazyBean httpTracing; - Tracer tracer; + CurrentTraceContext currentTraceContext; TracingMapConnect(LazyBean httpTracing) { this.httpTracing = httpTracing; @@ -81,16 +88,22 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor { @Override public Mono apply(Mono 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()))); + return mono.subscriberContext(context -> { + TraceContext invocationContext = currentTraceContext().get(); + if (invocationContext != null) { + // Read in this processor and also in ScopePassingSpanSubscriber + context = context.put(TraceContext.class, invocationContext); + } + return context.put(CurrentClientSpan.class, new CurrentClientSpan()); + }); } - Tracer tracer() { - if (this.tracer == null) { - this.tracer = this.httpTracing.get().tracing().tracer(); + CurrentTraceContext currentTraceContext() { + if (this.currentTraceContext == null) { + this.currentTraceContext = this.httpTracing.get().tracing() + .currentTraceContext(); } - return this.tracer; + return this.currentTraceContext; } } @@ -100,21 +113,12 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor { final LazyBean httpTracing; - List propagationKeys; - HttpClientHandler handler; TracingDoOnRequest(LazyBean httpTracing) { this.httpTracing = httpTracing; } - List propagationKeys() { - if (this.propagationKeys == null) { - this.propagationKeys = httpTracing.get().tracing().propagation().keys(); - } - return this.propagationKeys; - } - HttpClientHandler handler() { if (this.handler == null) { this.handler = HttpClientHandler.create(httpTracing.get()); @@ -122,30 +126,39 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor { return this.handler; } + CurrentTraceContext currentTraceContext() { + return httpTracing.get().tracing().currentTraceContext(); + } + @Override public void accept(HttpClientRequest req, Connection connection) { - // request already instrumented - // TODO: consider another, cheaper way, like flagging a context - // property. If not, comment why. - for (String key : propagationKeys()) { - if (req.requestHeaders().contains(key)) { - return; - } + CurrentClientSpan ref = req.currentContext() + .getOrDefault(CurrentClientSpan.class, null); + if (ref == null) { // Somehow TracingMapConnect was not invoked.. skip out + return; } - // Look for a parent propagated by TracingMapConnect - AtomicReference ref = req.currentContext() - .getOrDefault(AtomicReference.class, null); - Span parent = ref != null ? ref.get() : null; + // This might be re-entrant on auto-redirect or connection retry: + // See reactor/reactor-netty#1000 for follow-ups. + Span clientSpan = ref.getAndSet(null); + if (clientSpan != null) { + // Retry from a connect fail wouldn't have parsed the request, leading to + // an empty span with no data if we finished it. An auto-redirect would + // have parsed the request, but we have no idea which status code it + // finished with. Since we can't see the preceding request state, we + // abandon its span in favor of the next. + clientSpan.abandon(); + } // Start a new client span with the appropriate parent + TraceContext parent = req.currentContext().getOrDefault(TraceContext.class, + null); WrappedHttpClientRequest request = new WrappedHttpClientRequest(req); - Span clientSpan = parent != null ? handler().handleSend(request, parent) - : handler().handleSend(request); - parseConnectionAddress(connection, clientSpan); - // Swap the ref with the client span, so that other hooks can see it - if (ref != null) { + // Simplify after openzipkin/brave#1082 + try (Scope ws = currentTraceContext().maybeScope(parent)) { + clientSpan = handler().handleSend(request); + parseConnectionAddress(connection, clientSpan); ref.set(clientSpan); } } @@ -182,9 +195,8 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor { } @Override - public void accept(HttpClientRequest request, Throwable error) { - // TODO: the current context here does not have the AtomicReference - handle(request.currentContext(), null, error); + public void accept(HttpClientRequest req, Throwable error) { + handle(req.currentContext(), null, error); } } @@ -222,14 +234,18 @@ class HttpClientBeanPostProcessor implements BeanPostProcessor { void handle(Context context, @Nullable HttpClientResponse resp, @Nullable Throwable error) { - AtomicReference ref = context.getOrDefault(AtomicReference.class, null); - Span span = ref != null ? ref.get() : null; - if (span == null) { + CurrentClientSpan ref = context.getOrDefault(CurrentClientSpan.class, null); + if (ref == null) { // Somehow TracingMapConnect was not invoked.. skip out + return; + } + + Span clientSpan = ref.getAndSet(null); + if (clientSpan == null) { return; // Unexpected. In the handle method, without a span to finish! } WrappedHttpClientResponse response = resp != null ? new WrappedHttpClientResponse(resp) : null; - handler().handleReceive(response, error, span); + handler().handleReceive(response, error, clientSpan); } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/ReactorNettyHttpClientSpringBootTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/ReactorNettyHttpClientSpringBootTests.java index 791146de6..efecdd11b 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/ReactorNettyHttpClientSpringBootTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/client/ReactorNettyHttpClientSpringBootTests.java @@ -21,7 +21,10 @@ import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import brave.propagation.B3SinglePropagation; +import brave.propagation.CurrentTraceContext; +import brave.propagation.CurrentTraceContext.Scope; import brave.propagation.Propagation; +import brave.propagation.TraceContext; import brave.sampler.Sampler; import io.netty.handler.codec.http.HttpResponseStatus; import org.junit.After; @@ -70,6 +73,12 @@ public class ReactorNettyHttpClientSpringBootTests { @Autowired BlockingQueue spans; + @Autowired + CurrentTraceContext currentTraceContext; + + TraceContext context = TraceContext.newBuilder().traceId(1).spanId(1).sampled(true) + .build(); + @After public void tearDown() { if (disposableServer != null) { @@ -96,6 +105,26 @@ public class ReactorNettyHttpClientSpringBootTests { assertThat(clientSpan.remoteEndpoint().portAsInt()).isNotZero(); } + @Test + public void shouldUseInvocationContext() throws Exception { + disposableServer = HttpServer.create().port(0) + // this reads the trace context header, b3, returning it in the response + .handle((in, out) -> out + .sendString(Flux.just(in.requestHeaders().get("b3")))) + .bindNow(); + + String b3SingleHeaderReadByServer; + try (Scope ws = currentTraceContext.newScope(context)) { + b3SingleHeaderReadByServer = httpClient.port(disposableServer.port()).get() + .uri("/").responseContent().aggregate().asString().block(); + } + + Span clientSpan = takeClientSpan(); + + assertThat(b3SingleHeaderReadByServer).isEqualTo(context.traceIdString() + "-" + + clientSpan.id() + "-1-" + context.spanIdString()); + } + @Test public void shouldSendTraceContextToServer_rootSpan() throws Exception { disposableServer = HttpServer.create().port(0)