diff --git a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java index 29e33a3dd8c..518475fc853 100644 --- a/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java +++ b/dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java @@ -3,11 +3,12 @@ import static datadog.trace.api.cache.RadixTreeCache.UNSET_STATUS; import static datadog.trace.api.datastreams.DataStreamsContext.fromTags; import static datadog.trace.api.gateway.Events.EVENTS; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.traceConfig; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.DDTags; import datadog.trace.api.function.TriConsumer; @@ -26,6 +27,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities; import datadog.trace.bootstrap.instrumentation.api.TagContext; import datadog.trace.bootstrap.instrumentation.api.Tags; @@ -124,12 +126,14 @@ protected AgentTracer.TracerAPI tracer() { return AgentTracer.get(); } - public AgentSpanContext.Extracted extract(REQUEST_CARRIER carrier) { + public Context extract(REQUEST_CARRIER carrier) { AgentPropagation.ContextVisitor getter = getter(); if (null == carrier || null == getter) { return null; } - return extractContextAndGetSpanContext(carrier, getter); + + return Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getRootContext(), carrier, getter); } /** Deprecated. Use {@link #startSpan(String, Object, AgentSpanContext.Extracted)} instead. */ diff --git a/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java b/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java index a18d138c35e..39c478baf18 100644 --- a/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java @@ -20,6 +20,7 @@ import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import net.bytebuddy.asm.Advice; import scala.concurrent.Future; @@ -80,7 +81,11 @@ public static AgentScope methodEnter( if (request != null) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, headers); + defaultPropagator() + .inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), + request, + headers); // Request is immutable, so we have to assign new value once we update headers request = headers.getRequest(); } diff --git a/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java index a339eff7542..18db4d44c0f 100644 --- a/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java @@ -1,22 +1,25 @@ package datadog.trace.instrumentation.akkahttp; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.akkahttp.AkkaHttpServerDecorator.DECORATE; import akka.http.scaladsl.model.HttpRequest; import akka.http.scaladsl.model.HttpResponse; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; public class DatadogWrapperHelper { public static AgentScope createSpan(final HttpRequest request) { - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request); - final AgentSpan span = DECORATE.startSpan(request, extractedContext); + final Context extractedContext = DECORATE.extract(request); + AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(request, extractedSpanContext); DECORATE.afterStart(span); - DECORATE.onRequest(span, request, request, extractedContext); + DECORATE.onRequest(span, request, request, extractedSpanContext); - return activateSpan(span); + return (AgentScope) extractedContext.with(span).attach(); } public static void finishSpan(final AgentSpan span, final HttpResponse response) { diff --git a/dd-java-agent/instrumentation/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java b/dd-java-agent/instrumentation/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java index 105ac824f93..b828ed349c4 100644 --- a/dd-java-agent/instrumentation/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java +++ b/dd-java-agent/instrumentation/akka-http/akka-http-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java @@ -8,6 +8,7 @@ import akka.http.scaladsl.HttpExt; import akka.http.scaladsl.model.HttpRequest; import akka.http.scaladsl.model.HttpResponse; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -30,8 +31,12 @@ public static AgentScope methodEnter( if (request != null) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, headers); - // Request is immutable, so we have to assign new value once we update headers + defaultPropagator() + .inject( + Context.current().with(span).with(dsmContext), + request, + headers); // Request is immutable, so we have to assign new value once we update + // headers request = headers.getRequest(); } return activateSpan(span); diff --git a/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java b/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java index 99dfee7ba3c..d75d7026d90 100644 --- a/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java +++ b/dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java @@ -5,6 +5,7 @@ import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE; import static datadog.trace.instrumentation.apachehttpasyncclient.HttpHeadersInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; @@ -36,8 +37,7 @@ public HttpRequest generateRequest() throws IOException, HttpException { DECORATE.onRequest(span, new HostAndRequestAsHttpUriRequest(delegate.getTarget(), request)); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); - + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); return request; } diff --git a/dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java b/dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java index e9f726f6fdd..24f8f22238c 100644 --- a/dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java +++ b/dd-java-agent/instrumentation/apache-httpclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.apachehttpclient.ApacheHttpClientDecorator.HTTP_REQUEST; import static datadog.trace.instrumentation.apachehttpclient.HttpHeadersInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -47,7 +48,7 @@ private static AgentScope activateHttpSpan(final HttpUriRequest request) { // AWS calls are often signed, so we can't add headers without breaking the signature. if (!awsClientCall) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); } return scope; diff --git a/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java b/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java index 673656e75ea..b4beaa1b32f 100644 --- a/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java +++ b/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java @@ -5,6 +5,7 @@ import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE; import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.io.IOException; @@ -29,7 +30,7 @@ public void sendRequest(HttpRequest request, EntityDetails entityDetails, HttpCo DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); delegate.sendRequest(request, entityDetails, context); } } diff --git a/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java b/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java index 01b8ea4814d..3fa46229ad3 100644 --- a/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java +++ b/dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST; import static datadog.trace.instrumentation.apachehttpclient5.HttpHeadersInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -47,7 +48,7 @@ private static AgentScope activateHttpSpan(final HttpRequest request) { // AWS calls are often signed, so we can't add headers without breaking the signature. if (!awsClientCall) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); } return scope; diff --git a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java index ab192c56df6..f0b90f3acd4 100644 --- a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java +++ b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java @@ -16,6 +16,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.agent.tooling.muzzle.Reference; @@ -137,7 +138,8 @@ public static AgentScope before( span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call); if (null != span) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), headers, SETTER); + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), headers, SETTER); return activateSpan(span); } } diff --git a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/TracingServerInterceptor.java b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/TracingServerInterceptor.java index 2e3026f9cad..6a1bd4bb879 100644 --- a/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/TracingServerInterceptor.java +++ b/dd-java-agent/instrumentation/armeria-grpc/src/main/java/datadog/trace/instrumentation/armeria/grpc/server/TracingServerInterceptor.java @@ -2,7 +2,6 @@ import static datadog.trace.api.datastreams.DataStreamsContext.fromTags; import static datadog.trace.api.gateway.Events.EVENTS; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.armeria.grpc.server.GrpcExtractAdapter.GETTER; @@ -11,6 +10,8 @@ import static datadog.trace.instrumentation.armeria.grpc.server.GrpcServerDecorator.GRPC_SERVER; import static datadog.trace.instrumentation.armeria.grpc.server.GrpcServerDecorator.SERVER_PATHWAY_EDGE_TAGS; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; @@ -23,6 +24,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.api.TagContext; import io.grpc.ForwardingServerCall; import io.grpc.ForwardingServerCallListener; @@ -62,13 +64,19 @@ public ServerCall.Listener interceptCall( return next.startCall(call, headers); } - AgentSpanContext spanContext = extractContextAndGetSpanContext(headers, GETTER); + Context extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), headers, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + AgentSpanContext extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); AgentTracer.TracerAPI tracer = tracer(); - spanContext = callIGCallbackRequestStarted(tracer, spanContext); + extractedSpanContext = callIGCallbackRequestStarted(tracer, extractedSpanContext); CallbackProvider cbp = tracer.getCallbackProvider(RequestContextSlot.APPSEC); final AgentSpan span = - startSpan(DECORATE.instrumentationNames()[0], GRPC_SERVER, spanContext).setMeasured(true); + startSpan(DECORATE.instrumentationNames()[0], GRPC_SERVER, extractedSpanContext) + .setMeasured(true); AgentTracer.get() .getDataStreamsMonitoring() @@ -85,7 +93,7 @@ public ServerCall.Listener interceptCall( DECORATE.onCall(span, call); final ServerCall.Listener result; - try (AgentScope scope = activateSpan(span)) { + try (AgentScope scope = (AgentScope) extractedContext.with(span).attach()) { // Wrap the server call so that we can decorate the span // with the resulting status final TracingServerCall tracingServerCall = new TracingServerCall<>(span, call); diff --git a/dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java b/dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java index a7ca3e02bb7..8f2f8b8b184 100644 --- a/dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-eventbridge-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java @@ -87,7 +87,7 @@ private String getTraceContextToInject( jsonBuilder.append('{'); // Inject context - datadog.context.Context context = span; + datadog.context.Context context = datadog.context.Context.current().with(span); if (traceConfig().isDataStreamsEnabled()) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(eventBusName)); context = context.with(dsmContext); diff --git a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java index 04bd95b5bc8..817aa42f4fe 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java @@ -17,6 +17,7 @@ import com.amazonaws.Response; import com.amazonaws.handlers.HandlerContextKey; import com.amazonaws.handlers.RequestHandler2; +import datadog.context.Context; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.datastreams.AgentDataStreamsMonitoring; @@ -71,7 +72,8 @@ public void beforeRequest(final Request request) { request.addHandlerContext(SPAN_CONTEXT_KEY, span); if (Config.get().isAwsPropagationEnabled()) { try { - Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, request, DECORATE); + Propagators.forConcern(XRAY_TRACING_CONCERN) + .inject(Context.current().with(span), request, DECORATE); } catch (Throwable e) { log.warn("Unable to inject trace header", e); } diff --git a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java index a739a7bb5de..e6e64084417 100644 --- a/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java @@ -72,7 +72,8 @@ public SdkHttpRequest modifyHttpRequest( final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); if (span != null) { SdkHttpRequest.Builder requestBuilder = context.httpRequest().toBuilder(); - Propagators.forConcern(XRAY_TRACING_CONCERN).inject(span, requestBuilder, DECORATE); + Propagators.forConcern(XRAY_TRACING_CONCERN) + .inject(datadog.context.Context.current().with(span), requestBuilder, DECORATE); return requestBuilder.build(); } } catch (Throwable e) { diff --git a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java index 3d1333a777f..29cd937f206 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java @@ -38,7 +38,7 @@ private ByteBuffer getMessageAttributeValueToInject( final AgentSpan span = newSpan(request); StringBuilder jsonBuilder = new StringBuilder(); jsonBuilder.append('{'); - Context context = span; + Context context = Context.current().with(span); if (traceConfig().isDataStreamsEnabled()) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(snsTopicName)); context = context.with(dsmContext); diff --git a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java index 20143055b2a..d0208cacf77 100644 --- a/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java @@ -38,7 +38,7 @@ private SdkBytes getMessageAttributeValueToInject( final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); StringBuilder jsonBuilder = new StringBuilder(); jsonBuilder.append('{'); - datadog.context.Context context = span; + datadog.context.Context context = datadog.context.Context.current().with(span); if (traceConfig().isDataStreamsEnabled()) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(snsTopicName)); context = context.with(dsmContext); diff --git a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java index 4b353f12591..c602bbc54ba 100644 --- a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java @@ -23,6 +23,7 @@ import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; @@ -84,7 +85,7 @@ public AmazonWebServiceRequest beforeMarshalling(AmazonWebServiceRequest request private Context newContext(AmazonWebServiceRequest request, String queueUrl) { AgentSpan span = newSpan(request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(queueUrl)); - return span.with(dsmContext); + return Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext); } private AgentSpan newSpan(AmazonWebServiceRequest request) { diff --git a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java index 8eb154800c9..98168f606e1 100644 --- a/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.aws.v1.sqs; import static datadog.trace.api.datastreams.DataStreamsContext.create; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -19,10 +18,13 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS; import com.amazonaws.services.sqs.model.Message; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.Iterator; import java.util.LinkedHashMap; import org.slf4j.Logger; @@ -62,12 +64,18 @@ protected void startNewMessageSpan(Message message) { closePrevious(true); if (message != null) { AgentSpan queueSpan = null; + Context extractedContext = null; if (batchContext == null) { // first grab any incoming distributed context - AgentSpanContext spanContext = - Config.get().isSqsPropagationEnabled() - ? extractContextAndGetSpanContext(message, GETTER) - : null; + AgentSpanContext spanContext = null; + if (Config.get().isSqsPropagationEnabled()) { + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), message, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + spanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + } // next add a time-in-queue span for non-legacy SQS traces if (TIME_IN_QUEUE_ENABLED) { long timeInQueueStart = GETTER.extractTimeInQueueStart(message); diff --git a/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java b/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java index bc5729e49c1..af06d06e70c 100644 --- a/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java +++ b/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java @@ -14,6 +14,7 @@ import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.InstanceStore; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; @@ -96,7 +97,7 @@ private datadog.context.Context getContext( ExecutionAttributes executionAttributes, String queueUrl) { AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE); DataStreamsContext dsmContext = DataStreamsContext.fromTags(getTags(queueUrl)); - return span.with(dsmContext); + return Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext); } private LinkedHashMap getTags(String queueUrl) { diff --git a/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java b/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java index 2213ac64d6e..42f740392e4 100644 --- a/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java +++ b/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/TracingIterator.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.aws.v2.sqs; import static datadog.trace.api.datastreams.DataStreamsContext.create; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -18,10 +17,13 @@ import static datadog.trace.instrumentation.aws.v2.sqs.SqsDecorator.TIME_IN_QUEUE_ENABLED; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.Iterator; import java.util.LinkedHashMap; import org.slf4j.Logger; @@ -64,12 +66,18 @@ protected void startNewMessageSpan(Message message) { closePrevious(true); if (message != null) { AgentSpan queueSpan = null; + Context extractedContext = null; if (batchContext == null) { // first grab any incoming distributed context - AgentSpanContext spanContext = - Config.get().isSqsPropagationEnabled() - ? extractContextAndGetSpanContext(message, GETTER) - : null; + AgentSpanContext spanContext = null; + if (Config.get().isSqsPropagationEnabled()) { + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), message, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + spanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + } // next add a time-in-queue span for non-legacy SQS traces if (TIME_IN_QUEUE_ENABLED) { long timeInQueueStart = GETTER.extractTimeInQueueStart(message); @@ -99,7 +107,11 @@ protected void startNewMessageSpan(Message message) { CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, queueUrl, requestId); - activateNext(span); + if (extractedContext == null) { + activateNext(span); + } else { + extractedContext.with(span).attach(); + } if (queueSpan != null) { BROKER_DECORATE.beforeFinish(queueSpan); queueSpan.finish(); diff --git a/dd-java-agent/instrumentation/axis-2/src/main/java/datadog/trace/instrumentation/axis2/AxisTransportInstrumentation.java b/dd-java-agent/instrumentation/axis-2/src/main/java/datadog/trace/instrumentation/axis2/AxisTransportInstrumentation.java index 863b2c38009..0193ca8215d 100644 --- a/dd-java-agent/instrumentation/axis-2/src/main/java/datadog/trace/instrumentation/axis2/AxisTransportInstrumentation.java +++ b/dd-java-agent/instrumentation/axis-2/src/main/java/datadog/trace/instrumentation/axis2/AxisTransportInstrumentation.java @@ -12,6 +12,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; @@ -77,7 +78,7 @@ public static AgentScope beginTransport(@Advice.Argument(0) final MessageContext message.setProperty("TRANSPORT_HEADERS", headers); } try { - defaultPropagator().inject(span, headers, SETTER); + defaultPropagator().inject(Context.current().with(span), headers, SETTER); } catch (Throwable ignore) { } diff --git a/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java b/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java index 604b5d77c42..03133702412 100644 --- a/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java +++ b/dd-java-agent/instrumentation/azure-functions/src/main/java/datadog/trace/instrumentation/azure/functions/AzureFunctionsInstrumentation.java @@ -3,7 +3,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.declaresMethod; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.isAnnotatedWith; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.decorator.http.HttpResourceDecorator.HTTP_RESOURCE_DECORATOR; import static datadog.trace.instrumentation.azurefunctions.AzureFunctionsDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; @@ -14,6 +13,7 @@ import com.microsoft.azure.functions.ExecutionContext; import com.microsoft.azure.functions.HttpRequestMessage; import com.microsoft.azure.functions.HttpResponseMessage; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -67,13 +67,17 @@ public static class AzureFunctionsAdvice { public static AgentScope methodEnter( @Advice.Argument(0) final HttpRequestMessage request, @Advice.Argument(1) final ExecutionContext context) { - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request); - final AgentSpan span = DECORATE.startSpan(request, extractedContext); + final Context extractedContext = DECORATE.extract(request); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + + final AgentSpan span = DECORATE.startSpan(request, extractedSpanContext); DECORATE.afterStart(span, context.getFunctionName()); - DECORATE.onRequest(span, request, request, extractedContext); + DECORATE.onRequest(span, request, request, extractedSpanContext); HTTP_RESOURCE_DECORATOR.withRoute( span, request.getHttpMethod().name(), request.getUri().getPath()); - return activateSpan(span); + return (AgentScope) extractedContext.with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java b/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java index be1caba302d..9112b9b7fdd 100644 --- a/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/commons-httpclient-2/src/main/java/datadog/trace/instrumentation/commonshttpclient/CommonsHttpClientInstrumentation.java @@ -20,6 +20,7 @@ import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import net.bytebuddy.asm.Advice; import org.apache.commons.httpclient.HttpClient; import org.apache.commons.httpclient.HttpMethod; @@ -70,7 +71,11 @@ public static AgentScope methodEnter(@Advice.Argument(1) final HttpMethod httpMe DECORATE.afterStart(span); DECORATE.onRequest(span, httpMethod); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), httpMethod, SETTER); + defaultPropagator() + .inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), + httpMethod, + SETTER); return scope; } catch (BlockingException e) { diff --git a/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java b/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java index 931924a9637..5e13bbdcaf3 100644 --- a/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java +++ b/dd-java-agent/instrumentation/google-http-client/src/main/java/datadog/trace/instrumentation/googlehttpclient/GoogleHttpClientDecorator.java @@ -5,6 +5,7 @@ import com.google.api.client.http.HttpRequest; import com.google.api.client.http.HttpResponse; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.URIUtils; @@ -39,7 +40,7 @@ public AgentSpan prepareSpan(AgentSpan span, HttpRequest request) { DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); return span; } diff --git a/dd-java-agent/instrumentation/google-pubsub/src/main/java/datadog/trace/instrumentation/googlepubsub/PublisherInstrumentation.java b/dd-java-agent/instrumentation/google-pubsub/src/main/java/datadog/trace/instrumentation/googlepubsub/PublisherInstrumentation.java index 87f8413d80b..ae19dd68601 100644 --- a/dd-java-agent/instrumentation/google-pubsub/src/main/java/datadog/trace/instrumentation/googlepubsub/PublisherInstrumentation.java +++ b/dd-java-agent/instrumentation/google-pubsub/src/main/java/datadog/trace/instrumentation/googlepubsub/PublisherInstrumentation.java @@ -19,6 +19,7 @@ import com.google.auto.service.AutoService; import com.google.cloud.pubsub.v1.Publisher; import com.google.pubsub.v1.PubsubMessage; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -82,7 +83,7 @@ public static AgentScope before( PubsubMessage.Builder builder = msg.toBuilder(); DataStreamsContext dsmContext = DataStreamsContext.fromTags(sortedTags); - defaultPropagator().inject(span.with(dsmContext), builder, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), builder, SETTER); msg = builder.build(); return activateSpan(span); } diff --git a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java index f556f9bf920..e39a52b027c 100644 --- a/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly-2/src/main/java/datadog/trace/instrumentation/grizzly/GrizzlyHttpHandlerInstrumentation.java @@ -1,13 +1,13 @@ package datadog.trace.instrumentation.grizzly; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.grizzly.GrizzlyDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.CorrelationIdentifier; @@ -73,12 +73,15 @@ public static class HandleAdvice { return false; } - final AgentSpanContext.Extracted parentContext = DECORATE.extract(request); - final AgentSpan span = DECORATE.startSpan(request, parentContext); - DECORATE.afterStart(span); - DECORATE.onRequest(span, request, request, parentContext); + final Context extractedContext = DECORATE.extract(request); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); - scope = activateSpan(span); + final AgentSpan span = DECORATE.startSpan(request, extractedSpanContext); + DECORATE.afterStart(span); + DECORATE.onRequest(span, request, request, extractedSpanContext); + scope = (AgentScope) extractedContext.with(span).attach(); request.setAttribute(DD_SPAN_ATTRIBUTE, span); request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); diff --git a/dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java b/dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java index e21ca17dc43..9483372c4a2 100644 --- a/dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java +++ b/dd-java-agent/instrumentation/grizzly-client-1.9/src/main/java/datadog/trace/instrumentation/grizzly/client/AsyncHttpClientInstrumentation.java @@ -14,17 +14,23 @@ import com.google.auto.service.AutoService; import com.ning.http.client.AsyncHandler; import com.ning.http.client.Request; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.Collections; import net.bytebuddy.asm.Advice; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; @AutoService(InstrumenterModule.class) public final class AsyncHttpClientInstrumentation extends InstrumenterModule.Tracing implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + private static final Logger LOGGER = + LoggerFactory.getLogger(AsyncHttpClientInstrumentation.class); public AsyncHttpClientInstrumentation() { super("grizzly-client", "ning"); @@ -70,7 +76,8 @@ public static void onEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + Context current = Java8BytecodeBridge.getCurrentContext(); + defaultPropagator().inject(current.with(span).with(dsmContext), request, SETTER); handler = new AsyncHandlerAdapter<>(span, parentSpan, handler); } } diff --git a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java index 1f8029a0940..9dadfc03a66 100644 --- a/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java +++ b/dd-java-agent/instrumentation/grizzly-http-2.3.20/src/main/java/datadog/trace/instrumentation/grizzlyhttp232/GrizzlyDecorator.java @@ -1,8 +1,7 @@ package datadog.trace.instrumentation.grizzlyhttp232; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; - import datadog.appsec.api.blocking.BlockingContentType; +import datadog.context.Context; import datadog.trace.api.gateway.BlockResponseFunction; import datadog.trace.api.gateway.Flow; import datadog.trace.api.gateway.RequestContext; @@ -114,13 +113,17 @@ public static NextAction onHttpCodecFilterExit( } HttpRequestPacket httpRequest = (HttpRequestPacket) httpHeader; HttpResponsePacket httpResponse = httpRequest.getResponse(); - AgentSpanContext.Extracted context = DECORATE.extract(httpRequest); - AgentSpan span = DECORATE.startSpan(httpRequest, context); - AgentScope scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(httpRequest); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + + AgentSpan span = DECORATE.startSpan(httpRequest, extractedSpanContext); + AgentScope scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); ctx.getAttributes().setAttribute(DD_SPAN_ATTRIBUTE, span); ctx.getAttributes().setAttribute(DD_RESPONSE_ATTRIBUTE, httpResponse); - DECORATE.onRequest(span, httpRequest, httpRequest, context); + DECORATE.onRequest(span, httpRequest, httpRequest, extractedSpanContext); Flow.Action.RequestBlockingAction rba = span.getRequestBlockingAction(); if (rba != null && thiz instanceof HttpServerFilter) { diff --git a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/ClientCallImplInstrumentation.java b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/ClientCallImplInstrumentation.java index 229064cd30b..0882cd78302 100644 --- a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/ClientCallImplInstrumentation.java +++ b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/client/ClientCallImplInstrumentation.java @@ -11,6 +11,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.datastreams.DataStreamsContext; @@ -94,7 +95,7 @@ public static AgentScope before( span = InstrumentationContext.get(ClientCall.class, AgentSpan.class).get(call); if (null != span) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), headers, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), headers, SETTER); return activateSpan(span); } return null; diff --git a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/TracingServerInterceptor.java b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/TracingServerInterceptor.java index 5be9e13a5a3..2d763cfb7b9 100644 --- a/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/TracingServerInterceptor.java +++ b/dd-java-agent/instrumentation/grpc-1.5/src/main/java/datadog/trace/instrumentation/grpc/server/TracingServerInterceptor.java @@ -2,7 +2,6 @@ import static datadog.trace.api.datastreams.DataStreamsContext.fromTags; import static datadog.trace.api.gateway.Events.EVENTS; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.grpc.server.GrpcExtractAdapter.GETTER; @@ -11,6 +10,8 @@ import static datadog.trace.instrumentation.grpc.server.GrpcServerDecorator.GRPC_SERVER; import static datadog.trace.instrumentation.grpc.server.GrpcServerDecorator.SERVER_PATHWAY_EDGE_TAGS; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.function.TriConsumer; import datadog.trace.api.function.TriFunction; @@ -23,6 +24,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.api.TagContext; import io.grpc.ForwardingServerCall; import io.grpc.ForwardingServerCallListener; @@ -62,7 +64,12 @@ public ServerCall.Listener interceptCall( return next.startCall(call, headers); } - AgentSpanContext spanContext = extractContextAndGetSpanContext(headers, GETTER); + final Context extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), headers, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + AgentSpanContext spanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); AgentTracer.TracerAPI tracer = tracer(); spanContext = callIGCallbackRequestStarted(tracer, spanContext); @@ -84,7 +91,7 @@ public ServerCall.Listener interceptCall( DECORATE.onCall(span, call); final ServerCall.Listener result; - try (AgentScope scope = activateSpan(span)) { + try (AgentScope scope = (AgentScope) extractedContext.with(span).attach()) { // Wrap the server call so that we can decorate the span // with the resulting status final TracingServerCall tracingServerCall = new TracingServerCall<>(span, call); diff --git a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java index 9d5dd1fc9d4..f59e38aecd5 100644 --- a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java +++ b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java @@ -11,6 +11,7 @@ import static net.bytebuddy.matcher.ElementMatchers.isPublic; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.InstrumenterConfig; @@ -87,7 +88,7 @@ public static HttpUrlState methodEnter( final AgentSpan span = state.start(thiz); if (!connected) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), thiz, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), thiz, SETTER); } } return state; diff --git a/dd-java-agent/instrumentation/java-http-client/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java b/dd-java-agent/instrumentation/java-http-client/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java index 16e8630f7ec..aef3dc56e9e 100644 --- a/dd-java-agent/instrumentation/java-http-client/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java +++ b/dd-java-agent/instrumentation/java-http-client/src/main/java11/datadog/trace/instrumentation/httpclient/HeadersAdvice.java @@ -6,6 +6,7 @@ import static datadog.trace.instrumentation.httpclient.HttpHeadersInjectAdapter.KEEP; import static datadog.trace.instrumentation.httpclient.HttpHeadersInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import java.net.http.HttpHeaders; @@ -20,7 +21,7 @@ public static void methodExit(@Advice.Return(readOnly = false) HttpHeaders heade final Map> headerMap = new HashMap<>(headers.map()); final AgentSpan span = activeSpan(); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), headerMap, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), headerMap, SETTER); headers = HttpHeaders.of(headerMap, KEEP); } } diff --git a/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java b/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java index b01be762333..5148c3acbab 100644 --- a/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java +++ b/dd-java-agent/instrumentation/jax-rs-client-1.1/src/main/java/datadog/trace/instrumentation/jaxrs/v1/JaxRsClientV1Instrumentation.java @@ -18,6 +18,7 @@ import com.sun.jersey.api.client.ClientHandler; import com.sun.jersey.api.client.ClientRequest; import com.sun.jersey.api.client.ClientResponse; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.datastreams.DataStreamsContext; @@ -77,7 +78,8 @@ public static AgentScope onEnter( request.getProperties().put(DD_SPAN_ATTRIBUTE, span); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request.getHeaders(), SETTER); + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), request.getHeaders(), SETTER); return activateSpan(span); } return null; diff --git a/dd-java-agent/instrumentation/jax-rs-client-2.0/src/main/java/datadog/trace/instrumentation/jaxrs/ClientTracingFilter.java b/dd-java-agent/instrumentation/jax-rs-client-2.0/src/main/java/datadog/trace/instrumentation/jaxrs/ClientTracingFilter.java index 9e7963bf67d..8ec762ee5c2 100644 --- a/dd-java-agent/instrumentation/jax-rs-client-2.0/src/main/java/datadog/trace/instrumentation/jaxrs/ClientTracingFilter.java +++ b/dd-java-agent/instrumentation/jax-rs-client-2.0/src/main/java/datadog/trace/instrumentation/jaxrs/ClientTracingFilter.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.jaxrs.JaxRsClientDecorator.DECORATE; import static datadog.trace.instrumentation.jaxrs.JaxRsClientDecorator.JAX_RS_CLIENT_CALL; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -30,7 +31,9 @@ public void filter(final ClientRequestContext requestContext) { DECORATE.onRequest(span, requestContext); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), requestContext.getHeaders(), SETTER); + defaultPropagator() + .inject( + Context.current().with(span).with(dsmContext), requestContext.getHeaders(), SETTER); requestContext.setProperty(SPAN_PROPERTY_NAME, span); } diff --git a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java index 61c4e9be271..0c9040b00f2 100644 --- a/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java +++ b/dd-java-agent/instrumentation/jetty-11/src/main/java11/datadog/trace/instrumentation/jetty11/JettyServerAdvice.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty11.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.CorrelationIdentifier; import datadog.trace.api.GlobalTracer; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -26,12 +27,15 @@ public static AgentScope onEnter( return activateSpan((AgentSpan) existingSpan); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req); - span = DECORATE.startSpan(req, extractedContext); - final AgentScope scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(req, extractedSpanContext); + final AgentScope scope = (AgentScope) extractedContext.with(span).attach(); span.setMeasured(true); DECORATE.afterStart(span); - DECORATE.onRequest(span, req, req, extractedContext); + DECORATE.onRequest(span, req, req, extractedSpanContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); diff --git a/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java b/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java index 11b31cdffca..39cf8e765e3 100644 --- a/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java +++ b/dd-java-agent/instrumentation/jetty-12/src/main/java17/datadog/trace/instrumentation/jetty12/JettyServerAdvice.java @@ -3,6 +3,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; +import datadog.context.Context; import datadog.trace.api.CorrelationIdentifier; import datadog.trace.api.GlobalTracer; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -28,12 +29,15 @@ public static void onExit( } } - final AgentSpanContext.Extracted extractedContext = JettyDecorator.DECORATE.extract(req); - final AgentSpan span = JettyDecorator.DECORATE.startSpan(req, extractedContext); - try (final AgentScope scope = activateSpan(span)) { + final Context extractedContext = JettyDecorator.DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = JettyDecorator.DECORATE.startSpan(req, extractedSpanContext); + try (final AgentScope scope = (AgentScope) extractedContext.with(span).attach()) { span.setMeasured(true); JettyDecorator.DECORATE.afterStart(span); - JettyDecorator.DECORATE.onRequest(span, req, req, extractedContext); + JettyDecorator.DECORATE.onRequest(span, req, req, extractedSpanContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); diff --git a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java index e8a9e449963..16a235d5a83 100644 --- a/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.0/src/main/java/datadog/trace/instrumentation/jetty70/JettyServerInstrumentation.java @@ -10,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.Config; @@ -156,11 +157,14 @@ public static AgentScope onEnter( return activateSpan((AgentSpan) existingSpan); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req); - span = DECORATE.startSpan(req, extractedContext); - final AgentScope scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(req, extractedSpanContext); + final AgentScope scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, req, req, extractedContext); + DECORATE.onRequest(span, req, req, extractedSpanContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); diff --git a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java index 81225c7c820..8c9248f9df0 100644 --- a/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-7.6/src/main/java/datadog/trace/instrumentation/jetty76/JettyServerInstrumentation.java @@ -10,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.Config; @@ -157,11 +158,14 @@ public static AgentScope onEnter( return activateSpan((AgentSpan) existingSpan); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req); - span = DECORATE.startSpan(req, extractedContext); - final AgentScope scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(req, extractedSpanContext); + final AgentScope scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, req, req, extractedContext); + DECORATE.onRequest(span, req, req, extractedSpanContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java index 7aaea42fc4d..650711fa4fd 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java/datadog/trace/instrumentation/jetty9/JettyServerInstrumentation.java @@ -15,6 +15,7 @@ import com.google.auto.service.AutoService; import datadog.appsec.api.blocking.BlockingException; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -171,11 +172,14 @@ public static AgentScope onEnter( return activateSpan((AgentSpan) existingSpan); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req); - span = DECORATE.startSpan(req, extractedContext); - final AgentScope scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(req, extractedSpanContext); + final AgentScope scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, req, req, extractedContext); + DECORATE.onRequest(span, req, req, extractedSpanContext); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); diff --git a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java index 25459847c0d..274a5838900 100644 --- a/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java +++ b/dd-java-agent/instrumentation/jetty-9/src/main/java_jetty10/datadog/trace/instrumentation/jetty10/HandleAdvice.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.jetty10.JettyDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.CorrelationIdentifier; import datadog.trace.api.GlobalTracer; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -25,12 +26,15 @@ public static AgentScope onEnter( return activateSpan((AgentSpan) existingSpan); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req); - span = DECORATE.startSpan(req, extractedContext); + final Context extractedContext = DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(req, extractedSpanContext); DECORATE.afterStart(span); - DECORATE.onRequest(span, req, req, extractedContext); + DECORATE.onRequest(span, req, req, extractedSpanContext); - final AgentScope scope = activateSpan(span); + final AgentScope scope = (AgentScope) extractedContext.with(span).attach(); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId()); diff --git a/dd-java-agent/instrumentation/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java b/dd-java-agent/instrumentation/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java index 58c64511ee9..1170e060ac0 100644 --- a/dd-java-agent/instrumentation/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java +++ b/dd-java-agent/instrumentation/jetty-client/jetty-client-10.0/src/main/java11/datadog/trace/instrumentation/jetty_client10/SendAdvice.java @@ -7,6 +7,7 @@ import static datadog.trace.instrumentation.jetty_client10.JettyClientDecorator.DECORATE; import static datadog.trace.instrumentation.jetty_client10.JettyClientDecorator.HTTP_REQUEST; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -27,7 +28,7 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); return span; } diff --git a/dd-java-agent/instrumentation/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java b/dd-java-agent/instrumentation/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java index 5575808b567..ec5e326b65c 100644 --- a/dd-java-agent/instrumentation/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java +++ b/dd-java-agent/instrumentation/jetty-client/jetty-client-12.0/src/main/java17/datadog/trace/instrumentation/jetty_client12/SendAdvice.java @@ -7,6 +7,7 @@ import static datadog.trace.instrumentation.jetty_client12.HeadersInjectAdapter.SETTER; import static datadog.trace.instrumentation.jetty_client12.JettyClientDecorator.HTTP_REQUEST; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -23,7 +24,7 @@ public static AgentScope methodEnter(@Advice.This final HttpRequest request) { JettyClientDecorator.DECORATE.afterStart(span); JettyClientDecorator.DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); return activateSpan(span); } diff --git a/dd-java-agent/instrumentation/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java b/dd-java-agent/instrumentation/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java index e6328faefcc..d8d58a1c920 100644 --- a/dd-java-agent/instrumentation/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java +++ b/dd-java-agent/instrumentation/jetty-client/jetty-client-9.1/src/main/java/datadog/trace/instrumentation/jetty_client91/JettyClientInstrumentation.java @@ -15,6 +15,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -93,7 +94,7 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); return span; } diff --git a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/DatadogMessageListener.java b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/DatadogMessageListener.java index b58034f3d88..dbd5d254929 100644 --- a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/DatadogMessageListener.java +++ b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/DatadogMessageListener.java @@ -1,6 +1,5 @@ package datadog.trace.instrumentation.jms; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.jms.JMSDecorator.BROKER_DECORATE; @@ -11,10 +10,13 @@ import static datadog.trace.instrumentation.jms.MessageExtractAdapter.GETTER; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.jms.MessageConsumerState; import datadog.trace.bootstrap.instrumentation.jms.SessionState; import javax.jms.Message; @@ -38,9 +40,15 @@ public DatadogMessageListener( @Override public void onMessage(Message message) { AgentSpan span; + Context extractedContext = null; AgentSpanContext propagatedContext = null; if (!consumerState.isPropagationDisabled()) { - propagatedContext = extractContextAndGetSpanContext(message, GETTER); + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), message, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + propagatedContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); } long startMillis = GETTER.extractTimeInQueueStart(message); if (startMillis == 0 || !TIME_IN_QUEUE_ENABLED) { @@ -70,7 +78,10 @@ public void onMessage(Message message) { // span will be finished by Session.commit/rollback/close sessionState.finishOnCommit(span); } - try (AgentScope scope = activateSpan(span)) { + try (AgentScope scope = + extractedContext == null + ? activateSpan(span) + : (AgentScope) extractedContext.with(span).attach()) { messageListener.onMessage(message); } catch (RuntimeException | Error thrown) { CONSUMER_DECORATE.onError(span, thrown); diff --git a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java index 446dc32897f..7b882af4549 100644 --- a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageConsumerInstrumentation.java @@ -3,7 +3,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -19,11 +18,14 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.jms.MessageConsumerState; import datadog.trace.bootstrap.instrumentation.jms.SessionState; import javax.jms.Message; @@ -118,9 +120,15 @@ public static void afterReceive( } AgentSpan span; + Context extractedContext = null; AgentSpanContext propagatedContext = null; if (!consumerState.isPropagationDisabled()) { - propagatedContext = extractContextAndGetSpanContext(message, GETTER); + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), message, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + propagatedContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); } long startMillis = GETTER.extractTimeInQueueStart(message); if (startMillis == 0 || !TIME_IN_QUEUE_ENABLED) { @@ -144,9 +152,7 @@ public static void afterReceive( CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, message, consumerState.getConsumerResourceName()); CONSUMER_DECORATE.onError(span, throwable); - - activateNext(span); // scope is left open until next message or it times out - + activateNext(span); SessionState sessionState = consumerState.getSessionState(); if (sessionState.isClientAcknowledge()) { // consumed spans will be finished by a call to Message.acknowledge diff --git a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageProducerInstrumentation.java b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageProducerInstrumentation.java index 9c5ac41a4d0..6a07d70b1fb 100644 --- a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageProducerInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/JMSMessageProducerInstrumentation.java @@ -19,6 +19,7 @@ import datadog.trace.bootstrap.InstrumentationContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.jms.MessageProducerState; import javax.jms.Destination; import javax.jms.Message; @@ -94,7 +95,8 @@ public static AgentScope beforeSend( if (JMSDecorator.canInject(message)) { if (Config.get().isJmsPropagationEnabled() && (null == producerState || !producerState.isPropagationDisabled())) { - defaultPropagator().inject(span, message, SETTER); + defaultPropagator() + .inject(Java8BytecodeBridge.getCurrentContext().with(span), message, SETTER); } if (TIME_IN_QUEUE_ENABLED) { if (null != producerState) { @@ -141,7 +143,8 @@ public static AgentScope beforeSend( if (JMSDecorator.canInject(message)) { if (Config.get().isJmsPropagationEnabled() && !Config.get().isJmsPropagationDisabledForDestination(destinationName)) - defaultPropagator().inject(span, message, SETTER); + defaultPropagator() + .inject(Java8BytecodeBridge.getCurrentContext().with(span), message, SETTER); if (TIME_IN_QUEUE_ENABLED) { MessageProducerState producerState = InstrumentationContext.get(MessageProducer.class, MessageProducerState.class) diff --git a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java index 6f2cb2c16e1..58acfdf2a43 100644 --- a/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java +++ b/dd-java-agent/instrumentation/jms/src/main/java/datadog/trace/instrumentation/jms/MDBMessageConsumerInstrumentation.java @@ -4,8 +4,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.hasSuperType; import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.jms.JMSDecorator.CONSUMER_DECORATE; import static datadog.trace.instrumentation.jms.JMSDecorator.JMS_CONSUME; @@ -16,11 +14,14 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import javax.jms.Destination; import javax.jms.JMSException; import javax.jms.Message; @@ -67,7 +68,12 @@ public static AgentScope methodEnter(@Advice.Argument(0) final Message message) if (CallDepthThreadLocalMap.incrementCallDepth(MessageListener.class) > 0) { return null; } - AgentSpanContext propagatedContext = extractContextAndGetSpanContext(message, GETTER); + final Context extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), message, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext propagatedContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); AgentSpan span = startSpan(JMS_CONSUME, propagatedContext); CONSUMER_DECORATE.afterStart(span); CharSequence consumerResourceName; @@ -81,7 +87,7 @@ public static AgentScope methodEnter(@Advice.Argument(0) final Message message) consumerResourceName = "unknown JMS destination"; } CONSUMER_DECORATE.onConsume(span, message, consumerResourceName); - return activateSpan(span); + return (AgentScope) extractedContext.with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java index a85ee9ce95d..cd0e12a063c 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/KafkaProducerInstrumentation.java @@ -37,6 +37,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.LinkedHashMap; import java.util.Map; import net.bytebuddy.asm.Advice; @@ -153,7 +154,8 @@ public static AgentScope onEnter( sortedTags.put(TOPIC_TAG, record.topic()); sortedTags.put(TYPE_TAG, "kafka"); try { - defaultPropagator().inject(span, record.headers(), setter); + defaultPropagator() + .inject(Java8BytecodeBridge.getCurrentContext().with(span), record.headers(), setter); if (STREAMING_CONTEXT.isDisabledForTopic(record.topic()) || STREAMING_CONTEXT.isSinkTopic(record.topic())) { // inject the context in the headers, but delay sending the stats until we know the @@ -161,7 +163,10 @@ public static AgentScope onEnter( // The stats are saved in the pathway context and sent in PayloadSizeAdvice. Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags); - dsmPropagator.inject(span.with(dsmContext), record.headers(), setter); + dsmPropagator.inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), + record.headers(), + setter); AvroSchemaExtractor.tryExtractProducer(record, span); } } catch (final IllegalStateException e) { @@ -175,12 +180,16 @@ record = record.value(), record.headers()); - defaultPropagator().inject(span, record.headers(), setter); + defaultPropagator() + .inject(Java8BytecodeBridge.getCurrentContext().with(span), record.headers(), setter); if (STREAMING_CONTEXT.isDisabledForTopic(record.topic()) || STREAMING_CONTEXT.isSinkTopic(record.topic())) { Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags); - dsmPropagator.inject(span.with(dsmContext), record.headers(), setter); + dsmPropagator.inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), + record.headers(), + setter); AvroSchemaExtractor.tryExtractProducer(record, span); } } diff --git a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java index 3c6f26995a3..4f37b14df64 100644 --- a/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java +++ b/dd-java-agent/instrumentation/kafka-clients-0.11/src/main/java/datadog/trace/instrumentation/kafka_clients/TracingIterator.java @@ -2,7 +2,6 @@ import static datadog.trace.api.datastreams.DataStreamsContext.create; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -21,6 +20,7 @@ import static datadog.trace.instrumentation.kafka_common.Utils.computePayloadSizeBytes; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import datadog.context.Context; import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; @@ -29,6 +29,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.Iterator; import java.util.LinkedHashMap; import org.apache.kafka.clients.consumer.ConsumerRecord; @@ -83,9 +84,14 @@ protected void startNewRecordSpan(ConsumerRecord val) { closePrevious(true); AgentSpan span, queueSpan = null; if (val != null) { + Context extractedContext = null; if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) { + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), val.headers(), GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); final AgentSpanContext spanContext = - extractContextAndGetSpanContext(val.headers(), GETTER); + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); long timeInQueueStart = GETTER.extractTimeInQueueStart(val.headers()); if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) { span = startSpan(operationName, spanContext); @@ -124,7 +130,10 @@ protected void startNewRecordSpan(ConsumerRecord val) { // for DSM users Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = create(sortedTags, val.timestamp(), payloadSize); - dsmPropagator.inject(span.with(dsmContext), val.headers(), SETTER); + dsmPropagator.inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), + val.headers(), + SETTER); } } } else { diff --git a/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerAdvice.java b/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerAdvice.java index 0756f585b52..42522b0f586 100644 --- a/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerAdvice.java +++ b/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/ProducerAdvice.java @@ -16,6 +16,7 @@ import static datadog.trace.instrumentation.kafka_clients38.KafkaDecorator.TIME_IN_QUEUE_ENABLED; import static datadog.trace.instrumentation.kafka_common.StreamingContext.STREAMING_CONTEXT; +import datadog.context.Context; import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; @@ -75,7 +76,7 @@ public static AgentScope onEnter( sortedTags.put(TOPIC_TAG, record.topic()); sortedTags.put(TYPE_TAG, "kafka"); try { - defaultPropagator().inject(span, record.headers(), setter); + defaultPropagator().inject(Context.current().with(span), record.headers(), setter); if (STREAMING_CONTEXT.isDisabledForTopic(record.topic()) || STREAMING_CONTEXT.isSinkTopic(record.topic())) { // inject the context in the headers, but delay sending the stats until we know the @@ -83,7 +84,8 @@ public static AgentScope onEnter( // The stats are saved in the pathway context and sent in PayloadSizeAdvice. Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags); - dsmPropagator.inject(span.with(dsmContext), record.headers(), setter); + dsmPropagator.inject( + Context.current().with(span).with(dsmContext), record.headers(), setter); AvroSchemaExtractor.tryExtractProducer(record, span); } } catch (final IllegalStateException e) { @@ -97,12 +99,13 @@ record = record.value(), record.headers()); - defaultPropagator().inject(span, record.headers(), setter); + defaultPropagator().inject(Context.current().with(span), record.headers(), setter); if (STREAMING_CONTEXT.isDisabledForTopic(record.topic()) || STREAMING_CONTEXT.isSinkTopic(record.topic())) { Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = fromTagsWithoutCheckpoint(sortedTags); - dsmPropagator.inject(span.with(dsmContext), record.headers(), setter); + dsmPropagator.inject( + Context.current().with(span).with(dsmContext), record.headers(), setter); AvroSchemaExtractor.tryExtractProducer(record, span); } } diff --git a/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java b/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java index 9a922d46ad0..2a90f3460c1 100644 --- a/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java +++ b/dd-java-agent/instrumentation/kafka-clients-3.8/src/main/java17/datadog/trace/instrumentation/kafka_clients38/TracingIterator.java @@ -2,7 +2,6 @@ import static datadog.trace.api.datastreams.DataStreamsContext.create; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateNext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -16,6 +15,7 @@ import static datadog.trace.instrumentation.kafka_clients38.TextMapInjectAdapter.SETTER; import static java.util.concurrent.TimeUnit.MILLISECONDS; +import datadog.context.Context; import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.api.Config; @@ -24,6 +24,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.InstrumentationTags; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.instrumentation.kafka_common.StreamingContext; import datadog.trace.instrumentation.kafka_common.Utils; import java.util.Iterator; @@ -80,9 +81,14 @@ protected void startNewRecordSpan(ConsumerRecord val) { closePrevious(true); AgentSpan span, queueSpan = null; if (val != null) { + Context extractedContext = null; if (!Config.get().isKafkaClientPropagationDisabledForTopic(val.topic())) { + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), val.headers(), GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); final AgentSpanContext spanContext = - extractContextAndGetSpanContext(val.headers(), GETTER); + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); long timeInQueueStart = GETTER.extractTimeInQueueStart(val.headers()); if (timeInQueueStart == 0 || !KafkaDecorator.TIME_IN_QUEUE_ENABLED) { span = startSpan(operationName, spanContext); @@ -124,7 +130,8 @@ protected void startNewRecordSpan(ConsumerRecord val) { // for DSM users Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = create(sortedTags, val.timestamp(), payloadSize); - dsmPropagator.inject(span.with(dsmContext), val.headers(), SETTER); + dsmPropagator.inject( + Context.current().with(span).with(dsmContext), val.headers(), SETTER); } } } else { diff --git a/dd-java-agent/instrumentation/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/KafkaStreamTaskInstrumentation.java b/dd-java-agent/instrumentation/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/KafkaStreamTaskInstrumentation.java index 10d15eb8486..00686d825e3 100644 --- a/dd-java-agent/instrumentation/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/KafkaStreamTaskInstrumentation.java +++ b/dd-java-agent/instrumentation/kafka-streams-0.11/src/main/java/datadog/trace/instrumentation/kafka_streams/KafkaStreamTaskInstrumentation.java @@ -3,7 +3,6 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; import static datadog.trace.api.datastreams.DataStreamsContext.create; import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.DSM_CONCERN; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.core.datastreams.TagsProcessor.DIRECTION_IN; @@ -31,6 +30,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.context.propagation.Propagator; import datadog.context.propagation.Propagators; import datadog.trace.agent.tooling.Instrumenter; @@ -42,6 +42,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.instrumentation.kafka_clients.TracingIterableDelegator; import java.util.LinkedHashMap; import java.util.Map; @@ -233,15 +234,20 @@ public static void start( AgentSpan span, queueSpan = null; StreamTaskContext streamTaskContext = InstrumentationContext.get(StreamTask.class, StreamTaskContext.class).get(task); + Context extractedContext = null; if (!Config.get().isKafkaClientPropagationDisabledForTopic(record.topic())) { - final AgentSpanContext extractedContext = - extractContextAndGetSpanContext(record, SR_GETTER); + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), record, SR_GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext spanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); long timeInQueueStart = SR_GETTER.extractTimeInQueueStart(record); if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) { - span = startSpan(KAFKA_CONSUME, extractedContext); + span = startSpan(KAFKA_CONSUME, spanContext); } else { queueSpan = - startSpan(KAFKA_DELIVER, extractedContext, MILLISECONDS.toMicros(timeInQueueStart)); + startSpan(KAFKA_DELIVER, spanContext, MILLISECONDS.toMicros(timeInQueueStart)); BROKER_DECORATE.afterStart(queueSpan); BROKER_DECORATE.onTimeInQueue(queueSpan, record); span = startSpan(KAFKA_CONSUME, queueSpan.context()); @@ -271,7 +277,7 @@ public static void start( if (STREAMING_CONTEXT.isSourceTopic(record.topic())) { Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = create(sortedTags, record.timestamp, payloadSize); - dsmPropagator.inject(span.with(dsmContext), record, SR_SETTER); + dsmPropagator.inject(Context.current().with(span).with(dsmContext), record, SR_SETTER); } } } else { @@ -280,7 +286,10 @@ public static void start( CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, record, node); - AgentScope agentScope = activateSpan(span); + AgentScope agentScope = + extractedContext == null + ? activateSpan(span) + : (AgentScope) extractedContext.with(span).attach(); if (null != queueSpan) { queueSpan.finish(); } @@ -309,15 +318,20 @@ public static void start( AgentSpan span, queueSpan = null; StreamTaskContext streamTaskContext = InstrumentationContext.get(StreamTask.class, StreamTaskContext.class).get(task); + Context extractedContext = null; if (!Config.get().isKafkaClientPropagationDisabledForTopic(record.topic())) { - final AgentSpanContext extractedContext = - extractContextAndGetSpanContext(record, PR_GETTER); + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), record, PR_GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext spanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); long timeInQueueStart = PR_GETTER.extractTimeInQueueStart(record); if (timeInQueueStart == 0 || !TIME_IN_QUEUE_ENABLED) { - span = startSpan(KAFKA_CONSUME, extractedContext); + span = startSpan(KAFKA_CONSUME, spanContext); } else { queueSpan = - startSpan(KAFKA_DELIVER, extractedContext, MILLISECONDS.toMicros(timeInQueueStart)); + startSpan(KAFKA_DELIVER, spanContext, MILLISECONDS.toMicros(timeInQueueStart)); BROKER_DECORATE.afterStart(queueSpan); BROKER_DECORATE.onTimeInQueue(queueSpan, record); span = startSpan(KAFKA_CONSUME, queueSpan.context()); @@ -354,7 +368,7 @@ public static void start( if (STREAMING_CONTEXT.isSourceTopic(record.topic())) { Propagator dsmPropagator = Propagators.forConcern(DSM_CONCERN); DataStreamsContext dsmContext = create(sortedTags, record.timestamp(), payloadSize); - dsmPropagator.inject(span.with(dsmContext), record, PR_SETTER); + dsmPropagator.inject(Context.current().with(span).with(dsmContext), record, PR_SETTER); } } } else { @@ -363,7 +377,10 @@ public static void start( CONSUMER_DECORATE.afterStart(span); CONSUMER_DECORATE.onConsume(span, record, node); - AgentScope agentScope = activateSpan(span); + AgentScope agentScope = + extractedContext == null + ? activateSpan(span) + : (AgentScope) extractedContext.with(span).attach(); if (null != queueSpan) { queueSpan.finish(); } diff --git a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java index 1a4769710b3..826d2cfb2d0 100644 --- a/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-20/src/main/java/datadog/trace/instrumentation/liberty20/LibertyServerInstrumentation.java @@ -14,6 +14,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.ws.webcontainer.webapp.WebApp; import com.ibm.wsspi.webcontainer.webapp.IWebAppDispatcherContext; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.ClassloaderConfigurationOverrides; @@ -105,10 +106,13 @@ public static class HandleRequestAdvice { } catch (NullPointerException e) { } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request); + final Context extractedContext = DECORATE.extract(request); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext); - final AgentSpan span = DECORATE.startSpan(request, extractedContext); - scope = activateSpan(span); + final AgentSpan span = DECORATE.startSpan(request, extractedSpanContext); + scope = (AgentScope) extractedContext.with(span).attach(); if (Config.get().isJeeSplitByDeployment()) { final IWebAppDispatcherContext dispatcherContext = request.getWebAppDispatcherContext(); if (dispatcherContext != null) { @@ -122,7 +126,7 @@ public static class HandleRequestAdvice { } } DECORATE.afterStart(span); - DECORATE.onRequest(span, request, request, extractedContext); + DECORATE.onRequest(span, request, request, extractedSpanContext); request.setAttribute(DD_SPAN_ATTRIBUTE, span); request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId()); diff --git a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java index 30c8d7980cc..63cc46154f2 100644 --- a/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java +++ b/dd-java-agent/instrumentation/liberty-23/src/main/java/datadog/trace/instrumentation/liberty23/LibertyServerInstrumentation.java @@ -14,6 +14,7 @@ import com.ibm.ws.webcontainer.srt.SRTServletResponse; import com.ibm.ws.webcontainer.webapp.WebApp; import com.ibm.wsspi.webcontainer.webapp.IWebAppDispatcherContext; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.ClassloaderConfigurationOverrides; @@ -107,10 +108,13 @@ public static class HandleRequestAdvice { } catch (NullPointerException e) { } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request); + final Context extractedContext = DECORATE.extract(request); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); request.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext); - final AgentSpan span = DECORATE.startSpan(request, extractedContext); - scope = activateSpan(span); + final AgentSpan span = DECORATE.startSpan(request, extractedSpanContext); + scope = (AgentScope) extractedContext.with(span).attach(); if (Config.get().isJeeSplitByDeployment()) { final IWebAppDispatcherContext dispatcherContext = request.getWebAppDispatcherContext(); if (dispatcherContext != null) { @@ -124,7 +128,7 @@ public static class HandleRequestAdvice { } } DECORATE.afterStart(span); - DECORATE.onRequest(span, request, request, extractedContext); + DECORATE.onRequest(span, request, request, extractedSpanContext); request.setAttribute(DD_SPAN_ATTRIBUTE, span); request.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); request.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId()); diff --git a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java index db47fba6bbd..555a7faf7fc 100644 --- a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/client/HttpClientRequestTracingHandler.java @@ -10,6 +10,7 @@ import static datadog.trace.instrumentation.netty38.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty38.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -67,7 +68,8 @@ public void writeRequested(final ChannelHandlerContext ctx, final MessageEvent m } DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request.headers(), SETTER); + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), request.headers(), SETTER); channelTraceContext.setClientSpan(span); diff --git a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java index 20c685a581b..d78a85c5d42 100644 --- a/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-3.8/src/main/java/datadog/trace/instrumentation/netty38/server/HttpServerRequestTracingHandler.java @@ -3,6 +3,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.netty38.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.ContextStore; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -44,15 +45,18 @@ public void messageReceived(final ChannelHandlerContext ctx, final MessageEvent final HttpRequest request = (HttpRequest) msg.getMessage(); final HttpHeaders headers = request.headers(); - final AgentSpanContext.Extracted context = DECORATE.extract(headers); - final AgentSpan span = DECORATE.startSpan(headers, context); + final Context extractedContext = DECORATE.extract(headers); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(headers, extractedSpanContext); channelTraceContext.reset(); channelTraceContext.setRequestHeaders(headers); - try (final AgentScope scope = activateSpan(span)) { + try (final AgentScope scope = (AgentScope) extractedContext.with(span).attach()) { DECORATE.afterStart(span); - DECORATE.onRequest(span, ctx.getChannel(), request, context); + DECORATE.onRequest(span, ctx.getChannel(), request, extractedSpanContext); channelTraceContext.setServerSpan(span); diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java index dfcbcb27ec4..4d1bc76c269 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/client/HttpClientRequestTracingHandler.java @@ -13,6 +13,7 @@ import static datadog.trace.instrumentation.netty40.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty40.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.Config; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -90,7 +91,8 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann // AWS calls are often signed, so we can't add headers without breaking the signature. if (!awsClientCall) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request.headers(), SETTER); + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), request.headers(), SETTER); } ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java index d5b0aabd49e..9612bd9ca69 100644 --- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/server/HttpServerRequestTracingHandler.java @@ -7,6 +7,7 @@ import static datadog.trace.instrumentation.netty40.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty40.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -40,12 +41,15 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final HttpRequest request = (HttpRequest) msg; final HttpHeaders headers = request.headers(); - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(headers); - final AgentSpan span = DECORATE.startSpan(headers, extractedContext); + final Context extractedContext = DECORATE.extract(headers); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(headers, extractedSpanContext); - try (final AgentScope scope = activateSpan(span)) { + try (final AgentScope scope = (AgentScope) extractedContext.with(span).attach()) { DECORATE.afterStart(span); - DECORATE.onRequest(span, channel, request, extractedContext); + DECORATE.onRequest(span, channel, request, extractedSpanContext); channel.attr(ANALYZED_RESPONSE_KEY).set(null); channel.attr(BLOCKED_RESPONSE_KEY).set(null); diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java index 3672e678054..609b41d10fa 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/client/HttpClientRequestTracingHandler.java @@ -13,6 +13,7 @@ import static datadog.trace.instrumentation.netty41.client.NettyHttpClientDecorator.NETTY_CLIENT_REQUEST; import static datadog.trace.instrumentation.netty41.client.NettyResponseInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.Config; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -91,7 +92,8 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann // AWS calls are often signed, so we can't add headers without breaking the signature. if (!awsClientCall) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request.headers(), SETTER); + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), request.headers(), SETTER); } ctx.channel().attr(SPAN_ATTRIBUTE_KEY).set(span); diff --git a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java index 3fd4e640d44..92e7bc43cd8 100644 --- a/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java +++ b/dd-java-agent/instrumentation/netty-4.1/src/main/java/datadog/trace/instrumentation/netty41/server/HttpServerRequestTracingHandler.java @@ -7,6 +7,7 @@ import static datadog.trace.instrumentation.netty41.AttributeKeys.SPAN_ATTRIBUTE_KEY; import static datadog.trace.instrumentation.netty41.server.NettyHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.gateway.Flow; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -39,12 +40,15 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { final HttpRequest request = (HttpRequest) msg; final HttpHeaders headers = request.headers(); - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(headers); - final AgentSpan span = DECORATE.startSpan(headers, extractedContext); + final Context extractedContext = DECORATE.extract(headers); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(headers, extractedSpanContext); - try (final AgentScope scope = activateSpan(span)) { + try (final AgentScope scope = (AgentScope) extractedContext.with(span).attach()) { DECORATE.afterStart(span); - DECORATE.onRequest(span, channel, request, extractedContext); + DECORATE.onRequest(span, channel, request, extractedSpanContext); channel.attr(ANALYZED_RESPONSE_KEY).set(null); channel.attr(BLOCKED_RESPONSE_KEY).set(null); diff --git a/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/TracingInterceptor.java b/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/TracingInterceptor.java index c68cdf00fe4..3ab224d9221 100644 --- a/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/TracingInterceptor.java +++ b/dd-java-agent/instrumentation/okhttp-2/src/main/java/datadog/trace/instrumentation/okhttp2/TracingInterceptor.java @@ -11,6 +11,7 @@ import com.squareup.okhttp.Interceptor; import com.squareup.okhttp.Request; import com.squareup.okhttp.Response; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -28,7 +29,8 @@ public Response intercept(final Chain chain) throws IOException { final Request.Builder requestBuilder = chain.request().newBuilder(); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), requestBuilder, SETTER); + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), requestBuilder, SETTER); final Response response; try { diff --git a/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/TracingInterceptor.java b/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/TracingInterceptor.java index fa9b076cae8..41f3ca402ba 100644 --- a/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/TracingInterceptor.java +++ b/dd-java-agent/instrumentation/okhttp-3/src/main/java/datadog/trace/instrumentation/okhttp3/TracingInterceptor.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.okhttp3.OkHttpClientDecorator.OKHTTP_REQUEST; import static datadog.trace.instrumentation.okhttp3.RequestBuilderInjectAdapter.SETTER; +import datadog.context.Context; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -31,7 +32,9 @@ public Response intercept(final Chain chain) throws IOException { final Request.Builder requestBuilder = chain.request().newBuilder(); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), requestBuilder, SETTER); + + defaultPropagator() + .inject(Context.current().with(span).with(dsmContext), requestBuilder, SETTER); final Response response; try { diff --git a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/OtelContextPropagators.java b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/OtelContextPropagators.java index ad6042d0f0e..38eeb558dac 100644 --- a/dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/OtelContextPropagators.java +++ b/dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/main/java/datadog/trace/instrumentation/opentelemetry/OtelContextPropagators.java @@ -45,7 +45,8 @@ public void inject(final Context context, final C carrier, final Setter s } AgentSpan agentSpan = converter.toAgentSpan(span); if (agentSpan != null) { - defaultPropagator().inject(agentSpan, carrier, setter::set); + defaultPropagator() + .inject(datadog.context.Context.current().with(agentSpan), carrier, setter::set); } } diff --git a/dd-java-agent/instrumentation/opentracing/api-0.31/src/main/java/datadog/trace/instrumentation/opentracing31/OTTracer.java b/dd-java-agent/instrumentation/opentracing/api-0.31/src/main/java/datadog/trace/instrumentation/opentracing31/OTTracer.java index 7074a62b377..aa2ed041201 100644 --- a/dd-java-agent/instrumentation/opentracing/api-0.31/src/main/java/datadog/trace/instrumentation/opentracing31/OTTracer.java +++ b/dd-java-agent/instrumentation/opentracing/api-0.31/src/main/java/datadog/trace/instrumentation/opentracing31/OTTracer.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; @@ -53,7 +54,8 @@ public void inject(final SpanContext spanContext, final Format format, fi if (carrier instanceof TextMap) { final AgentSpanContext context = converter.toContext(spanContext); AgentSpan span = fromSpanContext(context); - defaultPropagator().inject(span, (TextMap) carrier, OTTextMapSetter.INSTANCE); + defaultPropagator() + .inject(Context.current().with(span), (TextMap) carrier, OTTextMapSetter.INSTANCE); } else { log.debug("Unsupported format for propagation - {}", format.getClass().getName()); } diff --git a/dd-java-agent/instrumentation/opentracing/api-0.32/src/main/java/datadog/trace/instrumentation/opentracing32/OTTracer.java b/dd-java-agent/instrumentation/opentracing/api-0.32/src/main/java/datadog/trace/instrumentation/opentracing32/OTTracer.java index 83094d05722..b8c738fa934 100644 --- a/dd-java-agent/instrumentation/opentracing/api-0.32/src/main/java/datadog/trace/instrumentation/opentracing32/OTTracer.java +++ b/dd-java-agent/instrumentation/opentracing/api-0.32/src/main/java/datadog/trace/instrumentation/opentracing32/OTTracer.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; @@ -64,7 +65,11 @@ public void inject(final SpanContext spanContext, final Format format, fi if (carrier instanceof TextMapInject) { final AgentSpanContext context = converter.toContext(spanContext); AgentSpan span = fromSpanContext(context); - defaultPropagator().inject(span, (TextMapInject) carrier, OTTextMapInjectSetter.INSTANCE); + defaultPropagator() + .inject( + Context.current().with(span), + (TextMapInject) carrier, + OTTextMapInjectSetter.INSTANCE); } else { log.debug("Unsupported format for propagation - {}", format.getClass().getName()); } diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java index 5c82e9bb020..512ba2ec918 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/DatadogWrapperHelper.java @@ -1,8 +1,8 @@ package datadog.trace.instrumentation.pekkohttp; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.pekkohttp.PekkoHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -11,12 +11,15 @@ public class DatadogWrapperHelper { public static AgentScope createSpan(final HttpRequest request) { - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(request); - final AgentSpan span = DECORATE.startSpan(request, extractedContext); + final Context extractedContext = DECORATE.extract(request); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(request, extractedSpanContext); DECORATE.afterStart(span); - DECORATE.onRequest(span, request, request, extractedContext); + DECORATE.onRequest(span, request, request, extractedSpanContext); - return activateSpan(span); + return (AgentScope) extractedContext.with(span).attach(); } public static void finishSpan(final AgentSpan span, final HttpResponse response) { diff --git a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java index 7e22ecd660d..a89b7a55473 100644 --- a/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java +++ b/dd-java-agent/instrumentation/pekko-http-1.0/src/main/java/datadog/trace/instrumentation/pekkohttp/PekkoHttpSingleRequestInstrumentation.java @@ -17,6 +17,7 @@ import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import net.bytebuddy.asm.Advice; import org.apache.pekko.http.scaladsl.HttpExt; import org.apache.pekko.http.scaladsl.model.HttpRequest; @@ -81,7 +82,11 @@ public static AgentScope methodEnter( if (request != null) { DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, headers); + defaultPropagator() + .inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), + request, + headers); // Request is immutable, so we have to assign new value once we update headers request = headers.getRequest(); } diff --git a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java index f21f313ce19..c3af8935a45 100644 --- a/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play-2.3/src/main/java/datadog/trace/instrumentation/play23/PlayAdvice.java @@ -7,6 +7,7 @@ import static datadog.trace.instrumentation.play23.PlayHttpServerDecorator.PLAY_REQUEST; import static datadog.trace.instrumentation.play23.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -21,18 +22,28 @@ public class PlayAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter(@Advice.Argument(0) final Request req) { final AgentSpan span; + final Context extractedContext; if (activeSpan() == null) { Headers headers = req.headers(); - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(headers); - span = DECORATE.startSpan(headers, extractedContext); + extractedContext = DECORATE.extract(headers); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(headers, extractedSpanContext); } else { // An upstream framework (e.g. akka-http, netty) has already started the span. // Do not extract the context. span = startSpan(PLAY_REQUEST); span.setMeasured(true); + extractedContext = null; } - final AgentScope scope = activateSpan(span); + final AgentScope scope; + if (extractedContext == null) { + scope = activateSpan(span); + } else { + scope = (AgentScope) extractedContext.with(span).attach(); + } DECORATE.afterStart(span); return scope; } diff --git a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java index 60b55add14f..9a27a51ba37 100644 --- a/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play-2.4/src/main/java/datadog/trace/instrumentation/play24/PlayAdvice.java @@ -7,6 +7,7 @@ import static datadog.trace.instrumentation.play24.PlayHttpServerDecorator.PLAY_REQUEST; import static datadog.trace.instrumentation.play24.PlayHttpServerDecorator.REPORT_HTTP_STATUS; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -21,6 +22,7 @@ public class PlayAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter(@Advice.Argument(value = 0, readOnly = false) Request req) { final AgentSpan span; + final Context extractedContext; // If we have already added a `play.request` span, then don't do it again if (req.tags().contains("_dd_HasPlayRequestSpan")) { @@ -29,16 +31,25 @@ public static AgentScope onEnter(@Advice.Argument(value = 0, readOnly = false) R if (activeSpan() == null) { final Headers headers = req.headers(); - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(headers); - span = DECORATE.startSpan(headers, extractedContext); + extractedContext = DECORATE.extract(headers); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(headers, extractedSpanContext); } else { // An upstream framework (e.g. akka-http, netty) has already started the span. // Do not extract the context. span = startSpan(PLAY_REQUEST); span.setMeasured(true); + extractedContext = null; } - final AgentScope scope = activateSpan(span); + final AgentScope scope; + if (extractedContext == null) { + scope = activateSpan(span); + } else { + scope = (AgentScope) extractedContext.with(span).attach(); + } DECORATE.afterStart(span); req = RequestHelper.withTag(req, "_dd_HasPlayRequestSpan", "true"); diff --git a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java index ecd5b903e19..f6bf0eff975 100644 --- a/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java +++ b/dd-java-agent/instrumentation/play-2.6/src/main/java/datadog/trace/instrumentation/play26/PlayAdvice.java @@ -6,6 +6,7 @@ import static datadog.trace.instrumentation.play26.PlayHttpServerDecorator.DECORATE; import static datadog.trace.instrumentation.play26.PlayHttpServerDecorator.PLAY_REQUEST; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -21,7 +22,7 @@ public class PlayAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope onEnter( @Advice.Argument(value = 0, readOnly = false) Request req, - @Advice.Local("extractedContext") AgentSpanContext.Extracted extractedContext) { + @Advice.Local("extractedContext") Context extractedContext) { final AgentSpan span; // If we have already added a `play.request` span, then don't do it again @@ -32,13 +33,23 @@ public static AgentScope onEnter( if (activeSpan() == null) { final Headers headers = req.headers(); extractedContext = DECORATE.extract(headers); - span = DECORATE.startSpan(headers, extractedContext); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(headers, extractedSpanContext); } else { // An upstream framework (e.g. akka-http, netty) has already started the span. // Do not extract the context. span = startSpan(PLAY_REQUEST); + extractedContext = null; + } + + final AgentScope scope; + if (extractedContext == null) { + scope = activateSpan(span); + } else { + scope = (AgentScope) extractedContext.with(span).attach(); } - final AgentScope scope = activateSpan(span); span.setMeasured(true); DECORATE.afterStart(span); diff --git a/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java index 556aa762cc6..4ed52cc5c9c 100644 --- a/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws-1/src/main/java/datadog/trace/instrumentation/playws1/PlayWSClientInstrumentation.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.PLAY_WS_REQUEST; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -31,7 +32,7 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); if (asyncHandler instanceof StreamedAsyncHandler) { asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); diff --git a/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java index 666b01c4bf4..400810121d9 100644 --- a/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws-2.1/src/main/java/datadog/trace/instrumentation/playws21/PlayWSClientInstrumentation.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.PLAY_WS_REQUEST; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -31,7 +32,7 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); if (asyncHandler instanceof StreamedAsyncHandler) { asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); diff --git a/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java b/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java index 668e7e1a47a..120b7412441 100644 --- a/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java +++ b/dd-java-agent/instrumentation/play-ws-2/src/main/java/datadog/trace/instrumentation/playws2/PlayWSClientInstrumentation.java @@ -8,6 +8,7 @@ import static datadog.trace.instrumentation.playws.PlayWSClientDecorator.PLAY_WS_REQUEST; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.datastreams.DataStreamsContext; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; @@ -31,7 +32,7 @@ public static AgentSpan methodEnter( DECORATE.afterStart(span); DECORATE.onRequest(span, request); DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), request, SETTER); if (asyncHandler instanceof StreamedAsyncHandler) { asyncHandler = new StreamedAsyncHandlerWrapper((StreamedAsyncHandler) asyncHandler, span); diff --git a/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitChannelInstrumentation.java b/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitChannelInstrumentation.java index 1c762039653..2c35e8552ab 100644 --- a/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitChannelInstrumentation.java +++ b/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitChannelInstrumentation.java @@ -37,6 +37,7 @@ import com.rabbitmq.client.Consumer; import com.rabbitmq.client.GetResponse; import com.rabbitmq.client.MessageProperties; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.Config; @@ -197,7 +198,7 @@ public static AgentScope setResourceNameAddHeaders( HAS_ROUTING_KEY_TAG, routingKey == null || routingKey.isEmpty() ? "false" : "true"); sortedTags.put(TYPE_TAG, "rabbitmq"); DataStreamsContext dsmContext = DataStreamsContext.fromTags(sortedTags); - defaultPropagator().inject(span.with(dsmContext), headers, SETTER); + defaultPropagator().inject(Context.current().with(span).with(dsmContext), headers, SETTER); props = new AMQP.BasicProperties( props.getContentType(), diff --git a/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitDecorator.java b/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitDecorator.java index c3d1797c4fc..8df6f3fe335 100644 --- a/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitDecorator.java +++ b/dd-java-agent/instrumentation/rabbitmq-amqp-2.7/src/main/java/datadog/trace/instrumentation/rabbitmq/amqp/RabbitDecorator.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.rabbitmq.amqp; import static datadog.trace.api.datastreams.DataStreamsContext.create; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.bootstrap.instrumentation.api.InstrumentationTags.AMQP_COMMAND; @@ -18,6 +17,8 @@ import com.rabbitmq.client.Command; import com.rabbitmq.client.Consumer; import com.rabbitmq.client.Envelope; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.api.Config; import datadog.trace.api.naming.SpanNaming; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -26,6 +27,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentTracer; import datadog.trace.bootstrap.instrumentation.api.ContextVisitors; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; import datadog.trace.bootstrap.instrumentation.decorator.MessagingClientDecorator; @@ -205,10 +207,19 @@ public static AgentScope startReceivingSpan( String queue) { final Map headers = propagate && null != properties ? properties.getHeaders() : null; - AgentSpanContext parentContext = - null != headers - ? extractContextAndGetSpanContext(headers, ContextVisitors.objectValuesMap()) - : null; + Context extractedContext = null; + AgentSpanContext parentContext = null; + if (null != headers) { + extractedContext = + Propagators.defaultPropagator() + .extract( + Java8BytecodeBridge.getCurrentContext(), + headers, + ContextVisitors.objectValuesMap()); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + parentContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + } // TODO: check dynamically bound queues - // https://github.com/DataDog/dd-trace-java/pull/2955#discussion_r677787875 @@ -260,7 +271,10 @@ public static AgentScope startReceivingSpan( } CONSUMER_DECORATE.afterStart(span); - AgentScope scope = activateSpan(span); + AgentScope scope = + extractedContext == null + ? activateSpan(span) + : (AgentScope) extractedContext.with(span).attach(); if (null != queueSpan) { queueSpan.finish(spanStartMicros); } diff --git a/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java b/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java index 68b808ab66f..ce6bba0f006 100644 --- a/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java +++ b/dd-java-agent/instrumentation/restlet-2.2/src/main/java/datadog/trace/instrumentation/restlet/RestletInstrumentation.java @@ -1,13 +1,13 @@ package datadog.trace.instrumentation.restlet; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.instrumentation.restlet.RestletDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; import com.sun.net.httpserver.HttpExchange; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -54,11 +54,14 @@ public String[] helperClassNames() { public static class RestletHandleAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope beginRequest(@Advice.Argument(0) final HttpExchange exchange) { - AgentSpanContext.Extracted context = DECORATE.extract(exchange); - AgentSpan span = DECORATE.startSpan(exchange, context); - AgentScope scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(exchange); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + AgentSpan span = DECORATE.startSpan(exchange, extractedSpanContext); + AgentScope scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, exchange, exchange, context); + DECORATE.onRequest(span, exchange, exchange, extractedSpanContext); DECORATE.onPeerConnection(span, exchange.getRemoteAddress()); return scope; diff --git a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java index 27797f8ba8f..5a2d12b320b 100644 --- a/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-2/src/main/java/datadog/trace/instrumentation/servlet2/Servlet2Advice.java @@ -1,9 +1,9 @@ package datadog.trace.instrumentation.servlet2; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.servlet2.Servlet2Decorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.ClassloaderConfigurationOverrides; import datadog.trace.api.Config; import datadog.trace.api.CorrelationIdentifier; @@ -52,11 +52,14 @@ public static boolean onEnter( InstrumentationContext.get(ServletResponse.class, Integer.class).put(response, 200); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(httpServletRequest); - final AgentSpan span = DECORATE.startSpan(httpServletRequest, extractedContext); - scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(httpServletRequest); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(httpServletRequest, extractedSpanContext); + scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedContext); + DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedSpanContext); httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span); httpServletRequest.setAttribute( diff --git a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java index 00e57312bce..4d6af7fd3a6 100644 --- a/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java +++ b/dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java @@ -7,6 +7,7 @@ import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_SPAN_ATTRIBUTE; import static datadog.trace.instrumentation.servlet3.Servlet3Decorator.DECORATE; +import datadog.context.Context; import datadog.trace.api.ClassloaderConfigurationOverrides; import datadog.trace.api.Config; import datadog.trace.api.CorrelationIdentifier; @@ -69,12 +70,15 @@ public static boolean onEnter( return false; } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(httpServletRequest); - final AgentSpan span = DECORATE.startSpan(httpServletRequest, extractedContext); - scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(httpServletRequest); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + final AgentSpan span = DECORATE.startSpan(httpServletRequest, extractedSpanContext); + scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedContext); + DECORATE.onRequest(span, httpServletRequest, httpServletRequest, extractedSpanContext); httpServletRequest.setAttribute(DD_SPAN_ATTRIBUTE, span); httpServletRequest.setAttribute( diff --git a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java index d2e3df70436..10d9d8fe8a8 100644 --- a/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java +++ b/dd-java-agent/instrumentation/servlet/src/main/java/datadog/trace/instrumentation/servlet/dispatcher/RequestDispatcherInstrumentation.java @@ -31,6 +31,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import java.util.Map; import javax.servlet.RequestDispatcher; import javax.servlet.ServletRequest; @@ -123,7 +124,9 @@ public static AgentScope start( // In case we lose context, inject trace into to the request. DataStreamsContext dsmContext = DataStreamsContext.fromTags(CLIENT_PATHWAY_EDGE_TAGS); - defaultPropagator().inject(span.with(dsmContext), request, SETTER); + defaultPropagator() + .inject( + Java8BytecodeBridge.getCurrentContext().with(span).with(dsmContext), request, SETTER); // temporarily replace from request to avoid spring resource name bubbling up: requestSpan = request.getAttribute(DD_SPAN_ATTRIBUTE); diff --git a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java index 53df15ffd76..b91311d99fc 100644 --- a/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java +++ b/dd-java-agent/instrumentation/spray-1.3/src/main/scala/datadog/trace/instrumentation/spray/SprayHttpServerRunSealedRouteAdvice.java @@ -5,6 +5,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; import static datadog.trace.instrumentation.spray.SprayHttpServerDecorator.DECORATE; +import datadog.context.Context; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; @@ -16,22 +17,29 @@ public class SprayHttpServerRunSealedRouteAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) public static AgentScope enter(@Advice.Argument(value = 1, readOnly = false) RequestContext ctx) { final AgentSpan span; - final AgentSpanContext.Extracted extractedContext; + Context extractedContext = null; + AgentSpanContext.Extracted extractedSpanContext = null; if (activeSpan() == null) { // Propagate context in case income request was going through several routes // TODO: Add test for it final HttpRequest request = ctx.request(); extractedContext = DECORATE.extract(request); - span = DECORATE.startSpan(request, extractedContext); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(request, extractedSpanContext); } else { extractedContext = null; span = startSpan(DECORATE.spanName()); } - final AgentScope scope = activateSpan(span); + final AgentScope scope = + extractedContext == null + ? activateSpan(span) + : (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - ctx = SprayHelper.wrapRequestContext(ctx, span, extractedContext); + ctx = SprayHelper.wrapRequestContext(ctx, span, extractedSpanContext); return scope; } diff --git a/dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageHandlerInstrumentation.java b/dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageHandlerInstrumentation.java index edc4bcfe589..1a90ced4c94 100644 --- a/dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/spring-messaging-4/src/main/java/datadog/trace/instrumentation/springmessaging/SpringMessageHandlerInstrumentation.java @@ -1,7 +1,6 @@ package datadog.trace.instrumentation.springmessaging; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; @@ -12,11 +11,14 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; +import datadog.context.propagation.Propagators; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext; +import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge; import net.bytebuddy.asm.Advice; import org.springframework.messaging.Message; import org.springframework.messaging.handler.invocation.InvocableHandlerMethod; @@ -59,17 +61,25 @@ public static AgentScope onEnter( @Advice.This InvocableHandlerMethod thiz, @Advice.Argument(0) Message message) { AgentSpanContext parentContext; AgentSpan parent = activeSpan(); + Context extractedContext = null; if (null != parent) { // prefer existing context, assume it was already extracted from this message parentContext = parent.context(); } else { // otherwise try to re-extract the message context to avoid disconnected trace - parentContext = extractContextAndGetSpanContext(message, GETTER); + extractedContext = + Propagators.defaultPropagator() + .extract(Java8BytecodeBridge.getCurrentContext(), message, GETTER); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + parentContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); } AgentSpan span = startSpan(SPRING_INBOUND, parentContext); DECORATE.afterStart(span); span.setResourceName(DECORATE.spanNameForMethod(thiz.getMethod())); - return activateSpan(span); + return extractedContext == null + ? activateSpan(span) + : (AgentScope) extractedContext.with(span).attach(); } @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) diff --git a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java index e860003559e..d2fae7a9dce 100644 --- a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseClientInstrumentation.java @@ -13,6 +13,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -86,7 +87,8 @@ public static AgentScope beginRequest( DECORATE.afterStart(span); // add trace id to client-side request before it gets submitted as an HttpRequest - defaultPropagator().inject(span, TargetContext.getRequest(connection), SETTER); + defaultPropagator() + .inject(Context.current().with(span), TargetContext.getRequest(connection), SETTER); // capture span to be finished by one of the various client response advices connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, span); diff --git a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java index d95d4668d1e..a251a9334e9 100644 --- a/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java +++ b/dd-java-agent/instrumentation/synapse-3/src/main/java/datadog/trace/instrumentation/synapse3/SynapseServerInstrumentation.java @@ -10,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -68,18 +69,20 @@ public static AgentScope beginRequest( // check incoming request for distributed trace ids HttpRequest request = connection.getHttpRequest(); - AgentSpanContext.Extracted extractedContext = DECORATE.extract(request); - + final Context extractedContext = DECORATE.extract(request); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); AgentSpan span; if (null != extractedContext) { - span = DECORATE.startSpan(request, extractedContext); + span = DECORATE.startSpan(request, extractedSpanContext); } else { span = startSpan(DECORATE.spanName()); span.setMeasured(true); } - AgentScope scope = activateSpan(span); + AgentScope scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, connection, request, extractedContext); + DECORATE.onRequest(span, connection, request, extractedSpanContext); // capture span to be finished by one of the various server response advices connection.getContext().setAttribute(SYNAPSE_SPAN_KEY, span); diff --git a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java index 3ae2690d000..e5df0886fa1 100644 --- a/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java +++ b/dd-java-agent/instrumentation/tomcat-5.5/src/main/java/datadog/trace/instrumentation/tomcat/TomcatServerInstrumentation.java @@ -12,6 +12,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.ExcludeFilterProvider; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; @@ -122,18 +123,22 @@ public static AgentScope onService(@Advice.Argument(0) org.apache.coyote.Request return activateSpan((AgentSpan) existingSpan); } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(req); - req.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedContext); + final Context extractedContext = DECORATE.extract(req); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + + req.setAttribute(DD_EXTRACTED_CONTEXT_ATTRIBUTE, extractedSpanContext); + + final AgentSpan span = DECORATE.startSpan(req, extractedSpanContext); - final AgentSpan span = DECORATE.startSpan(req, extractedContext); - final AgentScope scope = activateSpan(span); // This span is finished when Request.recycle() is called by RequestInstrumentation. DECORATE.afterStart(span); req.setAttribute(DD_SPAN_ATTRIBUTE, span); req.setAttribute(CorrelationIdentifier.getTraceIdKey(), GlobalTracer.get().getTraceId()); req.setAttribute(CorrelationIdentifier.getSpanIdKey(), GlobalTracer.get().getSpanId()); - return scope; + return (AgentScope) extractedContext.with(span).attach(); } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HandlerInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HandlerInstrumentation.java index 8192c5a9840..71de644b83d 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HandlerInstrumentation.java @@ -14,6 +14,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArguments; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.api.gateway.Flow.Action.RequestBlockingAction; @@ -93,11 +94,15 @@ public static void onEnter( return; } - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(exchange); - final AgentSpan span = DECORATE.startSpan(exchange, extractedContext).setMeasured(true); - scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(exchange); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + + final AgentSpan span = DECORATE.startSpan(exchange, extractedSpanContext).setMeasured(true); + scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, exchange, exchange, extractedContext); + DECORATE.onRequest(span, exchange, exchange, extractedSpanContext); exchange.putAttachment(DD_UNDERTOW_CONTINUATION, captureSpan(span)); diff --git a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java index cd205d013ef..39bf5ca6932 100644 --- a/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java +++ b/dd-java-agent/instrumentation/undertow/undertow-2.0/src/main/java/datadog/trace/instrumentation/undertow/HttpRequestParserInstrumentation.java @@ -2,13 +2,13 @@ import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass; import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; -import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan; import static datadog.trace.instrumentation.undertow.UndertowDecorator.DECORATE; import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import com.google.auto.service.AutoService; +import datadog.context.Context; import datadog.trace.agent.tooling.Instrumenter; import datadog.trace.agent.tooling.InstrumenterModule; import datadog.trace.bootstrap.instrumentation.api.AgentScope; @@ -75,11 +75,14 @@ public static void afterRequestParse( AgentScope scope = null; try { if (span == null) { - final AgentSpanContext.Extracted extractedContext = DECORATE.extract(exchange); - span = DECORATE.startSpan(exchange, extractedContext).setMeasured(true); - scope = activateSpan(span); + final Context extractedContext = DECORATE.extract(exchange); + final AgentSpan extractedSpan = AgentSpan.fromContext(extractedContext); + final AgentSpanContext.Extracted extractedSpanContext = + extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context(); + span = DECORATE.startSpan(exchange, extractedSpanContext).setMeasured(true); + scope = (AgentScope) extractedContext.with(span).attach(); DECORATE.afterStart(span); - DECORATE.onRequest(span, exchange, exchange, extractedContext); + DECORATE.onRequest(span, exchange, exchange, extractedSpanContext); } DECORATE.onError(span, throwable); // because we know that a http 400 will be thrown diff --git a/dd-trace-core/src/jmh/java/datadog/trace/core/baggage/BaggagePropagatorBenchmark.java b/dd-trace-core/src/jmh/java/datadog/trace/core/baggage/BaggagePropagatorBenchmark.java new file mode 100644 index 00000000000..f99f097cb6d --- /dev/null +++ b/dd-trace-core/src/jmh/java/datadog/trace/core/baggage/BaggagePropagatorBenchmark.java @@ -0,0 +1,108 @@ +package datadog.trace.core.baggage; + +import static java.util.concurrent.TimeUnit.MICROSECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; + +import datadog.context.Context; +import datadog.context.propagation.CarrierSetter; +import datadog.context.propagation.CarrierVisitor; +import datadog.trace.bootstrap.instrumentation.api.ContextVisitors; +import java.util.HashMap; +import java.util.Map; +import java.util.function.BiConsumer; +import java.util.logging.Logger; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.Blackhole; + +@State(Scope.Benchmark) +@Warmup(iterations = 4, time = 30, timeUnit = SECONDS) +@Measurement(iterations = 4, time = 30, timeUnit = SECONDS) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(MICROSECONDS) +@Fork(value = 1) +public class BaggagePropagatorBenchmark { + // Simulates different baggage header values + + private static final Logger logger = Logger.getLogger(BaggagePropagatorBenchmark.class.getName()); + + @Param("ignore this value") + String baggageHeaderValue; + + private BaggagePropagator propagator; + private Context context; + private Map carrier; + private CarrierSetter> setter; + private CarrierVisitor> visitor; + private boolean test; + + static class MapCarrierAccessor + implements CarrierSetter>, CarrierVisitor> { + @Override + public void set(Map carrier, String key, String value) { + if (carrier != null && key != null && value != null) { + carrier.put(key, value); + } + } + + @Override + public void forEachKeyValue(Map carrier, BiConsumer visitor) { + carrier.forEach(visitor); + } + } + + @Setup(Level.Trial) + public void setUp() { + propagator = new BaggagePropagator(true, true); + context = Context.root(); + + StringBuilder sb = new StringBuilder(8192); + for (int i = 0; i < 8192; i++) { + sb.append('a'); + } + final String aKey = "key1=" + sb.toString(); + // final String aKey = "key1=Amélie"; + + baggageHeaderValue = aKey; + carrier = new HashMap<>(); + carrier.put(BaggagePropagator.BAGGAGE_KEY, baggageHeaderValue); + visitor = ContextVisitors.stringValuesMap(); + setter = new MapCarrierAccessor(); + test = false; + } + + @Benchmark + public void propagate(Blackhole blackhole) { + Context extractedContext = propagator.extract(context, carrier, visitor); + carrier = new HashMap<>(); + // blackhole.consume(baggage.getW3cHeader()); + // blackhole.consume(extractedContext); + propagator.inject(extractedContext, carrier, setter); + // if (!test){ + // Baggage baggage = Baggage.fromContext(extractedContext); + // System.out.println("w3c header: " + baggage.getW3cHeader()); + // Map baggageMap = baggage.asMap(); + // for(String key : baggageMap.keySet()){ + // System.out.println(key + ":" + baggageMap.get(key)); + // } + // System.out.println("Carrier.size: " + carrier.size()); + // for (String key : carrier.keySet()){ + // System.out.println("Carrier.keyset: " + key); + // System.out.println("Carrier: " + carrier.get(key)); + // } + // test = true; + // } + + // blackhole.consume(baggage); + } +} diff --git a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java index fcf844baf0a..aa5ff7e803b 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/datastreams/DefaultDataStreamsMonitoring.java @@ -15,6 +15,7 @@ import datadog.communication.ddagent.DDAgentFeaturesDiscovery; import datadog.communication.ddagent.SharedCommunicationObjects; +import datadog.context.Context; import datadog.context.propagation.Propagator; import datadog.trace.api.Config; import datadog.trace.api.TraceConfig; @@ -294,7 +295,9 @@ public void setProduceCheckpoint( DataStreamsContext dsmContext = fromTags(sortedTags); this.propagator.inject( - span.with(dsmContext), carrier, DataStreamsContextCarrierAdapter.INSTANCE); + Context.current().with(span).with(dsmContext), + carrier, + DataStreamsContextCarrierAdapter.INSTANCE); } @Override diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java index 9835d2e1c03..19b3df2ce53 100644 --- a/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java +++ b/dd-trace-core/src/main/java/datadog/trace/core/propagation/HttpCodec.java @@ -221,7 +221,6 @@ public TagContext extract( TagContext partialContext = null; // Extract and cache all headers in advance ExtractionCache extractionCache = new ExtractionCache<>(carrier, getter); - for (final Extractor extractor : this.extractors) { TagContext extracted = extractor.extract(extractionCache, extractionCache); // Check if context is valid @@ -260,7 +259,6 @@ else if (extracted != null && partialContext == null) { partialContext = extracted; } } - if (context != null) { log.debug("Extract complete context {}", context); return context; diff --git a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java index fdea6629942..7ea71644eec 100644 --- a/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java +++ b/dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java @@ -4,6 +4,7 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext; import static datadog.trace.bootstrap.instrumentation.api.AgentSpan.fromSpanContext; +import datadog.context.Context; import datadog.context.propagation.CarrierSetter; import datadog.trace.api.Config; import datadog.trace.api.DDTags; @@ -486,7 +487,8 @@ public void inject(final SpanContext spanContext, final Format format, fi if (carrier instanceof TextMap) { final AgentSpanContext context = converter.toContext(spanContext); AgentSpan span = fromSpanContext(context); - defaultPropagator().inject(span, (TextMap) carrier, TextMapSetter.INSTANCE); + defaultPropagator() + .inject(Context.current().with(span), (TextMap) carrier, TextMapSetter.INSTANCE); } else { log.debug("Unsupported format for propagation - {}", format.getClass().getName()); } diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java index d6526b29716..c02fc18a85e 100644 --- a/internal-api/src/main/java/datadog/trace/api/Config.java +++ b/internal-api/src/main/java/datadog/trace/api/Config.java @@ -2369,11 +2369,13 @@ public boolean isTracePropagationExtractFirst() { } public boolean isBaggageExtract() { - return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE); + return tracePropagationStylesToExtract.contains(TracePropagationStyle.BAGGAGE) + && tracePropagationBehaviorExtract != TracePropagationBehaviorExtract.IGNORE; } public boolean isBaggageInject() { - return tracePropagationStylesToInject.contains(TracePropagationStyle.BAGGAGE); + return tracePropagationStylesToInject.contains(TracePropagationStyle.BAGGAGE) + && tracePropagationBehaviorExtract != TracePropagationBehaviorExtract.IGNORE; } public boolean isBaggagePropagationEnabled() {