feat(aws-lambda): Link SQS and SNS invocations to their producers

sentry_sdk/integrations/aws_lambda/_handler.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import sys
 from os import environ
 from typing import TYPE_CHECKING
@@ -21,6 +22,7 @@
 )
 from sentry_sdk.scope import Scope
 from sentry_sdk.traces import SENTRY_TRACE_HEADER_NAME, SegmentNameSource
+from sentry_sdk.tracing_utils import extract_sentrytrace_data
 from sentry_sdk.utils import (
     TimeoutThread,
     capture_internal_exceptions,
@@ -31,12 +33,79 @@
 if TYPE_CHECKING:
     from typing import Any, Callable, Mapping, TypeVar
 
-    from sentry_sdk._types import Attributes
+    from sentry_sdk._types import Attributes, SpanLink
 
     F = TypeVar("F", bound=Callable[..., Any])
 
 TIMEOUT_WARNING_BUFFER = 1500
 MILLIS_TO_SECONDS = 1000.0
+# OTel's default link limit; Lambda delivers up to 10,000 SQS messages at once.
+# https://opentelemetry.io/docs/specs/otel/trace/sdk/#span-limits
+MAX_PRODUCER_LINKS = 128
+
+
+def _get_producer_links(aws_event: "Any", trace_id: str) -> "list[SpanLink]":
+    """Links to the spans that sent the SQS or SNS messages this invocation got."""
+    records: "list[Any]" = []
+    if isinstance(aws_event, dict):
+        records = aws_event.get("Records") or []
+
+    links: "list[SpanLink]" = []
+    for record in records:
+        if "Sns" in record:
+            # https://docs.aws.amazon.com/lambda/latest/dg/with-sns.html
+            sentry_trace = (
+                record["Sns"]
+                .get("MessageAttributes", {})
+                .get(SENTRY_TRACE_HEADER_NAME, {})
+                .get("Value")
+            )
+        else:
+            # https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
+            sentry_trace = (
+                record.get("messageAttributes", {})
+                .get(SENTRY_TRACE_HEADER_NAME, {})
+                .get("stringValue")
+            )
+            # SNS wraps the message and its attributes in a JSON envelope when raw message delivery is off.
+            # https://docs.aws.amazon.com/sns/latest/dg/sns-large-payload-raw-message-delivery.html
+            body = record.get("body", "")
+            if sentry_trace is None and f'"{SENTRY_TRACE_HEADER_NAME}"' in body:
+                try:
+                    envelope = json.loads(body)
+                except ValueError:
+                    envelope = None
+                if (
+                    isinstance(envelope, dict)
+                    and envelope.get("Type") == "Notification"
+                ):
+                    sentry_trace = (
+                        envelope.get("MessageAttributes", {})
+                        .get(SENTRY_TRACE_HEADER_NAME, {})
+                        .get("Value")
+                    )
+
+        data = extract_sentrytrace_data(sentry_trace)
+        if data is None or not data["trace_id"] or not data["parent_span_id"]:
+            continue
+        # the invocation continues this trace, so a link would point into its own trace.
+        if data["trace_id"] == trace_id:
+            continue
+
+        link: "SpanLink" = {
+            "trace_id": str(data["trace_id"]),
+            "span_id": str(data["parent_span_id"]),
+            "attributes": {SPANDATA.SENTRY_LINK_TYPE: "message_origin"},
+        }
+        if data["parent_sampled"] is not None:
+            link["sampled"] = bool(data["parent_sampled"])
+        # a batch send gives many messages the same producer span.
+        if link not in links:
+            links.append(link)
+        if len(links) >= MAX_PRODUCER_LINKS:
+            break
+
+    return links
 
 
 def _wrap_handler(handler: "F") -> "F":
@@ -196,11 +265,23 @@ def sentry_handler(
                 sentry_sdk.continue_trace(custom)
             else:
                 sentry_sdk.continue_trace(headers)
+
+            # SQS and SNS messages link to their producers; they never parent the invocation.
+            producer_links: "list[SpanLink]" = []
+            with capture_internal_exceptions():
+                producer_links = _get_producer_links(
+                    aws_event,
+                    sentry_sdk.get_current_scope()
+                    .get_active_propagation_context()
+                    .trace_id,
+                )
+
             Scope.set_custom_sampling_context(sampling_context)
             with sentry_sdk.start_span(
                 name=function_name,
                 parent_span=None,
                 attributes=attributes,
+                links=producer_links,
             ):
                 try:
                     return handler(aws_event, aws_context, *args, **kwargs)

tests/integrations/aws_lambda/test_aws_lambda.py

@@ -870,6 +870,141 @@ def test_trace_continuation(lambda_client, test_environment):
     assert SPANDATA.FAAS_INVOCATION_ID in attrs
 
 
+MESSAGE_ORIGIN_LINK = {
+    SPANDATA.SENTRY_LINK_TYPE: {"value": "message_origin", "type": "string"}
+}
+
+
+def _sqs_record(sentry_trace=None, body="{}"):
+    record = {
+        "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
+        "body": body,
+        "messageAttributes": {},
+        "eventSource": "aws:sqs",
+        "eventSourceARN": "arn:aws:sqs:us-east-1:012345678912:orders",
+    }
+    if sentry_trace is not None:
+        record["messageAttributes"]["sentry-trace"] = {
+            "stringValue": sentry_trace,
+            "stringListValues": [],
+            "binaryListValues": [],
+            "dataType": "String",
+        }
+    return record
+
+
+def _sns_notification(sentry_trace):
+    return {
+        "Type": "Notification",
+        "MessageId": "95df01b4-ee98-5cb9-9903-4c221d41eb5e",
+        "TopicArn": "arn:aws:sns:us-east-1:012345678912:alerts",
+        "Message": "{}",
+        "MessageAttributes": {
+            "sentry-trace": {"Type": "String", "Value": sentry_trace}
+        },
+    }
+
+
+@pytest.mark.parametrize(
+    "aws_event,links",
+    [
+        pytest.param(
+            {
+                "Records": [
+                    _sqs_record("471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"),
+                    # the same producer span: linked once.
+                    _sqs_record("471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"),
+                    _sqs_record("0f4c2a5b8d9e4f1a2b3c4d5e6f7a8b9c-1a2b3c4d5e6f7a8b-0"),
+                    _sqs_record(),
+                    # incomplete context: not linked.
+                    _sqs_record("-1"),
+                    # no sampling decision: linked without `sampled`.
+                    _sqs_record("5c6d7e8f9a0b4c1d8e2f3a4b5c6d7e8f-2b3c4d5e6f7a8b9c"),
+                    # a body that isn't JSON: skipped, other links kept.
+                    _sqs_record(body='not json: "sentry-trace"'),
+                ]
+            },
+            [
+                {
+                    "trace_id": "471a43a4192642f0b136d5159a501701",
+                    "span_id": "6e8f22c393e68f19",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": True,
+                },
+                {
+                    "trace_id": "0f4c2a5b8d9e4f1a2b3c4d5e6f7a8b9c",
+                    "span_id": "1a2b3c4d5e6f7a8b",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": False,
+                },
+                {
+                    "trace_id": "5c6d7e8f9a0b4c1d8e2f3a4b5c6d7e8f",
+                    "span_id": "2b3c4d5e6f7a8b9c",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                },
+            ],
+            id="sqs",
+        ),
+        pytest.param(
+            {
+                "Records": [
+                    {
+                        "EventSource": "aws:sns",
+                        "Sns": _sns_notification(
+                            "471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"
+                        ),
+                    }
+                ]
+            },
+            [
+                {
+                    "trace_id": "471a43a4192642f0b136d5159a501701",
+                    "span_id": "6e8f22c393e68f19",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": True,
+                }
+            ],
+            id="sns",
+        ),
+        pytest.param(
+            {
+                "Records": [
+                    # SNS to SQS without raw message delivery.
+                    _sqs_record(
+                        body=json.dumps(
+                            _sns_notification(
+                                "471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"
+                            )
+                        )
+                    )
+                ]
+            },
+            [
+                {
+                    "trace_id": "471a43a4192642f0b136d5159a501701",
+                    "span_id": "6e8f22c393e68f19",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": True,
+                }
+            ],
+            id="sns-to-sqs",
+        ),
+    ],
+)
+def test_producer_links(lambda_client, test_environment, aws_event, links):
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreaming",
+        Payload=json.dumps(aws_event),
+    )
+
+    (segment_span,) = [
+        s for s in test_environment["server"].span_items if s["is_segment"]
+    ]
+    # messages are linked, never the parent of the invocation.
+    assert "parent_span_id" not in segment_span
+    assert segment_span["links"] == links
+
+
 @pytest.mark.parametrize(
     "headers",
     [
@@ -892,6 +1027,11 @@ def test_trace_continuation_from_client_context(
                 "sentry-trace": "471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"
             }
         },
+        "Records": [
+            # a producer in the continued trace: not linked, it's already in the trace.
+            _sqs_record("471a43a4192642f0b136d5159a501701-1a2b3c4d5e6f7a8b-1"),
+            _sqs_record("5c6d7e8f9a0b4c1d8e2f3a4b5c6d7e8f-2b3c4d5e6f7a8b9c-1"),
+        ],
     }
 
     lambda_client.invoke(FunctionName="ClientContextTrace", Payload=json.dumps(payload))
@@ -901,6 +1041,14 @@ def test_trace_continuation_from_client_context(
     ]
     assert segment_span["trace_id"] == "471a43a4192642f0b136d5159a501701"
     assert segment_span["parent_span_id"] == "6e8f22c393e68f19"
+    assert segment_span["links"] == [
+        {
+            "trace_id": "5c6d7e8f9a0b4c1d8e2f3a4b5c6d7e8f",
+            "span_id": "2b3c4d5e6f7a8b9c",
+            "attributes": MESSAGE_ORIGIN_LINK,
+            "sampled": True,
+        }
+    ]
 
 
 def test_request_attributes(lambda_client, test_environment):