X-Ray trace header parsing and Lambda tracing propagation #11708

joe4dev · 2024-10-17T15:55:53Z

Motivation

We do not parse incoming X-Ray trace headers. Sample trace header (with Lambda lineage):

X-Amzn-Trace-Id: Root=1-5759e988-bd862e3fe1be46a994272793;Sampled=1;Lineage=a87bd80c:1|68fd508a:5|c512fbe3:2

Changes

Add trace context parser to LocalStack handler chain with an implementation for X-Ray trace header parsing
Add a trace_context field to the LocalStack RequestContext
Propage trace_context through the Lambda service (including encoding and decoding hacks for async invokes)
Add aws-validated test case test_xray_trace_propagation
Add TraceHeader helper class implementation based on the open source AWS Python X-Ray SDK (Apache2.0 license-attributed): https://github.com/aws/aws-xray-sdk-python/blob/master/aws_xray_sdk/core/models/trace_header.py
Add TraceId helper class implementation based on the open source AWS Python X-Ray SDK (Apache2.0 license-attributed): https://github.com/aws/aws-xray-sdk-python/blob/master/aws_xray_sdk/core/models/traceid.py

Testing

ext integration tests:

Usage

Use context.trace_context within any provider to access the parsed X-Ray trace header. Check out the TraceHeader class for more details.

Create new TraceHeader:

TraceHeader()  # all fields are None !

h1 = TraceHeader().ensure_root_exists()
h1.to_header_str() == "Root=1-67126ec5-45ded1c6463bac5ef4be1d69"

h2 = TraceHeader(root="1-3152b799-8954dae64eda91bc9a23a7e8", parent="7fa8c0f79203be72", sampled=1)
h2.to_header_str() == "Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1"

github-actions · 2024-10-17T16:04:35Z

S3 Image Test Results (AMD64 / ARM64)

2 files 2 suites 3m 34s ⏱️
423 tests 369 ✅ 54 💤 0 ❌
846 runs 738 ✅ 108 💤 0 ❌

Results for commit 006b34b.

♻️ This comment has been updated with latest results.

joe4dev

Happy to hear your feedback about X-Ray trace propagation (presented in today's demo) :)

joe4dev · 2024-10-17T15:58:15Z

localstack-core/localstack/aws/app.py

@@ -45,6 +45,7 @@ def __init__(self, service_manager: ServiceManager = None) -> None:
                handlers.add_region_from_header,
                handlers.rewrite_region,
                handlers.add_account_id,
+                handlers.parse_trace_context,


@bentsku @thrau What do you think about the ordering in the handler chain?

I think this is good! At first, I thought it'd be nicer to have it before the edge routes, because we could also make use of this in API Gateway execute-api routes for example. However, we do not pass the RequestContext to the edge routes at the moment. (@thrau laments this design decision still today 😄) so there is no point in making it happen early for now. We can always move it later on if anything changes! But for now, this looks to me 👌

I see. Feel free to change later :)

localstack-core/localstack/aws/api/core.py

localstack-core/localstack/aws/handlers/tracing.py

joe4dev · 2024-10-17T16:02:31Z

localstack-core/localstack/services/lambda_/invocation/event_manager.py

@@ -57,6 +58,11 @@ class SQSInvocation:
    exception_retries: int = 0

    def encode(self) -> str:
+        # Encode TraceHeader
+        aws_trace_header_str = TraceHeader(


@dfangl not super happy about this encoding/decoding hack. Feel free to share your suggestion.

This doesn't go well if aws_trace_header is not of type TraceHeader

localstack-core/localstack/utils/xray/trace_header.py

github-actions · 2024-10-18T09:17:09Z

LocalStack Community integration with Pro

2 files ±0 2 suites ±0 1h 42m 12s ⏱️ -12s
3 510 tests +1 3 097 ✅ +1 413 💤 ±0 0 ❌ ±0
3 512 runs +1 3 097 ✅ +1 415 💤 ±0 0 ❌ ±0

Results for commit 006b34b. ± Comparison against base commit 4a775c4.

♻️ This comment has been updated with latest results.

This fixes the warning "Subsegment ## handler discarded due to Lambda worker still initializing" Fixes the regression for the ext test: `tests.aws.services.lambda_.test_lambda_xray.TestLambdaXrayIntegration.test_basic_xray_integration`

bentsku

Awesome, this looks great! 🚀 I've only reviewed the parts that are not related to lambda, as I'm lacking a bit of context and knowledge there.

The handler looks great, I only have a few nits in the code taken from the AWS Python X-Ray SDK, as I think we could simplify it a bit by using new available Python standard lib utils and f-strings. Great work! 🚀

bentsku · 2024-10-28T10:17:23Z

localstack-core/localstack/aws/app.py

@@ -45,6 +45,7 @@ def __init__(self, service_manager: ServiceManager = None) -> None:
                handlers.add_region_from_header,
                handlers.rewrite_region,
                handlers.add_account_id,
+                handlers.parse_trace_context,


I think this is good! At first, I thought it'd be nicer to have it before the edge routes, because we could also make use of this in API Gateway execute-api routes for example. However, we do not pass the RequestContext to the edge routes at the moment. (@thrau laments this design decision still today 😄) so there is no point in making it happen early for now. We can always move it later on if anything changes! But for now, this looks to me 👌

bentsku · 2024-10-28T10:21:34Z

localstack-core/localstack/utils/xray/trace_header.py

+    Generate a random 16-digit hex str.
+    This is used for generating segment/subsegment id.
+    """
+    return binascii.b2a_hex(os.urandom(8)).decode("utf-8")


I know this is taken from AWS code, but I think we could use some Python utils to make it a bit nicer.

from secrets import token_hex

Suggested change

return binascii.b2a_hex(os.urandom(8)).decode("utf-8")

return token_hex(8)

nice suggestions ✨ This looks much cleaner 🤩

I fully agree on a technical level. For licensing and maintenance reasons, I opted to minimize changes to the original AWS code here (it is easier to match the authoritative AWS implementation regarding header parsing). Do you think it's worth deviating here?

I suppose not, we could also keep the code as it is, if it makes it easier to port new changes as they come (not sure how often this happens?).

I'm also crying every time I see %s%s%s%s%s 😭 . Nevertheless, I tend to prefer keeping the original for full parity (it's easier to relate to the AWS repo). Yeah, I don't expect direct changes here, maybe indirect ones for other X-Ray parts from the same source repo.

bentsku · 2024-10-28T10:23:04Z

localstack-core/localstack/utils/xray/traceid.py

+        Generate a random trace id.
+        """
+        self.start_time = int(time.time())
+        self.__number = binascii.b2a_hex(os.urandom(12)).decode("utf-8")


nit: not sure about the double __ here, should we just use a single underscore?
And we can also use token_hex instead of binascii and os.random

bentsku · 2024-10-28T10:25:18Z

localstack-core/localstack/utils/xray/traceid.py

+        return "%s%s%s%s%s" % (
+            TraceId.VERSION,
+            TraceId.DELIMITER,
+            format(self.start_time, "x"),
+            TraceId.DELIMITER,
+            self.__number,
+        )


We should probably use f-strings here, so we could even avoid the call to format:

Suggested change

return "%s%s%s%s%s" % (

TraceId.VERSION,

TraceId.DELIMITER,

format(self.start_time, "x"),

TraceId.DELIMITER,

self.__number,

)

return f"{TraceId.VERSION}{TraceId.DELIMITER}{self.start_time:x}{TraceId.DELIMITER}{self._number}"

bentsku · 2024-10-28T10:32:12Z

localstack-core/localstack/aws/handlers/tracing.py

This looks great 👌

dominikschubert

LGTM, most points have already been discussed and the current state seems very well documented and thought through. There's a few things we still need to tackle in the future like the sampling strategies but I actually quite like that we leave that out for now and can delegate this in the future.

Great work on re-using existing components. It makes a lot of sense to stay close to the official tooling here to make sure our behavior isn't diverging from AWS.

In terms of adjusting the internalized code from the xray lib, I'd generally prefer to stay closer to the original code if possible but as long as it doesn't impede porting new changes too much, it's of course also ok to modernize it a bit.

dominikschubert · 2024-10-29T13:33:01Z

localstack-core/localstack/services/lambda_/invocation/execution_environment.py

+        )
+        # The Lambda RIE requires a full tracing header including Root, Parent, and Samples. Otherwise, tracing fails
+        # with the warning "Subsegment ## handler discarded due to Lambda worker still initializing"
+        aws_trace_header.ensure_sampled_exists()


so if I understand that right, for lambda it's always Sampled=1 while other services will by default get a Sampled=0? 🤔

no, the sampled field is missing altogether (i.e., not present rather than sampled=0) for other services if no incoming tracing header is present with the sampled field.

That's not ideal, but this depends on the (service-dependent) sampling strategies and other services. We're lacking parity tests here, so idk what's the best behavior. Just leaving it open for now ...

dominikschubert · 2024-10-29T13:34:57Z

localstack-core/localstack/services/lambda_/invocation/execution_environment.py

+        # TODO: replace this random parent id with actual parent segment created within the Lambda provider using X-Ray
+        aws_trace_header.ensure_parent_exists()
+        # TODO: test and implement Active and PassThrough tracing and sampling decisions.
+        # TODO: implement Lambda lineage: https://docs.aws.amazon.com/lambda/latest/dg/invocation-recursion.html


Thanks for leaving the TODOs here. Those seem like really interesting tasks for engineers getting started at LocalStack 😉

dominikschubert · 2024-10-29T13:36:06Z

localstack-core/localstack/services/lambda_/invocation/lambda_service.py

        :param payload: Invocation payload
        :return: The invocation result
        """
+        # NOTE: consider making the trace_context mandatory once we update all usages (should be easier after v4.0)


agree, wanted to point out the same here 🙌

dominikschubert · 2024-10-29T13:42:07Z

tests/aws/services/lambda_/test_lambda_integration_xray.py

+    root_trace_id = "1-3152b799-8954dae64eda91bc9a23a7e8"
+    xray_trace_header = TraceHeader(root=root_trace_id, parent="7fa8c0f79203be72", sampled=1)


Does this present any issues when testing repeatedly against AWS with the same trace id / parent?

nope, works fine.
AWS does not perform any magic validation trying to re-construct the implicit start timestamp encoded here :)

PS: I fixed the repeated AWS validation issues in other X-Ray tests in this ext PR (awaiting review, low prio): https://github.com/localstack/localstack-ext/pull/3596

joe4dev added the semver: minor Non-breaking changes which can be included in minor releases, but not in patch releases label Oct 17, 2024

joe4dev added this to the 4.0 milestone Oct 17, 2024

joe4dev self-assigned this Oct 17, 2024

joe4dev requested review from dominikschubert, dfangl, gregfurman and thrau as code owners October 17, 2024 15:55

joe4dev commented Oct 17, 2024

View reviewed changes

joe4dev marked this pull request as draft October 18, 2024 07:47

joe4dev force-pushed the feature/lambda-xray-trace-propagation branch from d35c6d4 to 8fbd77c Compare October 18, 2024 14:27

joe4dev mentioned this pull request Oct 18, 2024

bug: Lambda invocations don't support X-Ray Passive Tracing (via x-amzn-trace-id header) #11336

Closed

1 task

joe4dev marked this pull request as ready for review October 18, 2024 17:06

joe4dev force-pushed the feature/lambda-xray-trace-propagation branch from d68d520 to 5a9693f Compare October 21, 2024 08:54

joe4dev added 9 commits October 22, 2024 10:55

Add xray trace context propagation

7d44384

Add aws-validated XRay trace propagation Lambda test

3acfd6b

Improve None value handling

455136c

Clarify license information

9c9b835

Make trace_context optional to quickfix regressions

1005e8e

Fix async invoke trace header encoding regression

5d8367a

Generate root trace id if it doesn't exist

4087d43

Generate root trace id for async/event invokes without trace header

9d361d4

joe4dev force-pushed the feature/lambda-xray-trace-propagation branch from 0a44ee9 to 006b34b Compare October 22, 2024 08:55

joe4dev requested a review from bentsku October 25, 2024 14:53

bentsku approved these changes Oct 28, 2024

View reviewed changes

dominikschubert approved these changes Oct 29, 2024

View reviewed changes

joe4dev merged commit e675014 into master Oct 29, 2024
47 checks passed

joe4dev deleted the feature/lambda-xray-trace-propagation branch October 29, 2024 19:52

bentsku mentioned this pull request Apr 4, 2025

feat: propagate x-ray trace id to event bridge targets #12481

Merged

	return binascii.b2a_hex(os.urandom(8)).decode("utf-8")
	return token_hex(8)

		root_trace_id = "1-3152b799-8954dae64eda91bc9a23a7e8"
		xray_trace_header = TraceHeader(root=root_trace_id, parent="7fa8c0f79203be72", sampled=1)

Uh oh!

X-Ray trace header parsing and Lambda tracing propagation #11708

X-Ray trace header parsing and Lambda tracing propagation #11708

Uh oh!

Conversation

joe4dev commented Oct 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Changes

Testing

Usage

Uh oh!

github-actions bot commented Oct 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

S3 Image Test Results (AMD64 / ARM64)

Uh oh!

joe4dev left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Oct 18, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

LocalStack Community integration with Pro

Uh oh!

bentsku left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dominikschubert left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

joe4dev commented Oct 17, 2024 •

edited

Loading

github-actions bot commented Oct 17, 2024 •

edited

Loading

github-actions bot commented Oct 18, 2024 •

edited

Loading