Enable tracing for Spring Webflux WebClient
ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java/datadog/trace/instrumentation/springwebflux/client/DefaultWebClientInstrumentation.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Map;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public class DefaultWebClientInstrumentation extends Instrumenter.Default {
+
+ public DefaultWebClientInstrumentation() {
+ super("spring-webflux");
+ }
+
+ @Override
+ public String[] helperClassNames() {
+ return new String[] {
+ "datadog.trace.agent.decorator.BaseDecorator",
+ "datadog.trace.agent.decorator.ClientDecorator",
+ "datadog.trace.agent.decorator.HttpClientDecorator",
+ packageName + ".SpringWebfluxHttpClientDecorator",
+ packageName + ".HttpHeadersInjectAdapter",
+ packageName + ".TracingClientResponseSubscriber",
+ packageName + ".TracingClientResponseSubscriber$1",
+ packageName + ".TracingClientResponseMono",
+ packageName + ".ClientResponseWrapper"
+ };
+ }
+
+ @Override
+ public ElementMatcher<? super TypeDescription> typeMatcher() {
+ return named(
+ "org.springframework.web.reactive.function.client.ExchangeFunctions$DefaultExchangeFunction");
+ }
+
+ @Override
+ public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+ return singletonMap(
+ isMethod()
+ .and(isPublic())
+ .and(named("exchange"))
+ .and(
+ takesArgument(
+ 0, named("org.springframework.web.reactive.function.client.ClientRequest"))),
+ packageName + ".DefaultWebClientAdvice");
+ }
+}

ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/ClientResponseWrapper.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import java.util.List;
+import org.springframework.core.ParameterizedTypeReference;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.ResponseCookie;
+import org.springframework.http.ResponseEntity;
+import org.springframework.http.client.reactive.ClientHttpResponse;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.reactive.function.BodyExtractor;
+import org.springframework.web.reactive.function.client.ClientResponse;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
+
+public class ClientResponseWrapper implements ClientResponse {
+
+ private final ClientResponse clientResponse;
+ private final Context context;
+
+ public ClientResponseWrapper(final ClientResponse clientResponse, final Context context) {
+ this.clientResponse = clientResponse;
+ this.context = context;
+ }
+
+ @Override
+ public HttpStatus statusCode() {
+ return clientResponse.statusCode();
+ }
+
+ @Override
+ public Headers headers() {
+ return clientResponse.headers();
+ }
+
+ @Override
+ public MultiValueMap<String, ResponseCookie> cookies() {
+ return clientResponse.cookies();
+ }
+
+ @Override
+ public <T> T body(final BodyExtractor<T, ? super ClientHttpResponse> extractor) {
+ return clientResponse.body(extractor);
+ }
+
+ @Override
+ public <T> Mono<T> bodyToMono(final Class<? extends T> elementClass) {
+ return clientResponse.<T>bodyToMono(elementClass).subscriberContext(context);
+ }
+
+ @Override
+ public <T> Mono<T> bodyToMono(final ParameterizedTypeReference<T> typeReference) {
+ return clientResponse.<T>bodyToMono(typeReference).subscriberContext(context);
+ }
+
+ @Override
+ public <T> Flux<T> bodyToFlux(final Class<? extends T> elementClass) {
+ return clientResponse.<T>bodyToFlux(elementClass).subscriberContext(context);
+ }
+
+ @Override
+ public <T> Flux<T> bodyToFlux(final ParameterizedTypeReference<T> typeReference) {
+ return clientResponse.<T>bodyToFlux(typeReference).subscriberContext(context);
+ }
+
+ @Override
+ public <T> Mono<ResponseEntity<T>> toEntity(final Class<T> bodyType) {
+ return clientResponse.toEntity(bodyType);
+ }
+
+ @Override
+ public <T> Mono<ResponseEntity<T>> toEntity(final ParameterizedTypeReference<T> typeReference) {
+ return clientResponse.toEntity(typeReference);
+ }
+
+ @Override
+ public <T> Mono<ResponseEntity<List<T>>> toEntityList(final Class<T> elementType) {
+ return clientResponse.toEntityList(elementType);
+ }
+
+ @Override
+ public <T> Mono<ResponseEntity<List<T>>> toEntityList(
+ final ParameterizedTypeReference<T> typeReference) {
+ return clientResponse.toEntityList(typeReference);
+ }
+
+ public int rawStatusCode() {
+ return clientResponse.statusCode().value();
+ }
+}

ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/DefaultWebClientAdvice.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import io.opentracing.util.GlobalTracer;
+import net.bytebuddy.asm.Advice;
+import org.springframework.web.reactive.function.client.ClientRequest;
+import org.springframework.web.reactive.function.client.ClientResponse;
+import org.springframework.web.reactive.function.client.ExchangeFunction;
+import reactor.core.publisher.Mono;
+
+public class DefaultWebClientAdvice {
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void methodExit(
+ @Advice.Thrown final Throwable throwable,
+ @Advice.This final ExchangeFunction exchangeFunction,
+ @Advice.Argument(0) final ClientRequest clientRequest,
+ @Advice.Return(readOnly = false) Mono<ClientResponse> mono) {
+ if (throwable == null
+ && mono != null
+ && !clientRequest.headers().keySet().contains("x-datadog-trace-id")) {
+ mono = new TracingClientResponseMono(clientRequest, exchangeFunction, GlobalTracer.get());
+ }
+ }
+}

ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/HttpHeadersInjectAdapter.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import io.opentracing.propagation.TextMap;
+import java.util.Iterator;
+import java.util.Map;
+import org.springframework.http.HttpHeaders;
+
+public class HttpHeadersInjectAdapter implements TextMap {
+
+ private final HttpHeaders headers;
+
+ public HttpHeadersInjectAdapter(final HttpHeaders headers) {
+ this.headers = headers;
+ }
+
+ @Override
+ public Iterator<Map.Entry<String, String>> iterator() {
+ throw new UnsupportedOperationException("This class should be used only with Tracer.inject()!");
+ }
+
+ @Override
+ public void put(final String key, final String value) {
+ headers.set(key, value);
+ }
+}

ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/SpringWebfluxHttpClientDecorator.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import datadog.trace.agent.decorator.HttpClientDecorator;
+import io.opentracing.Span;
+import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.web.reactive.function.client.ClientRequest;
+import org.springframework.web.reactive.function.client.ClientResponse;
+
+@Slf4j
+public class SpringWebfluxHttpClientDecorator
+ extends HttpClientDecorator<ClientRequest, ClientResponse> {
+ public static final SpringWebfluxHttpClientDecorator DECORATE =
+ new SpringWebfluxHttpClientDecorator();
+
+ public void onCancel(final Span span) {
+ final Map<String, Object> logs = new HashMap<>(2);
+ logs.put("event", "cancelled");
+ logs.put("message", "The subscription was cancelled");
+ span.log(logs);
+ }
+
+ @Override
+ protected String[] instrumentationNames() {
+ return new String[] {"spring-webflux-client"};
+ }
+
+ @Override
+ protected String component() {
+ return "spring-webflux-client";
+ }
+
+ @Override
+ protected String method(final ClientRequest httpRequest) {
+ return httpRequest.method().name();
+ }
+
+ @Override
+ protected URI url(final ClientRequest httpRequest) {
+ return httpRequest.url();
+ }
+
+ @Override
+ protected String hostname(final ClientRequest httpRequest) {
+ return httpRequest.url().getHost();
+ }
+
+ @Override
+ protected Integer port(final ClientRequest httpRequest) {
+ return httpRequest.url().getPort();
+ }
+
+ @Override
+ protected Integer status(final ClientResponse httpResponse) {
+ return httpResponse.statusCode().value();
+ }
+}

ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import io.opentracing.Tracer;
+import io.opentracing.propagation.Format;
+import io.opentracing.tag.Tags;
+import org.springframework.web.reactive.function.client.ClientRequest;
+import org.springframework.web.reactive.function.client.ClientResponse;
+import org.springframework.web.reactive.function.client.ExchangeFunction;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
+
+public class TracingClientResponseMono extends Mono<ClientResponse> {
+
+ private final ClientRequest clientRequest;
+ private final ExchangeFunction exchangeFunction;
+ private final Tracer tracer;
+
+ public TracingClientResponseMono(
+ final ClientRequest clientRequest,
+ final ExchangeFunction exchangeFunction,
+ final Tracer tracer) {
+ this.clientRequest = clientRequest;
+ this.exchangeFunction = exchangeFunction;
+ this.tracer = tracer;
+ }
+
+ @Override
+ public void subscribe(final CoreSubscriber<? super ClientResponse> subscriber) {
+ final Context context = subscriber.currentContext();
+ final Span parentSpan = context.<Span>getOrEmpty(Span.class).orElseGet(tracer::activeSpan);
+
+ final Span span =
+ tracer
+ .buildSpan(clientRequest.method().toString())
+ .asChildOf(parentSpan)
+ .withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
+ .start();
+
+ try (final Scope scope = tracer.scopeManager().activate(span, false)) {
+ final ClientRequest mutatedRequest =
+ ClientRequest.from(clientRequest)
+ .headers(
+ httpHeaders ->
+ tracer.inject(
+ span.context(),
+ Format.Builtin.HTTP_HEADERS,
+ new HttpHeadersInjectAdapter(httpHeaders)))
+ .build();
+ exchangeFunction
+ .exchange(mutatedRequest)
+ .subscribe(
+ new TracingClientResponseSubscriber(subscriber, mutatedRequest, context, span));
+ }
+ }
+}

ADD
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java
+package datadog.trace.instrumentation.springwebflux.client;
+
+import static datadog.trace.instrumentation.springwebflux.client.SpringWebfluxHttpClientDecorator.DECORATE;
+
+import io.opentracing.Span;
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import org.reactivestreams.Subscription;
+import org.springframework.web.reactive.function.client.ClientRequest;
+import org.springframework.web.reactive.function.client.ClientResponse;
+import reactor.core.CoreSubscriber;
+import reactor.util.context.Context;
+
+public class TracingClientResponseSubscriber implements CoreSubscriber<ClientResponse> {
+
+ private final CoreSubscriber<? super ClientResponse> subscriber;
+ private final ClientRequest clientRequest;
+ private final Context context;
+ private final Span span;
+
+ public TracingClientResponseSubscriber(
+ final CoreSubscriber<? super ClientResponse> subscriber,
+ final ClientRequest clientRequest,
+ final Context context,
+ final Span span) {
+ this.subscriber = subscriber;
+ this.clientRequest = clientRequest;
+ this.context = context.put(Span.class, span);
+ this.span = span;
+ }
+
+ @Override
+ public void onSubscribe(final Subscription subscription) {
+ DECORATE.onRequest(span, clientRequest);
+ DECORATE.onPeerConnection(span, remoteAddress());
+
+ subscriber.onSubscribe(
+ new Subscription() {
+ @Override
+ public void request(final long n) {
+ subscription.request(n);
+ }
+
+ @Override
+ public void cancel() {
+ DECORATE.onCancel(span);
+ subscription.cancel();
+ span.finish();
+ }
+ });
+ }
+
+ private InetAddress remoteAddress() {
+ try {
+ return InetAddress.getByName(clientRequest.url().getHost());
+ } catch (final UnknownHostException e) {
+ return null;
+ }
+ }
+
+ @Override
+ public void onNext(final ClientResponse clientResponse) {
+ try {
+ subscriber.onNext(new ClientResponseWrapper(clientResponse, context));
+ } finally {
+ DECORATE.onResponse(span, clientResponse);
+ }
+ }
+
+ @Override
+ public void onError(final Throwable throwable) {
+ try {
+ subscriber.onError(throwable);
+ } finally {
+ DECORATE.onError(span, throwable);
+ span.finish();
+ }
+ }
+
+ @Override
+ public void onComplete() {
+ try {
+ subscriber.onComplete();
+ } finally {
+ span.finish();
+ }
+ }
+
+ @Override
+ public Context currentContext() {
+ return context;
+ }
+}

MODIFY
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/server/AdviceUtils.java
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/server/AdviceUtils.java
@@ -5,6 +5,7 @@ import static datadog.trace.instrumentation.springwebflux.server.SpringWebfluxHt
import datadog.trace.instrumentation.reactor.core.ReactorCoreAdviceUtils;
import io.opentracing.Span;
import lombok.extern.slf4j.Slf4j;
+import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
@@ -39,4 +40,10 @@ public class AdviceUtils {
ReactorCoreAdviceUtils.finishSpanIfPresent(
(Span) serverRequest.attributes().remove(SPAN_ATTRIBUTE), throwable);
}
+
+ public static void finishSpanIfPresent(
+ final ClientRequest clientRequest, final Throwable throwable) {
+ ReactorCoreAdviceUtils.finishSpanIfPresent(
+ (Span) clientRequest.attributes().remove(SPAN_ATTRIBUTE), throwable);
+ }
}

Enable tracing for Spring Webflux WebClient