Repository: elastic/apm-agent-python

Issue:
[META 605] Add span links for message tracing
Meta issue:
- https://github.com/elastic/apm/issues/605

Summary of spec changes:

- If your agent instruments any of Azure Service Bus, JMS, Kafka, RabbitMQ, SQS, or SNS, then instrumentation for receiving messages or batches of messages "SHOULD":
    - Create a transaction (with trace-context from message metadata, if any) for the handling of each message, if possible.
    - Otherwise, when creating the transaction or span for the message or batch of messages, look for trace-context in metadata for up to 1000 messages and create a span link for each.
- If your agent supports AWS Lambda instrumentation, then in addition to the creation of span links, this spec change has a number of changes to the context captured for SQS and SNS-triggered invocations. (The special-casing of a message batch with a single message has been dropped.)


Discussion:
Need to add handling of distributed tracing / span links for SQS and SNS for AWS Lambda instrumentation

Resolving patch:
diff --git a/elasticapm/base.py b/elasticapm/base.py
--- a/elasticapm/base.py
+++ b/elasticapm/base.py
@@ -302,6 +302,7 @@ def begin_transaction(
         :param trace_parent: an optional TraceParent object for distributed tracing
         :param start: override the start timestamp, mostly useful for testing
         :param auto_activate: whether to set this transaction in execution_context
+        :param links: a sequence of traceparent objects to causally link this transaction with
         :return: the started transaction object
         """
         if self.config.is_recording:
diff --git a/elasticapm/contrib/serverless/aws.py b/elasticapm/contrib/serverless/aws.py
--- a/elasticapm/contrib/serverless/aws.py
+++ b/elasticapm/contrib/serverless/aws.py
@@ -184,7 +184,16 @@ def __enter__(self):
                 transaction_type = "messaging"
                 transaction_name = "RECEIVE {}".format(record["eventSourceARN"].split(":")[5])
 
-        self.transaction = self.client.begin_transaction(transaction_type, trace_parent=trace_parent)
+        if "Records" in self.event:
+            links = [
+                TraceParent.from_string(record["messageAttributes"]["traceparent"]["stringValue"])
+                for record in self.event["Records"][:1000]
+                if "messageAttributes" in record and "traceparent" in record["messageAttributes"]
+            ]
+        else:
+            links = []
+
+        self.transaction = self.client.begin_transaction(transaction_type, trace_parent=trace_parent, links=links)
         elasticapm.set_transaction_name(transaction_name, override=False)
         if self.source in SERVERLESS_HTTP_REQUEST:
             elasticapm.set_context(


Test patch:
diff --git a/tests/contrib/serverless/aws_sqs_test_data.json b/tests/contrib/serverless/aws_sqs_test_data.json
--- a/tests/contrib/serverless/aws_sqs_test_data.json
+++ b/tests/contrib/serverless/aws_sqs_test_data.json
@@ -28,6 +28,12 @@
                     "stringListValues": [],
                     "binaryListValues": [],
                     "dataType": "String"
+                },
+                "traceparent": {
+                    "stringValue": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00",
+                    "stringListValues": [],
+                    "binaryListValues": [],
+                    "dataType": "String"
                 }
             },
             "md5OfBody": "5eb63bbbe01eeed093cb22bb8f5acdc3",
@@ -36,4 +42,4 @@
             "awsRegion": "us-east-1"
         }
     ]
-}
\ No newline at end of file
+}
diff --git a/tests/contrib/serverless/aws_tests.py b/tests/contrib/serverless/aws_tests.py
--- a/tests/contrib/serverless/aws_tests.py
+++ b/tests/contrib/serverless/aws_tests.py
@@ -324,6 +324,8 @@ def test_func(event, context):
     assert transaction["span_count"]["started"] == 1
     assert transaction["context"]["message"]["headers"]["Population"] == "1250800"
     assert transaction["context"]["message"]["headers"]["City"] == "Any City"
+    assert len(transaction["links"]) == 1
+    assert transaction["links"][0] == {"trace_id": "0af7651916cd43dd8448eb211c80319c", "span_id": "b7ad6b7169203331"}
 
 
 def test_capture_serverless_s3_batch(event_s3_batch, context, elasticapm_client):