feat(aws-lambda): Add queue.process spans for messaging triggers

sentry_sdk/integrations/aws_lambda/_handler.py

@@ -1,5 +1,6 @@
 import functools
 import sys
+from contextlib import nullcontext
 from os import environ
 from typing import TYPE_CHECKING
 from urllib.parse import urlencode
@@ -12,7 +13,10 @@
     _get_user_from_event,
     _make_request_event_processor,
 )
-from sentry_sdk.integrations.aws_lambda._trigger import _get_trigger
+from sentry_sdk.integrations.aws_lambda._trigger import (
+    _get_trigger,
+    _start_process_span,
+)
 from sentry_sdk.integrations.aws_lambda.consts import (
     CLOUD_PLATFORM,
     CLOUD_PROVIDER,
@@ -30,7 +34,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Mapping, TypeVar
+    from typing import Any, Callable, ContextManager, Mapping, TypeVar
 
     from sentry_sdk._types import Attributes
 
@@ -204,7 +208,11 @@ def sentry_handler(
                 attributes=attributes,
             ):
                 try:
-                    return handler(aws_event, aws_context, *args, **kwargs)
+                    process_span: "ContextManager[Any]" = nullcontext()
+                    with capture_internal_exceptions():
+                        process_span = _start_process_span(aws_event)
+                    with process_span:
+                        return handler(aws_event, aws_context, *args, **kwargs)
                 except Exception:
                     exc_info = sys.exc_info()
                     sentry_event, hint = event_from_exception(

sentry_sdk/integrations/aws_lambda/_trigger.py

@@ -1,7 +1,18 @@
+import json
+from contextlib import nullcontext
 from typing import TYPE_CHECKING
 
+import sentry_sdk
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.integrations.aws_lambda.consts import ORIGIN
+from sentry_sdk.traces import SENTRY_TRACE_HEADER_NAME
+from sentry_sdk.tracing_utils import extract_sentrytrace_data
+from sentry_sdk.utils import capture_internal_exceptions
+
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, ContextManager, Optional
+
+    from sentry_sdk._types import SpanLink
 
 # `faas.trigger` of each `eventSource` of `Records` or of a top-level `eventSource`.
 # https://docs.aws.amazon.com/lambda/latest/dg/lambda-services.html
@@ -19,6 +30,13 @@
     "aws:docdb": "datasource",
 }
 
+# `messaging.system` of the event sources whose producers can propagate a trace.
+_MESSAGING_SYSTEMS = {"aws:sqs": "aws_sqs", "aws:sns": "aws.sns"}
+
+# OTel's default 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
+
 # https://docs.aws.amazon.com/eventbridge/latest/ref/overiew-event-structure.html
 _EVENTBRIDGE_KEYS = frozenset(
     (
@@ -87,3 +105,94 @@ def _get_trigger(aws_event: "Any") -> str:
         return "pubsub"
 
     return _TRIGGERS.get(aws_event.get("eventSource"), "other")
+
+
+def _get_producer_trace(record: "dict[str, Any]") -> "Optional[str]":
+    """The `sentry-trace` the producer put in a record's message attributes."""
+    # https://docs.aws.amazon.com/lambda/latest/dg/with-sns.html
+    if record.get("EventSource") == "aws:sns":
+        attribute = (
+            record["Sns"].get("MessageAttributes", {}).get(SENTRY_TRACE_HEADER_NAME, {})
+        )
+        return attribute.get("Value")
+
+    # https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
+    attribute = record.get("messageAttributes", {}).get(SENTRY_TRACE_HEADER_NAME)
+    if attribute is not None:
+        return attribute.get("stringValue")
+
+    # without raw message delivery, an SNS subscription delivers the notification,
+    # with its message attributes, as the body.
+    # https://docs.aws.amazon.com/sns/latest/dg/sns-sqs-as-subscriber.html
+    if '"TopicArn"' in record["body"]:
+        notification = json.loads(record["body"])
+        attribute = notification.get("MessageAttributes", {}).get(
+            SENTRY_TRACE_HEADER_NAME, {}
+        )
+        return attribute.get("Value")
+
+    return None
+
+
+def _get_producer_links(records: "list[Any]") -> "list[SpanLink]":
+    """One link per distinct producer context; messages never parent the invocation."""
+    links: "list[SpanLink]" = []
+    linked = set()
+    for record in records:
+        with capture_internal_exceptions():
+            producer = extract_sentrytrace_data(_get_producer_trace(record))
+            if producer is None:
+                continue
+
+            trace_id = str(producer["trace_id"])
+            span_id = str(producer["parent_span_id"])
+            if (trace_id, span_id) in linked:
+                continue
+            linked.add((trace_id, span_id))
+
+            link: "SpanLink" = {
+                "trace_id": trace_id,
+                "span_id": span_id,
+                "attributes": {SPANDATA.SENTRY_LINK_TYPE: "message_origin"},
+            }
+            if producer["parent_sampled"] is not None:
+                link["sampled"] = bool(producer["parent_sampled"])
+            links.append(link)
+            if len(links) == _MAX_PRODUCER_LINKS:
+                break
+
+    return links
+
+
+def _start_process_span(aws_event: "Any") -> "ContextManager[Any]":
+    """A `queue.process` span around the handler for an SQS or SNS delivery."""
+    records = aws_event.get("Records") if isinstance(aws_event, dict) else None
+    if not records:
+        return nullcontext()
+
+    event_source = _get_records_source(records)
+    if event_source not in _MESSAGING_SYSTEMS:
+        return nullcontext()
+
+    # one queue or topic per event: `arn:aws:sqs:<region>:<account>:<queue>`.
+    arn = (
+        records[0]["eventSourceARN"]
+        if event_source == "aws:sqs"
+        else records[0]["Sns"]["TopicArn"]
+    )
+    destination = arn.split(":")[5]
+
+    return sentry_sdk.start_span(
+        name=f"process {destination}",
+        attributes={
+            SPANDATA.SENTRY_OP: OP.QUEUE_PROCESS,
+            SPANDATA.SENTRY_ORIGIN: ORIGIN,
+            SPANDATA.SENTRY_KIND: "consumer",
+            SPANDATA.MESSAGING_SYSTEM: _MESSAGING_SYSTEMS[event_source],
+            SPANDATA.MESSAGING_OPERATION_TYPE: "process",
+            SPANDATA.MESSAGING_OPERATION_NAME: "process",
+            SPANDATA.MESSAGING_DESTINATION_NAME: destination,
+            SPANDATA.MESSAGING_BATCH_MESSAGE_COUNT: len(records),
+        },
+        links=_get_producer_links(records),
+    )

tests/integrations/aws_lambda/test_aws_lambda.py

@@ -24,6 +24,9 @@
 DOCKER_NETWORK_NAME = "lambda-test-network"
 SAM_TEMPLATE_FILE = "sam.template.yaml"
 SAM_SHUTDOWN_TIMEOUT = 10
+MESSAGE_ORIGIN_LINK = {
+    SPANDATA.SENTRY_LINK_TYPE: {"value": "message_origin", "type": "string"}
+}
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -830,20 +833,41 @@ def test_error(lambda_client, test_environment):
     assert SPANDATA.MESSAGING_BATCH_MESSAGE_COUNT not in attrs
 
 
-def _sqs_record():
-    return {
+def _sqs_record(sentry_trace=None):
+    record = {
         "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
         "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_record():
-    return {
+def _sns_record(sentry_trace=None):
+    record = {
         "EventSource": "aws:sns",
-        "Sns": {"TopicArn": "arn:aws:sns:us-east-1:012345678912:alerts"},
+        "Sns": {
+            "Type": "Notification",
+            "MessageId": "95df01b4-ee98-5cb9-9903-4c221d41eb5e",
+            "TopicArn": "arn:aws:sns:us-east-1:012345678912:alerts",
+            "Message": "{}",
+            "MessageAttributes": {},
+        },
     }
+    if sentry_trace is not None:
+        record["Sns"]["MessageAttributes"]["sentry-trace"] = {
+            "Type": "String",
+            "Value": sentry_trace,
+        }
+    return record
 
 
 def _appsync_event():
@@ -929,6 +953,151 @@ def test_trigger(lambda_client, test_environment, aws_event, trigger):
     assert _get_span_attr(attrs, SPANDATA.FAAS_TRIGGER) == trigger
 
 
+@pytest.mark.parametrize(
+    "aws_event,system,destination,links",
+    [
+        pytest.param(
+            {
+                "Records": [
+                    _sqs_record("471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"),
+                    # the same send context: linked once.
+                    _sqs_record("471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"),
+                    _sqs_record("0f4c2a5b8d9e4f1a2b3c4d5e6f7a8b9c-1a2b3c4d5e6f7a8b-0"),
+                    _sqs_record(),
+                ]
+            },
+            "aws_sqs",
+            "orders",
+            [
+                {
+                    "trace_id": "471a43a4192642f0b136d5159a501701",
+                    "span_id": "6e8f22c393e68f19",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": True,
+                },
+                {
+                    "trace_id": "0f4c2a5b8d9e4f1a2b3c4d5e6f7a8b9c",
+                    "span_id": "1a2b3c4d5e6f7a8b",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": False,
+                },
+            ],
+            id="sqs",
+        ),
+        pytest.param(
+            {
+                "Records": [
+                    _sns_record("471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1")
+                ]
+            },
+            "aws.sns",
+            "alerts",
+            [
+                {
+                    "trace_id": "471a43a4192642f0b136d5159a501701",
+                    "span_id": "6e8f22c393e68f19",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": True,
+                }
+            ],
+            id="sns",
+        ),
+        pytest.param(
+            {
+                "Records": [
+                    {
+                        **_sqs_record(),
+                        "body": json.dumps(
+                            _sns_record(
+                                "471a43a4192642f0b136d5159a501701-6e8f22c393e68f19-1"
+                            )["Sns"]
+                        ),
+                    }
+                ]
+            },
+            "aws_sqs",
+            "orders",
+            [
+                {
+                    "trace_id": "471a43a4192642f0b136d5159a501701",
+                    "span_id": "6e8f22c393e68f19",
+                    "attributes": MESSAGE_ORIGIN_LINK,
+                    "sampled": True,
+                }
+            ],
+            id="sns-to-sqs",
+        ),
+    ],
+)
+def test_messaging_process_span(
+    lambda_client, test_environment, aws_event, system, destination, links
+):
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreaming",
+        Payload=json.dumps(aws_event),
+    )
+
+    span_items = test_environment["server"].span_items
+    (segment_span,) = [s for s in span_items if s["is_segment"]]
+    (process_span,) = [
+        s
+        for s in span_items
+        if _get_span_attr(s["attributes"], SPANDATA.SENTRY_OP) == OP.QUEUE_PROCESS
+    ]
+    # messages are linked, never the parent of the invocation.
+    assert "parent_span_id" not in segment_span
+    assert process_span["parent_span_id"] == segment_span["span_id"]
+    assert process_span["trace_id"] == segment_span["trace_id"]
+    assert process_span["name"] == f"process {destination}"
+    assert process_span["links"] == links
+
+    attrs = process_span["attributes"]
+    assert _get_span_attr(attrs, SPANDATA.SENTRY_KIND) == "consumer"
+    assert _get_span_attr(attrs, SPANDATA.MESSAGING_SYSTEM) == system
+    assert _get_span_attr(attrs, SPANDATA.MESSAGING_OPERATION_TYPE) == "process"
+    assert _get_span_attr(attrs, SPANDATA.MESSAGING_OPERATION_NAME) == "process"
+    assert _get_span_attr(attrs, SPANDATA.MESSAGING_DESTINATION_NAME) == destination
+    assert _get_span_attr(attrs, SPANDATA.MESSAGING_BATCH_MESSAGE_COUNT) == len(
+        aws_event["Records"]
+    )
+
+
+def test_no_process_span_without_messages(lambda_client, test_environment):
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreaming",
+        Payload=json.dumps({"Records": [{"eventSource": "aws:kinesis"}]}),
+    )
+
+    (span,) = test_environment["server"].span_items
+    assert span["is_segment"]
+
+
+def test_process_span_error(lambda_client, test_environment):
+    response = lambda_client.invoke(
+        FunctionName="RaiseErrorSpanStreaming",
+        Payload=json.dumps({"Records": [_sqs_record()]}),
+    )
+
+    # the handler's exception still reaches the Lambda runtime.
+    assert json.loads(response["Payload"].read())["errorMessage"] == "Oh!"
+
+    (error_event,) = test_environment["server"].envelopes
+    (exception,) = error_event["exception"]["values"]
+    assert exception["value"] == "Oh!"
+    assert not exception["mechanism"]["handled"]
+
+    span_items = test_environment["server"].span_items
+    (segment_span,) = [s for s in span_items if s["is_segment"]]
+    (process_span,) = [
+        s
+        for s in span_items
+        if _get_span_attr(s["attributes"], SPANDATA.SENTRY_OP) == OP.QUEUE_PROCESS
+    ]
+    assert process_span["parent_span_id"] == segment_span["span_id"]
+    assert process_span["status"] == "error"
+    assert segment_span["status"] == "error"
+
+
 def test_trace_continuation(lambda_client, test_environment):
     trace_id = "471a43a4192642f0b136d5159a501701"
     parent_span_id = "6e8f22c393e68f19"