Don't send x-ray trace metadata for when x-ray has filtered trace (#193)

Author: DarcyRaynerDD

datadog_lambda/xray.py

@@ -10,11 +10,11 @@
 logger = logging.getLogger(__name__)
 
 
-def get_xray_host_port(adress):
-    if adress == "":
+def get_xray_host_port(address):
+    if address == "":
         logger.debug("X-Ray daemon env var not set, not sending sub-segment")
         return None
-    parts = adress.split(":")
+    parts = address.split(":")
     if len(parts) <= 1:
         logger.debug("X-Ray daemon env var not set, not sending sub-segment")
         return None
@@ -108,6 +108,11 @@ def send_segment(key, metadata):
             "Failed to create segment since it was not possible to get trace context from header"
         )
         return None
+
+    # Skip adding segment, if the xray trace is going to be sampled away.
+    if context["sampled"] == "0":
+        logger.debug("Skipping sending metadata, x-ray trace was sampled out")
+        return None
     segment = build_segment(context, key, metadata)
     segment_payload = build_segment_payload(segment)
     send(host_port_tuple, segment_payload)

tests/Dockerfile

@@ -7,7 +7,8 @@ RUN mkdir -p /test/datadog_lambda
 WORKDIR /test
 
 # Copy minimal subset of files to make pip install succeed and be cached (next docker builds will be way faster)
-COPY setup.py .
+COPY pyproject.toml .
+COPY poetry.lock .
 COPY README.md .
 COPY datadog_lambda/__init__.py datadog_lambda/__init__.py 
 

tests/test_xray.py

@@ -1,12 +1,25 @@
 import unittest
 import json
+import os
 
 from unittest.mock import MagicMock, patch
 
-from datadog_lambda.xray import get_xray_host_port, build_segment_payload, build_segment
+from datadog_lambda.xray import (
+    get_xray_host_port,
+    build_segment_payload,
+    build_segment,
+    send_segment,
+)
 
 
 class TestXRay(unittest.TestCase):
+    def tearDown(self):
+        if os.environ.get("_X_AMZN_TRACE_ID"):
+            os.environ.pop("_X_AMZN_TRACE_ID")
+        if os.environ.get("AWS_XRAY_DAEMON_ADDRESS"):
+            os.environ.pop("AWS_XRAY_DAEMON_ADDRESS")
+        return super().tearDown()
+
     def test_get_xray_host_port_empty_(self):
         result = get_xray_host_port("")
         self.assertIsNone(result)
@@ -20,6 +33,31 @@ def test_get_xray_host_port_success(self):
         self.assertEqual("mySuperHost", result[0])
         self.assertEqual(1000, result[1])
 
+    def test_send_segment_sampled_out(self):
+        os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080"
+        os.environ[
+            "_X_AMZN_TRACE_ID"
+        ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0"
+
+        with patch(
+            "datadog_lambda.xray.send", MagicMock(return_value=None)
+        ) as mock_send:
+            # XRay trace won't be sampled according to the trace header.
+            send_segment("my_key", {"data": "value"})
+            self.assertFalse(mock_send.called)
+
+    def test_send_segment_sampled(self):
+        os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080"
+        os.environ[
+            "_X_AMZN_TRACE_ID"
+        ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"
+        with patch(
+            "datadog_lambda.xray.send", MagicMock(return_value=None)
+        ) as mock_send:
+            # X-Ray trace will be sampled according to the trace header.
+            send_segment("my_key", {"data": "value"})
+            self.assertTrue(mock_send.called)
+
     def test_build_segment_payload_ok(self):
         exected_text = '{"format": "json", "version": 1}\nmyPayload'
         self.assertEqual(exected_text, build_segment_payload("myPayload"))