From 86913c0f55b2fc210d5a14b1c192050134e5f805 Mon Sep 17 00:00:00 2001
From: DarcyRaynerDD <darcy.rayner@datadoghq.com>
Date: Mon, 20 Sep 2021 16:23:53 -0400
Subject: [PATCH 1/3] Address X-Ray/dd-trace trace mergeing edge case

---
 datadog_lambda/tracing.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
index 9870f756..2c3c109d 100644
--- a/datadog_lambda/tracing.py
+++ b/datadog_lambda/tracing.py
@@ -362,7 +362,13 @@ def is_lambda_context():
 
 def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
     if trace_context_source == TraceContextSource.EVENT or merge_xray_traces:
-        headers = _context_obj_to_headers(dd_trace_context)
+        context = dict(dd_trace_context)
+        if merge_xray_traces:
+            xray_context = _get_xray_trace_context()
+            if xray_context != None:
+                context["parent-id"] = xray_context["parent-id"]
+
+        headers = _context_obj_to_headers(context)
         span_context = propagator.extract(headers)
         tracer.context_provider.activate(span_context)
         logger.debug(

From 258d1f68e91602b41a9ecbc3487d8198ac3351a0 Mon Sep 17 00:00:00 2001
From: DarcyRaynerDD <darcy.rayner@datadoghq.com>
Date: Mon, 20 Sep 2021 17:43:30 -0400
Subject: [PATCH 2/3] Add test for x-ray tracing edge case

---
 datadog_lambda/tracing.py |  2 +-
 tests/test_tracing.py     | 49 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
index 2c3c109d..4f679fe1 100644
--- a/datadog_lambda/tracing.py
+++ b/datadog_lambda/tracing.py
@@ -365,7 +365,7 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
         context = dict(dd_trace_context)
         if merge_xray_traces:
             xray_context = _get_xray_trace_context()
-            if xray_context != None:
+            if xray_context is not None:
                 context["parent-id"] = xray_context["parent-id"]
 
         headers = _context_obj_to_headers(context)
diff --git a/tests/test_tracing.py b/tests/test_tracing.py
index 82e11f3f..0bd50143 100644
--- a/tests/test_tracing.py
+++ b/tests/test_tracing.py
@@ -8,6 +8,7 @@
     from mock import MagicMock, patch, call
 
 from ddtrace.helpers import get_correlation_ids
+from ddtrace.context import Context
 
 from datadog_lambda.constants import SamplingPriority, TraceHeader, XraySubsegment
 from datadog_lambda.tracing import (
@@ -16,6 +17,7 @@
     create_function_execution_span,
     get_dd_trace_context,
     set_correlation_ids,
+    set_dd_trace_py_root,
     _convert_xray_trace_id,
     _convert_xray_entity_id,
     _convert_xray_sampling,
@@ -493,3 +495,50 @@ def test_function_with_trigger_tags(self):
         self.assertEqual(
             span.get_tag("function_trigger.event_source"), "cloudwatch-logs"
         )
+
+
+class TestSetTraceRootSpan(unittest.TestCase):
+    def setUp(self):
+        global dd_tracing_enabled
+        dd_tracing_enabled = False
+        os.environ["_X_AMZN_TRACE_ID"] = fake_xray_header_value
+        patcher = patch("datadog_lambda.tracing.send_segment")
+        self.mock_send_segment = patcher.start()
+        self.addCleanup(patcher.stop)
+        patcher = patch("datadog_lambda.tracing.is_lambda_context")
+        self.mock_is_lambda_context = patcher.start()
+        self.mock_is_lambda_context.return_value = True
+        self.addCleanup(patcher.stop)
+        patcher = patch("ddtrace.tracer.context_provider.activate")
+        self.mock_activate = patcher.start()
+        self.mock_activate.return_value = True
+        self.addCleanup(patcher.stop)
+
+    def tearDown(self):
+        global dd_tracing_enabled
+        dd_tracing_enabled = False
+        del os.environ["_X_AMZN_TRACE_ID"]
+
+    def test_set_parent_to_incomming_dd_trace_id(self):
+        lambda_ctx = get_mock_context()
+        ctx, source = extract_dd_trace_context(
+            {
+                "headers": {
+                    TraceHeader.TRACE_ID: "123",
+                    TraceHeader.PARENT_ID: "321",
+                    TraceHeader.SAMPLING_PRIORITY: "1",
+                }
+            },
+            lambda_ctx,
+        )
+        set_dd_trace_py_root(
+            source, True
+        )  # When merging is off, always use dd-trace-context
+
+        expected_context = Context(
+            trace_id=123,  # Trace Id from incomming context
+            span_id=int(fake_xray_header_value_parent_decimal),  # Parent Id from x-ray
+            sampling_priority=1,  # Sampling priority from incomming context
+        )
+        self.mock_activate.assert_called()
+        self.mock_activate.assert_has_calls([call(expected_context)])

From d2050dbc5ef2f547d340c373556692f1efbc88c0 Mon Sep 17 00:00:00 2001
From: DarcyRaynerDD <darcy.rayner@datadoghq.com>
Date: Mon, 20 Sep 2021 17:57:31 -0400
Subject: [PATCH 3/3] Add comment

---
 tests/test_tracing.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/test_tracing.py b/tests/test_tracing.py
index 0bd50143..e0ba12c4 100644
--- a/tests/test_tracing.py
+++ b/tests/test_tracing.py
@@ -519,7 +519,10 @@ def tearDown(self):
         dd_tracing_enabled = False
         del os.environ["_X_AMZN_TRACE_ID"]
 
-    def test_set_parent_to_incomming_dd_trace_id(self):
+    def test_mixed_parent_context_when_merging(self):
+        # When trace merging is enabled, and dd_trace headers are present,
+        # use the dd-trace trace-id and the x-ray parent-id
+        # This allows parenting relationships like dd-trace -> x-ray -> dd-trace
         lambda_ctx = get_mock_context()
         ctx, source = extract_dd_trace_context(
             {