Skip to content
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

Jump to bottom

Inject trace context into AWS Step Functions input #7585

Open
wants to merge 16 commits into
base: master
Choose a base branch
from
+333 −0
Open

Inject trace context into AWS Step Functions input #7585

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
7f8721d
inject step functions trace context
DylanLovesCoffee Sep 5, 2024
a974597
Merge branch 'master' into dylan/sfn-trace-ctx
DylanLovesCoffee Sep 9, 2024
7ecfa8c
remove previous logic
DylanLovesCoffee Sep 13, 2024
fa36832
aws sfn instrumentation
DylanLovesCoffee Sep 16, 2024
22c652b
Merge branch 'master' into dylan/sfn-trace-ctx
DylanLovesCoffee Sep 16, 2024
eb642b0
muzzle
DylanLovesCoffee Sep 16, 2024
0778f3d
single char
DylanLovesCoffee Sep 17, 2024
f9d0eea
Merge branch 'master' into dylan/sfn-trace-ctx
DylanLovesCoffee Oct 8, 2024
421856b
update muzzle
DylanLovesCoffee Oct 8, 2024
895341a
refactor input attr injector
DylanLovesCoffee Oct 8, 2024
9d76dec
refactor interceptor
DylanLovesCoffee Oct 9, 2024
aba1119
cleanup test
DylanLovesCoffee Oct 9, 2024
e56f2d3
Merge branch 'master' into dylan/sfn-trace-ctx
DylanLovesCoffee Oct 23, 2024
3039780
Merge branch 'master' into dylan/sfn-trace-ctx
DylanLovesCoffee Oct 23, 2024
370df6d
use JsonBuffer
DylanLovesCoffee Oct 24, 2024
a644a8b
catch exceptions
DylanLovesCoffee Oct 24, 2024
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
36 changes: 36 additions & 0 deletions ...ava-sdk-2.2/src/main/java/datadog/trace/instrumentation/aws/v2/AwsSdkClientDecorator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
import static datadog.trace.core.datastreams.TagsProcessor.TOPIC_TAG;
import static datadog.trace.core.datastreams.TagsProcessor.TYPE_TAG;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import datadog.trace.api.Config;
import datadog.trace.api.DDTags;
import datadog.trace.api.cache.DDCache;
Expand All @@ -23,6 +26,8 @@
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
import datadog.trace.bootstrap.instrumentation.decorator.HttpClientDecorator;
import datadog.trace.core.datastreams.TagsProcessor;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.net.URI;
import java.time.Instant;
import java.util.Collections;
Expand Down Expand Up @@ -398,6 +403,37 @@ public AgentSpan onSdkResponse(
return span;
}

public static void injectTraceToStepFunctionInput(SdkRequest request, AgentSpan span) {
Class<?> clazz = request.getClass();
if ((clazz.getSimpleName().equals("StartExecutionRequest"))
Copy link
Collaborator

Choose a reason for hiding this comment

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

May instanceof be used here?

Copy link
Contributor

Choose a reason for hiding this comment

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

Main issue is that not all AWS types will be available on the classpath - it depends which parts of the SDK the user deploys. So you can't use types directly like with instanceof.

However, I'd prefer to see this use the fully qualified class name, because there is a cost to calling getSimpleName and we know the package(s) ahead of time.

|| clazz.getSimpleName().equals("StartSyncExecutionRequest")) {
try {
Method method = clazz.getMethod("input");
String input = (String) method.invoke(request);

ObjectMapper mapper = new ObjectMapper();
DylanLovesCoffee marked this conversation as resolved.
Show resolved Hide resolved
JsonNode inputJsonNode = mapper.readTree(input);
if (inputJsonNode.isObject()) {
ObjectNode traceContext = mapper.createObjectNode();
traceContext.put("x-datadog-trace-id", span.getTraceId().toString());
traceContext.put("x-datadog-parent-id", String.valueOf(span.getSpanId()));
traceContext.put("x-datadog-sampling-priority", span.getSamplingPriority().toString());
traceContext.put("x-datadog-tags", mapper.writeValueAsString(span.getTags()));

ObjectNode ddInputNode = (ObjectNode) inputJsonNode;
ddInputNode.put("_datadog", traceContext);
}

Field inputField = clazz.getDeclaredField("input");
Copy link
Collaborator

Choose a reason for hiding this comment

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

This is pretty inefficient and is discouraged. Instead you might obtain (once a MethodHandle) and unreflect it in order to avoid paying that extra cost

Copy link
Contributor

Choose a reason for hiding this comment

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

Agreed - I would prefer this to be moved to a step-function specific instrumentation module, like the SNS specific instrumentation module(s) added in #6908 - you can still intercept the request and modify it as appropriate, but you'll have access to the right types and can directly target the appropriate part of the AWS SDK

Trying to do this in the generic AWS SDK instrumentation will make this code convoluted and hard to maintain.

I'd also recommend avoiding adding a dependency to com.fasterxml.jackson.databind in the instrumentation, because that might not be accessible in certain deployments - ideally take the same approach as SnsInterceptor and just build up the simple string (the JSON isn't that complex)

In fact it's probably a good idea to refactor that _datadog JSON building logic out so it can be shared across all the different AWS SDK interceptors that want to add _datadog to their requests...

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Ah thanks! That helps a ton. I'll give that a shot to try to clean things up + add test coverage

inputField.setAccessible(true);
inputField.set(request, mapper.writeValueAsString(inputJsonNode));

} catch (Throwable e) {
// Failed to inject trace context
}
}
}

@Override
protected String[] instrumentationNames() {
return new String[] {"aws-sdk"};
Expand Down
12 changes: 12 additions & 0 deletions ...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 @@ -63,6 +63,18 @@ public void afterMarshalling(
}
}

@Override
public SdkRequest modifyRequest(
Context.ModifyRequest context, ExecutionAttributes executionAttributes) {
SdkRequest sdkRequest = context.request();
final AgentSpan span = executionAttributes.getAttribute(SPAN_ATTRIBUTE);
if (span != null) {
DECORATE.injectTraceToStepFunctionInput(sdkRequest, span);
}

return sdkRequest;
}

@Override
public SdkHttpRequest modifyHttpRequest(
Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
Expand Down