[serverless] Inject trace context into SQS/SNS/EventBridge

[nhulston]

What does this PR do?

This PR injects trace context into:

• SQS SendMessage and SendMessageBatch messages
• SNS Publish and PublishBatch messages
• EventBridge PutEvents events

This allows the Lambda extension to use this trace context to create inferred spans and combine these spans into a single trace. (See PR 1 and PR 2 for these changes)

Motivation

SNS and SQS are already supported in Java and .NET. However, this Go tracer did not have support for SQS, SNS, or EventBridge.

I am also working on a similar PR in dd-trace-java and dd-trace-dotnet to add EventBridge trace context injection. This PR differs because it adds support for SQS, SNS, and EventBridge.

Notes for reviewers

• I organized the PR commit by commit, so you can review that way if you want.
• I tried to structure the code in a way that makes it easy to (1) add "monkey patching" for new services, and (2) add "monkey patching" to new operations to these existing services in the future. Hence, in eventbridge.go, we have a switch statement with a single case right now. But this will make it easier for future operations to be added.

Testing

I added unit tests, which can be ran with the commands:

go test ./contrib/aws/aws-sdk-go-v2/aws -v
go test ./contrib/aws/internal/sqs -v
go test ./contrib/aws/internal/sns -v
go test ./contrib/aws/internal/eventbridge -v 

I also did manual tests. Here are what the flame graphs look like after these changes:

System tests covering SNS/SQS: DataDog/system-tests#3204

Reviewer's Checklist

• Changed code has unit tests for its functionality at or near 100% coverage.
• System-Tests covering this feature have been added and enabled with the va.b.c-dev version tag.
• There is a benchmark for any new code, or changes to existing code.
• If this interacts with the agent in a new way, a system test has been added.
• Add an appropriate team label so this PR gets put in the right place for the release notes.
• Non-trivial go.mod changes, e.g. adding new modules, are reviewed by @DataDog/dd-trace-go-guild.

Unsure? Have a question? Request a review!

[pr-commenter]

Benchmarks

Benchmark execution time: 2024-10-16 21:48:16

Comparing candidate commit edb0584 in PR branch nicholas.hulston/sqs-sns-eventbridge-inject-trace-context with baseline commit 51abc7b in branch main.

Found 0 performance improvements and 0 performance regressions! Performance is the same for 58 metrics, 1 unstable metrics.

[rarguelloF]

some minor changes to make sure we don't expose packages if its not necessary 🙏

[purple4reina]

Apologies if it seems like I am nit picking you a lot here. Let me explain.

Customers care a lot about the performance of their apps. Since the instrumentation you are writing here gets run for every message send performed, it's important that we keep the overhead as low as possible.

Since our customers are expecting low overhead, this can sometimes lead to code that is more complex than otherwise necessary.

[rarguelloF]

Overall LGTM, just echoing @purple4reina's comments about extracting the repeated logic out of the for loop.

[purple4reina]

nit.

I think you could simplify this out even further. Since both handlePublish and handlePublishBatch call getTraceContext, you could make the getTraceContext call just once, outside of the handle functions and then pass the context objects into the handle functions.

This makes no functional change, but it would DRY up the code a bit. That would be an improvement, since if we need to update the getTraceContext call in the future, we'd right now need to update it in two locations. Making the getTraceContext call just once would mean updating code in a single location.

[nhulston]

Thanks, this is an easy fix and I'm open to doing it. But I'm trying to structure the EnrichOperation func in a way that allows other teams to modify other operations in SQS/SNS. I think there could be other operations where they won't need the trace context, so it might make more sense to keep the getTraceContext() calls in the handle funcs. For example, this in the java tracer:
https://github.com/DataDog/dd-trace-java/blob/master/dd-java-agent/instrumentation/aws-java-sqs-2.0/src/main/java/datadog/trace/instrumentation/aws/v2/sqs/SqsInterceptor.java#L84-L90

What do you think?

[purple4reina]

I think you do in fact have an opportunity to pull a lot of work in eventbridge out of the for loop so it is only performed a single time.

Additionally, you did a great job with the java and dotnet tracers to inspect json payloads as strings and adding trace context manually. I think you can do the same thing here too.

[purple4reina]

Can you add a test that sets the trace prop style to none? I'm concerned in that case we'd end up with invalid json, like {,"start":123}.

I think the env var you want is DD_TRACE_PROPAGATION_STYLE=none.

[purple4reina]

And actually, could probably avoid any such issues if you manually add the start time key to the carrier.

carrier[startTimeKey] = strconv.Itoa(startTimeMillis)
carrierJSON, err := json.Marshal(carrier)
etc...

This way you're always guaranteed at least one item in the map.

[nhulston]

Good idea

[rarguelloF]

LGTM!

[purple4reina]

This is dope