Fix trace propagation
MODIFY
dd-java-agent/instrumentation/spring-webflux/src/main/java/datadog/trace/instrumentation/springwebflux/client/DefaultWebClientInstrumentation.java
dd-java-agent/instrumentation/spring-webflux/src/main/java/datadog/trace/instrumentation/springwebflux/client/DefaultWebClientInstrumentation.java
@@ -32,7 +32,6 @@ public class DefaultWebClientInstrumentation extends Instrumenter.Default {
packageName + ".TracingClientResponseSubscriber",
packageName + ".TracingClientResponseSubscriber$1",
packageName + ".TracingClientResponseMono",
- packageName + ".ClientResponseWrapper"
};
}

DELETE
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;
-
-/**
- * Wrapper class for ClientResponse that adds Context to the body Publisher
- */
-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);
- }
-
- /**
- * ClientResponseWrapper is based on the ClientResponse from
- * spring-webflux-5.0.0.RELEASE. Since spring-webflux 5.1 ClientResponse
- * contains extra methods like rawStatusCode and gives methodNotFound
- * exceptions at runtime if used in a project with the latest spring-webflux
- * 5.1 or higher.
- * <p>
- * See https://docs.spring.io/spring/docs/5.1.x/javadoc-api/org/springframework/web/reactive/function/client/ClientResponse.html#rawStatusCode--
- */
- public int rawStatusCode() {
- return clientResponse.statusCode().value();
- }
-}

MODIFY
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseMono.java
@@ -2,6 +2,7 @@ package datadog.trace.instrumentation.springwebflux.client;
import static datadog.trace.instrumentation.springwebflux.client.SpringWebfluxHttpClientDecorator.DECORATE;
+import datadog.trace.context.TraceScope;
import io.opentracing.Scope;
import io.opentracing.Span;
import io.opentracing.Tracer;
@@ -36,13 +37,18 @@ public class TracingClientResponseMono extends Mono<ClientResponse> {
final Span span =
tracer
- .buildSpan("webflux.request")
+ .buildSpan("http.request")
.asChildOf(parentSpan)
.withTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT)
.start();
DECORATE.afterStart(span);
try (final Scope scope = tracer.scopeManager().activate(span, false)) {
+
+ if (scope instanceof TraceScope) {
+ ((TraceScope) scope).setAsyncPropagation(true);
+ }
+
final ClientRequest mutatedRequest =
ClientRequest.from(clientRequest)
.headers(
@@ -55,7 +61,7 @@ public class TracingClientResponseMono extends Mono<ClientResponse> {
exchangeFunction
.exchange(mutatedRequest)
.subscribe(
- new TracingClientResponseSubscriber(subscriber, mutatedRequest, context, span));
+ new TracingClientResponseSubscriber(subscriber, mutatedRequest, context, span, parentSpan));
}
}
}

MODIFY
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java
dd-java-agent/instrumentation/spring-webflux/src/main/java8/datadog/trace/instrumentation/springwebflux/client/TracingClientResponseSubscriber.java
@@ -2,7 +2,13 @@ package datadog.trace.instrumentation.springwebflux.client;
import static datadog.trace.instrumentation.springwebflux.client.SpringWebfluxHttpClientDecorator.DECORATE;
+import datadog.trace.context.TraceScope;
+import io.opentracing.Scope;
import io.opentracing.Span;
+import io.opentracing.Tracer;
+import io.opentracing.noop.NoopSpan;
+import io.opentracing.util.GlobalTracer;
+import java.util.concurrent.atomic.AtomicReference;
import org.reactivestreams.Subscription;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.ClientResponse;
@@ -11,32 +17,50 @@ import reactor.util.context.Context;
public class TracingClientResponseSubscriber implements CoreSubscriber<ClientResponse> {
+ private final Tracer tracer = GlobalTracer.get();
private final CoreSubscriber<? super ClientResponse> subscriber;
private final ClientRequest clientRequest;
private final Context context;
- private final Span span;
+ private final AtomicReference<Span> spanRef;
+ private final Span parentSpan;
public TracingClientResponseSubscriber(
final CoreSubscriber<? super ClientResponse> subscriber,
final ClientRequest clientRequest,
final Context context,
- final Span span) {
+ final Span span,
+ final Span parentSpan) {
this.subscriber = subscriber;
this.clientRequest = clientRequest;
- this.context = context.put(Span.class, span);
- this.span = span;
+ this.context = context;
+ spanRef = new AtomicReference<>(span);
+ this.parentSpan = parentSpan == null ? NoopSpan.INSTANCE : parentSpan;
}
@Override
public void onSubscribe(final Subscription subscription) {
+ final Span span = spanRef.get();
+ if (span == null) {
+ subscriber.onSubscribe(subscription);
+ return;
+ }
+
+ try (final Scope scope = tracer.scopeManager().activate(span, false)) {
+
+ if (scope instanceof TraceScope) {
+ ((TraceScope) scope).setAsyncPropagation(true);
+ }
+
DECORATE.onRequest(span, clientRequest);
subscriber.onSubscribe(
new Subscription() {
@Override
public void request(final long n) {
+ try (final Scope scope = tracer.scopeManager().activate(span, false)) {
subscription.request(n);
}
+ }
@Override
public void cancel() {
@@ -47,35 +71,62 @@ public class TracingClientResponseSubscriber implements CoreSubscriber<ClientRes
}
});
}
+ }
@Override
public void onNext(final ClientResponse clientResponse) {
- try {
- subscriber.onNext(new ClientResponseWrapper(clientResponse, context));
- } finally {
+ final Span span = spanRef.getAndSet(null);
+ if (span != null) {
DECORATE.onResponse(span, clientResponse);
+ DECORATE.beforeFinish(span);
+ span.finish();
+ }
+
+ try (final Scope scope = tracer.scopeManager().activate(parentSpan, false)) {
+
+ if (scope instanceof TraceScope) {
+ ((TraceScope) scope).setAsyncPropagation(true);
+ }
+
+ subscriber.onNext(clientResponse);
}
}
@Override
public void onError(final Throwable throwable) {
- try {
- subscriber.onError(throwable);
- } finally {
+ final Span span = spanRef.getAndSet(null);
+ if (span != null) {
DECORATE.onError(span, throwable);
DECORATE.beforeFinish(span);
span.finish();
}
+
+ try (final Scope scope = tracer.scopeManager().activate(parentSpan, false)) {
+
+ if (scope instanceof TraceScope) {
+ ((TraceScope) scope).setAsyncPropagation(true);
+ }
+
+ subscriber.onError(throwable);
+ }
}
@Override
public void onComplete() {
- try {
- subscriber.onComplete();
- } finally {
+ final Span span = spanRef.getAndSet(null);
+ if (span != null) {
DECORATE.beforeFinish(span);
span.finish();
}
+
+ try (final Scope scope = tracer.scopeManager().activate(parentSpan, false)) {
+
+ if (scope instanceof TraceScope) {
+ ((TraceScope) scope).setAsyncPropagation(true);
+ }
+
+ subscriber.onComplete();
+ }
}
@Override

MODIFY
dd-java-agent/instrumentation/spring-webflux/src/test/groovy/dd/trace/instrumentation/springwebflux/client/SpringWebfluxHttpClientTest.groovy
dd-java-agent/instrumentation/spring-webflux/src/test/groovy/dd/trace/instrumentation/springwebflux/client/SpringWebfluxHttpClientTest.groovy
package dd.trace.instrumentation.springwebflux.client
+import datadog.trace.agent.test.asserts.TraceAssert
import datadog.trace.agent.test.base.HttpClientTest
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.DDTags
+import datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator
import datadog.trace.instrumentation.springwebflux.client.SpringWebfluxHttpClientDecorator
+import io.opentracing.tag.Tags
+import io.opentracing.util.GlobalTracer
+import org.springframework.http.HttpMethod
import org.springframework.web.reactive.function.client.ClientResponse
import org.springframework.web.reactive.function.client.WebClient
import spock.lang.Shared
@@ -13,14 +20,20 @@ class SpringWebfluxHttpClientTest extends HttpClientTest<SpringWebfluxHttpClient
@Override
int doRequest(String method, URI uri, Map<String, String> headers, Closure callback) {
- assert method == "GET"
- ClientResponse response = client.get()
- .headers({ h -> headers.forEach({ key, value -> h.add(key, value) }) })
+ def hasParent = GlobalTracer.get().activeSpan() != null
+ ClientResponse response = client.method(HttpMethod.resolve(method))
+ .headers { h -> headers.forEach({ key, value -> h.add(key, value) }) }
.uri(uri)
.exchange()
+ .doOnSuccessOrError { success, error ->
+ blockUntilChildSpansFinished(1)
+ callback?.call()
+ }
.block()
- callback?.call()
+ if(hasParent) {
+ blockUntilChildSpansFinished(callback ? 3 : 2)
+ }
response.statusCode().value()
}
@@ -28,4 +41,54 @@ class SpringWebfluxHttpClientTest extends HttpClientTest<SpringWebfluxHttpClient
SpringWebfluxHttpClientDecorator decorator() {
return SpringWebfluxHttpClientDecorator.DECORATE
}
+
+
+ @Override
+ // parent spanRef must be cast otherwise it breaks debugging classloading (junit loads it early)
+ void clientSpan(TraceAssert trace, int index, Object parentSpan, String method = "GET", boolean renameService = false, boolean tagQueryString = false, URI uri = server.address.resolve("/success"), Integer status = 200, Throwable exception = null) {
+ super.clientSpan(trace, index, parentSpan, method, renameService, tagQueryString, uri, status, exception)
+ if (!exception) {
+ trace.span(index + 1) {
+ childOf(trace.span(index))
+ serviceName renameService ? "localhost" : "unnamed-java-app"
+ operationName "netty.client.request"
+ resourceName "$method $uri.path"
+ spanType DDSpanTypes.HTTP_CLIENT
+ errored exception != null
+ tags {
+ defaultTags()
+ if (exception) {
+ errorTags(exception.class, exception.message)
+ }
+ "$Tags.COMPONENT.key" NettyHttpClientDecorator.DECORATE.component()
+ if (status) {
+ "$Tags.HTTP_STATUS.key" status
+ }
+ "$Tags.HTTP_URL.key" "${uri.resolve(uri.path)}"
+ if (tagQueryString) {
+ "$DDTags.HTTP_QUERY" uri.query
+ "$DDTags.HTTP_FRAGMENT" { it == null || it == uri.fragment } // Optional
+ }
+ "$Tags.PEER_HOSTNAME.key" "localhost"
+ "$Tags.PEER_PORT.key" uri.port
+ "$Tags.PEER_HOST_IPV4.key" { it == null || it == "127.0.0.1" } // Optional
+ "$Tags.HTTP_METHOD.key" method
+ "$Tags.SPAN_KIND.key" Tags.SPAN_KIND_CLIENT
+ }
+ }
+ }
+ }
+
+ @Override
+ int size(int size) {
+ return size + 1
+ }
+
+ boolean testRedirects() {
+ false
+ }
+
+ boolean testConnectionFailure() {
+ false
+ }
}

MODIFY
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java
@@ -18,6 +18,7 @@ import groovy.lang.Closure;
import groovy.lang.DelegatesTo;
import groovy.transform.stc.ClosureParams;
import groovy.transform.stc.SimpleType;
+import io.opentracing.Span;
import io.opentracing.Tracer;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
@@ -202,17 +203,15 @@ public abstract class AgentTestRunner extends Specification {
}
public void blockUntilChildSpansFinished(final int numberOfSpans) throws InterruptedException {
- final DDSpan span = (DDSpan) io.opentracing.util.GlobalTracer.get().activeSpan();
- if (span == null) {
- // If there is no active span avoid getting an NPE
- return;
- }
- final PendingTrace pendingTrace = span.context().getTrace();
+ final Span span = io.opentracing.util.GlobalTracer.get().activeSpan();
+ if (span instanceof DDSpan) {
+ final PendingTrace pendingTrace = ((DDSpan) span).context().getTrace();
while (pendingTrace.size() < numberOfSpans) {
Thread.sleep(10);
}
}
+ }
public static class TestRunnerListener implements AgentBuilder.Listener {
private static final List<AgentTestRunner> activeTests = new CopyOnWriteArrayList<>();

MODIFY
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy
dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/base/HttpClientTest.groovy
@@ -188,7 +188,7 @@ abstract class HttpClientTest<T extends HttpClientDecorator> extends AgentTestRu
def "trace request with callback and no parent"() {
when:
def status = doRequest(method, server.address.resolve("/success"), ["is-dd-server": "false"]) {
- runUnderTrace("child") {
+ runUnderTrace("callback") {
// Ensure consistent ordering of traces for assertion.
TEST_WRITER.waitForTraces(1)
}
@@ -202,7 +202,7 @@ abstract class HttpClientTest<T extends HttpClientDecorator> extends AgentTestRu
clientSpan(it, 0, null, method, false)
}
trace(1, 1) {
- basicSpan(it, 0, "child")
+ basicSpan(it, 0, "callback")
}
}

Fix trace propagation