Skip to content

Baggage E2E "working" #8699

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mhlidd wants to merge 18 commits into from
Closed

Baggage E2E "working" #8699

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
56802d6
saving work
mhlidd Apr 7, 2025
ff0a3be
accomadating baggage
mhlidd Apr 9, 2025
2ea2788
cleanup comments
mhlidd Apr 15, 2025
210aebe
saving context value
mhlidd Apr 15, 2025
2fe727e
nit changes
mhlidd Apr 15, 2025
6bc47ce
supporting serverside-instrumentations
mhlidd Apr 16, 2025
c7ba3d4
adding support for generic context in instrumentation
mhlidd Apr 16, 2025
8ccb7b5
removing logs
mhlidd Apr 16, 2025
cb6a6f9
adding support for injection of generic context
mhlidd Apr 16, 2025
4d5ed09
fixing errors
mhlidd Apr 16, 2025
2dab314
visibility
mhlidd Apr 21, 2025
c78a5c8
spotless
mhlidd Apr 24, 2025
0487371
benchmark visibility
mhlidd Apr 25, 2025
ea19707
spotless
mhlidd Apr 30, 2025
7b258fc
using Bytecode bridge for instrumentations
mhlidd May 1, 2025
882611b
reverting activateNext changes
mhlidd May 2, 2025
f3e8a5f
attempt to use BytecodeBridge in extract
mhlidd May 2, 2025
6d2a261
updating import for rabbit decorator
mhlidd May 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions .../src/main/java/datadog/trace/bootstrap/instrumentation/decorator/HttpServerDecorator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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<REQUEST_CARRIER> 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. */
Expand Down
7 changes: 6 additions & 1 deletion ...ain/java/datadog/trace/instrumentation/akkahttp/AkkaHttpSingleRequestInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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();
}
Expand Down
13 changes: 8 additions & 5 deletions ...-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/DatadogWrapperHelper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
Copy link
Contributor

@PerfectSlayer PerfectSlayer May 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It feels like HttpServerDecorator (DECORATE here) should migrate from AgentSpan.Context to Context.
This should greatly simplify the related instrumentation calls.


return activateSpan(span);
return (AgentScope) extractedContext.with(span).attach();
}

public static void finishSpan(final AgentSpan span, final HttpResponse response) {
Expand Down
9 changes: 7 additions & 2 deletions ...p-10.6/src/main/java11/datadog/trace/instrumentation/akkahttp106/SingleRequestAdvice.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand Down
4 changes: 2 additions & 2 deletions ...n/java/datadog/trace/instrumentation/apachehttpasyncclient/DelegatingRequestProducer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}

Expand Down
3 changes: 2 additions & 1 deletion ...pclient-4/src/main/java/datadog/trace/instrumentation/apachehttpclient/HelperMethods.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
3 changes: 2 additions & 1 deletion ...c/main/java/datadog/trace/instrumentation/apachehttpclient5/DelegatingRequestChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
}
}
3 changes: 2 additions & 1 deletion ...client-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/HelperMethods.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
4 changes: 3 additions & 1 deletion ...java/datadog/trace/instrumentation/armeria/grpc/client/ClientCallImplInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -137,7 +138,8 @@ public static <T> 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);
}
}
Expand Down
18 changes: 13 additions & 5 deletions ...main/java/datadog/trace/instrumentation/armeria/grpc/server/TracingServerInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -62,13 +64,19 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> 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()
Expand All @@ -85,7 +93,7 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
DECORATE.onCall(span, call);

final ServerCall.Listener<ReqT> result;
try (AgentScope scope = activateSpan(span)) {
try (AgentScope scope = (AgentScope) extractedContext.with(span).attach()) {
Copy link

Copilot AI May 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] Casting the result of attach() to AgentScope may be risky if the context API changes. Consider using a dedicated method for obtaining an AgentScope directly to avoid potential ClassCastExceptions.

Suggested change
try (AgentScope scope = (AgentScope) extractedContext.with(span).attach()) {
try (AgentScope scope = getAgentScopeFromContext(extractedContext, span)) {

Copilot uses AI. Check for mistakes.
// Wrap the server call so that we can decorate the span
// with the resulting status
final TracingServerCall<ReqT, RespT> tracingServerCall = new TracingServerCall<>(span, call);
Expand Down
2 changes: 1 addition & 1 deletion ...rc/main/java/datadog/trace/instrumentation/aws/v2/eventbridge/EventBridgeInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
4 changes: 3 additions & 1 deletion ...-sdk-1.11.0/src/main/java/datadog/trace/instrumentation/aws/v0/TracingRequestHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand Down
3 changes: 2 additions & 1 deletion ...k-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/TracingExecutionInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
2 changes: 1 addition & 1 deletion ...s-java-sns-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sns/SnsInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion ...s-java-sns-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sns/SnsInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
3 changes: 2 additions & 1 deletion ...s-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/SqsInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Copy link
Contributor

@PerfectSlayer PerfectSlayer May 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't SqsInterceptor it's own class and not injected?
Did you observe bytecode issue when using direct calls to the context API?

Copy link
Contributor Author

@mhlidd mhlidd May 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I did? I don't remember tbh but I can locally revert and see what happens

}

private AgentSpan newSpan(AmazonWebServiceRequest request) {
Expand Down
18 changes: 13 additions & 5 deletions ...-java-sqs-1.0/src/main/java/datadog/trace/instrumentation/aws/v1/sqs/TracingIterator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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);
Copy link
Contributor

@PerfectSlayer PerfectSlayer May 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

extractContextAndGetSpanContext originally used a root context as default context value.
You might end up having weird behavior re-using current context for extraction.
I haven't check all calls but you would better have a quick pass to check this kind of change as it might not show up during testing... 😞

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);
Expand Down
3 changes: 2 additions & 1 deletion ...s-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Copy link
Contributor

@PerfectSlayer PerfectSlayer May 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Span seemed to be retrieved from a carrier. Meaning the current context might not apply here.
Need to check / investigate the expected behavior.

}

private LinkedHashMap<String, String> getTags(String queueUrl) {
Expand Down
Loading