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"))