diff --git a/sentry_sdk/integrations/aws_lambda/_handler.py b/sentry_sdk/integrations/aws_lambda/_handler.py index 3d006ad19e..bde707c339 100644 --- a/sentry_sdk/integrations/aws_lambda/_handler.py +++ b/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) diff --git a/tests/integrations/aws_lambda/test_aws_lambda.py b/tests/integrations/aws_lambda/test_aws_lambda.py index 88a4bf3cfe..aa548d3241 100644 --- a/tests/integrations/aws_lambda/test_aws_lambda.py +++ b/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):