feat(aws-lambda): Continue the trace from Invoke ClientContext

sentry_sdk/integrations/aws_lambda/_handler.py

@@ -20,7 +20,7 @@
     ORIGIN,
 )
 from sentry_sdk.scope import Scope
-from sentry_sdk.traces import SegmentNameSource
+from sentry_sdk.traces import SENTRY_TRACE_HEADER_NAME, SegmentNameSource
 from sentry_sdk.utils import (
     TimeoutThread,
     capture_internal_exceptions,
@@ -187,7 +187,15 @@ def sentry_handler(
                     }
                 )
 
-            sentry_sdk.continue_trace(headers)
+            # a synchronous `Invoke` passes the invoking span in `ClientContext`; it wins over
+            # event headers, which a forwarded HTTP event keeps from further upstream.
+            # https://docs.aws.amazon.com/lambda/latest/dg/python-context.html
+            client_context = getattr(aws_context, "client_context", None)
+            custom = getattr(client_context, "custom", None)
+            if custom and SENTRY_TRACE_HEADER_NAME in custom:
+                sentry_sdk.continue_trace(custom)
+            else:
+                sentry_sdk.continue_trace(headers)
             Scope.set_custom_sampling_context(sampling_context)
             with sentry_sdk.start_span(
                 name=function_name,

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/ClientContextTrace/.gitignore

@@ -0,0 +1,11 @@
+# Need to add some ignore rules in this directory, because the unit tests will add the Sentry SDK and its dependencies
+# into this directory to create a Lambda function package that contains everything needed to instrument a Lambda function using Sentry.
+
+# Ignore everything
+*
+
+# But not index.py
+!index.py
+
+# And not .gitignore itself
+!.gitignore

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/ClientContextTrace/index.py

@@ -0,0 +1,45 @@
+import json
+import os
+
+import sentry_sdk
+from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration
+from sentry_sdk.integrations.aws_lambda._runtime import get_lambda_bootstrap
+
+sentry_sdk.init(
+    dsn=os.environ.get("SENTRY_DSN"),
+    traces_sample_rate=1.0,
+    integrations=[AwsLambdaIntegration()],
+)
+
+# `sam local` drops `ClientContext`; pass the event's `client_context` to the runtime
+# instead, like lambda does for a synchronous `Invoke`.
+bootstrap = get_lambda_bootstrap()
+handle_event_request = bootstrap.handle_event_request
+
+
+def handle_event_request_with_client_context(
+    lambda_runtime_client,
+    request_handler,
+    invoke_id,
+    event_body,
+    content_type,
+    client_context_json,
+    *args,
+):
+    client_context_json = json.dumps(json.loads(event_body)["client_context"])
+    return handle_event_request(
+        lambda_runtime_client,
+        request_handler,
+        invoke_id,
+        event_body,
+        content_type,
+        client_context_json,
+        *args,
+    )
+
+
+bootstrap.handle_event_request = handle_event_request_with_client_context
+
+
+def handler(event, context):
+    return {"event": event}

tests/integrations/aws_lambda/test_aws_lambda.py

@@ -870,6 +870,39 @@ def test_trace_continuation(lambda_client, test_environment):
     assert SPANDATA.FAAS_INVOCATION_ID in attrs
 
 
+@pytest.mark.parametrize(
+    "headers",
+    [
+        pytest.param({}, id="client-context"),
+        # `ClientContext` wins over event headers.
+        pytest.param(
+            {"sentry-trace": "0f4c2a5b8d9e4f1a2b3c4d5e6f7a8b9c-1a2b3c4d5e6f7a8b-1"},
+            id="client-context-and-event-headers",
+        ),
+    ],
+)
+def test_trace_continuation_from_client_context(
+    lambda_client, test_environment, headers
+):
+    payload = {
+        "headers": headers,
+        # `ClientContextTrace` passes it to the runtime as the invocation's `ClientContext`.
+        "client_context": {
+            "custom": {
+                "sentry-trace": "471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"
+            }
+        },
+    }
+
+    lambda_client.invoke(FunctionName="ClientContextTrace", Payload=json.dumps(payload))
+
+    (segment_span,) = [
+        s for s in test_environment["server"].span_items if s["is_segment"]
+    ]
+    assert segment_span["trace_id"] == "471a43a4192642f0b136d5159a501701"
+    assert segment_span["parent_span_id"] == "6e8f22c393e68f19"
+
+
 def test_request_attributes(lambda_client, test_environment):
     payload = {
         "headers": {