fix(lambda): align Python context propagation with JS implementation

CHANGELOG.md

@@ -12,6 +12,8 @@ If your change does not need a CHANGELOG entry, add the "skip changelog" label t
 
 ## Unreleased
 
+- fix(lambda-layer): align context propagation with JS — delegate to global propagator so W3C traceparent is no longer ignored when X-Ray active tracing is enabled
+  ([#727](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/727))
 - feat: support environment-configured endpoint visibility for HTTP operation names
   ([#718](https://github.com/aws-observability/aws-otel-python-instrumentation/pull/718))
 - fix(lambda-layer): Disable all agentic instrumentation in Lambda by default

lambda-layer/src/otel_wrapper.py

@@ -41,7 +41,6 @@
 from opentelemetry.context import Context
 from opentelemetry.instrumentation.aws_lambda import _X_AMZN_TRACE_ID, AwsLambdaInstrumentor
 from opentelemetry.propagate import get_global_textmap
-from opentelemetry.propagators.aws import AwsXRayPropagator
 from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_HEADER_KEY
 from opentelemetry.trace import get_current_span
 
@@ -57,21 +56,24 @@ class HandlerError(Exception):
 
 def custom_event_context_extractor(lambda_event: Any) -> Context:
     xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
-    lambda_trace_context = AwsXRayPropagator().extract({TRACE_HEADER_KEY: xray_env_var})
-    parent_span_context = get_current_span(lambda_trace_context).get_span_context()
 
-    if parent_span_context is None or not parent_span_context.is_valid:
+    try:
+        headers = lambda_event["headers"]
+    except (TypeError, KeyError):
         headers = None
-        try:
-            headers = lambda_event["headers"]
-        except (TypeError, KeyError):
-            pass
-        if not isinstance(headers, dict):
-            headers = {}
+    if not isinstance(headers, dict):
+        headers = {}
+    else:
+        headers = {k: v for k, v in headers.items()}
 
-        return get_global_textmap().extract(headers)
+    if xray_env_var:
+        headers = {k: v for k, v in headers.items() if k.lower() != TRACE_HEADER_KEY.lower()}
+        headers[TRACE_HEADER_KEY] = xray_env_var
 
-    return lambda_trace_context
+    extracted_context = get_global_textmap().extract(headers)
+    if get_current_span(extracted_context).get_span_context():
+        return extracted_context
+    return Context()
 
 
 AwsLambdaInstrumentor().instrument(event_context_extractor=custom_event_context_extractor)
@@ -82,7 +84,7 @@ def custom_event_context_extractor(lambda_event: Any) -> Context:
     raise HandlerError("ORIG_HANDLER is not defined.")
 
 try:
-    (mod_name, handler_name) = path.rsplit(".", 1)
+    mod_name, handler_name = path.rsplit(".", 1)
 except ValueError as e:
     raise HandlerError("Bad path '{}' for ORIG_HANDLER: {}".format(path, str(e)))
 

lambda-layer/src/tests/test_lambda_instrumentation.py

@@ -25,9 +25,11 @@
 from shutil import which
 from unittest import mock
 
-from opentelemetry.environment_variables import OTEL_PROPAGATORS
 from opentelemetry.instrumentation.aws_lambda import _HANDLER, _X_AMZN_TRACE_ID, ORIG_HANDLER, AwsLambdaInstrumentor
+from opentelemetry.propagate import get_global_textmap, set_global_textmap
+from opentelemetry.propagators.aws import AwsXRayPropagator
 from opentelemetry.propagators.aws.aws_xray_propagator import TRACE_ID_FIRST_PART_LENGTH, TRACE_ID_VERSION
+from opentelemetry.propagators.composite import CompositePropagator
 from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.test_base import TestBase
@@ -196,6 +198,9 @@ def tearDownClass(cls):
         )
 
     def test_active_tracing(self):
+        original_propagator = get_global_textmap()
+        set_global_textmap(AwsXRayPropagator())
+
         test_env_patch = mock.patch.dict(
             "os.environ",
             {
@@ -206,71 +211,162 @@ def test_active_tracing(self):
         )
         test_env_patch.start()
 
-        mock_execute_lambda()
-
-        spans = self.memory_exporter.get_finished_spans()
+        try:
+            mock_execute_lambda()
+
+            spans = self.memory_exporter.get_finished_spans()
+
+            assert spans
+
+            self.assertEqual(len(spans), 1)
+            span = spans[0]
+            self.assertEqual(span.name, os.environ[ORIG_HANDLER])
+            self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
+            self.assertEqual(span.kind, SpanKind.SERVER)
+            self.assertSpanHasAttributes(
+                span,
+                {
+                    ResourceAttributes.CLOUD_RESOURCE_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
+                    SpanAttributes.FAAS_INVOCATION_ID: MOCK_LAMBDA_CONTEXT.aws_request_id,
+                },
+            )
+
+            parent_context = span.parent
+            self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
+            self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
+            self.assertTrue(parent_context.is_remote)
+        finally:
+            test_env_patch.stop()
+            set_global_textmap(original_propagator)
 
-        assert spans
+    def test_parent_context_from_lambda_event(self):
+        original_propagator = get_global_textmap()
+        set_global_textmap(TraceContextTextMapPropagator())
 
-        self.assertEqual(len(spans), 1)
-        span = spans[0]
-        self.assertEqual(span.name, os.environ[ORIG_HANDLER])
-        self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
-        self.assertEqual(span.kind, SpanKind.SERVER)
-        self.assertSpanHasAttributes(
-            span,
+        test_env_patch = mock.patch.dict(
+            "os.environ",
             {
-                ResourceAttributes.CLOUD_RESOURCE_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
-                SpanAttributes.FAAS_INVOCATION_ID: MOCK_LAMBDA_CONTEXT.aws_request_id,
+                **os.environ,
+                # NOT Active Tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH,
             },
         )
+        test_env_patch.start()
 
-        parent_context = span.parent
-        self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
-        self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
-        self.assertTrue(parent_context.is_remote)
-
-        test_env_patch.stop()
+        try:
+            trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2"
+            mock_execute_lambda(
+                {
+                    "headers": {
+                        TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header,
+                    }
+                }
+            )
+
+            spans = self.memory_exporter.get_finished_spans()
+
+            assert spans
+
+            self.assertEqual(len(spans), 1)
+            span = spans[0]
+            self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+
+            parent_context = span.parent
+            self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
+            self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+            self.assertEqual(len(parent_context.trace_state), 3)
+            self.assertEqual(
+                parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
+                MOCK_W3C_TRACE_STATE_VALUE,
+            )
+            self.assertTrue(parent_context.is_remote)
+        finally:
+            test_env_patch.stop()
+            set_global_textmap(original_propagator)
+
+    def test_xray_ignored_when_propagator_does_not_include_xray(self):
+        """When X-Ray active tracing is on but the user only configures
+        tracecontext propagator (no xray), the X-Ray env var is injected
+        into headers but the propagator cannot parse it. The W3C context
+        from event headers should be used instead."""
+        original_propagator = get_global_textmap()
+        set_global_textmap(TraceContextTextMapPropagator())
 
-    def test_parent_context_from_lambda_event(self):
         test_env_patch = mock.patch.dict(
             "os.environ",
             {
                 **os.environ,
-                # NOT Active Tracing
-                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_PASSTHROUGH,
-                # NOT using the X-Ray Propagator
-                OTEL_PROPAGATORS: "tracecontext",
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
             },
         )
         test_env_patch.start()
 
-        trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2"
-        mock_execute_lambda(
-            {
-                "headers": {
-                    TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
-                    TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header,
+        try:
+            mock_execute_lambda(
+                {
+                    "headers": {
+                        TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                    }
                 }
-            }
-        )
+            )
+
+            spans = self.memory_exporter.get_finished_spans()
 
-        spans = self.memory_exporter.get_finished_spans()
+            assert spans
+            self.assertEqual(len(spans), 1)
+            span = spans[0]
 
-        assert spans
+            self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
 
-        self.assertEqual(len(spans), 1)
-        span = spans[0]
-        self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+            parent_context = span.parent
+            self.assertEqual(parent_context.trace_id, MOCK_W3C_TRACE_ID)
+            self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+            self.assertTrue(parent_context.is_remote)
+        finally:
+            test_env_patch.stop()
+            set_global_textmap(original_propagator)
 
-        parent_context = span.parent
-        self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
-        self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
-        self.assertEqual(len(parent_context.trace_state), 3)
-        self.assertEqual(
-            parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
-            MOCK_W3C_TRACE_STATE_VALUE,
+    def test_w3c_takes_precedence_over_xray_when_both_present(self):
+        """When both X-Ray active tracing and W3C traceparent headers are
+        present, the W3C context should be used as the parent because
+        tracecontext comes after xray in the composite propagator — the
+        last propagator to extract a valid context wins."""
+        original_propagator = get_global_textmap()
+        set_global_textmap(CompositePropagator([AwsXRayPropagator(), TraceContextTextMapPropagator()]))
+
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+            },
         )
-        self.assertTrue(parent_context.is_remote)
+        test_env_patch.start()
+
+        try:
+            trace_state_header = f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2"
+            mock_execute_lambda(
+                {
+                    "headers": {
+                        TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: trace_state_header,
+                    }
+                }
+            )
+
+            spans = self.memory_exporter.get_finished_spans()
+
+            assert spans
+            self.assertEqual(len(spans), 1)
+            span = spans[0]
+
+            self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
 
-        test_env_patch.stop()
+            parent_context = span.parent
+            self.assertEqual(parent_context.trace_id, MOCK_W3C_TRACE_ID)
+            self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+            self.assertTrue(parent_context.is_remote)
+        finally:
+            test_env_patch.stop()
+            set_global_textmap(original_propagator)