Revert "update awslambda to use _X_AMZN_TRACE_ID as a Span Link" (open-telemetry#1911)

rapphil · ocelotl · web-flow · commit 0871dd455c0a · 2023-08-21T18:42:12.000+02:00
Co-authored-by: Diego Hurtado &lt;ocelotl@users.noreply.github.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -162,9 +162,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ([#1592](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1592))
 - `opentelemetry-instrumentation-django` Allow explicit `excluded_urls` configuration through `instrument()`
   ([#1618](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1618))
-- `opentelemetry-instrumentation-aws-lambda` Use env var `_X_AMZN_TRACE_ID` as a
-  Span Link instead of parent
-  ([#1657](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1657))
 
 ### Fixed
 
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py
@@ -71,7 +71,7 @@ def custom_event_context_extractor(lambda_event):
 import os
 import time
 from importlib import import_module
-from typing import Any, Callable, Collection, Optional, Sequence
+from typing import Any, Callable, Collection
 from urllib.parse import urlencode
 
 from wrapt import wrap_function_wrapper
@@ -90,7 +90,6 @@ def custom_event_context_extractor(lambda_event):
 from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import (
-    Link,
     Span,
     SpanKind,
     TracerProvider,
@@ -107,6 +106,9 @@ def custom_event_context_extractor(lambda_event):
 OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT = (
     "OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT"
 )
+OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION = (
+    "OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION"
+)
 
 
 def _default_event_context_extractor(lambda_event: Any) -> Context:
@@ -140,12 +142,14 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
 
 
 def _determine_parent_context(
-    lambda_event: Any, event_context_extractor: Callable[[Any], Context]
+    lambda_event: Any,
+    event_context_extractor: Callable[[Any], Context],
+    disable_aws_context_propagation: bool = False,
 ) -> Context:
     """Determine the parent context for the current Lambda invocation.
 
     See more:
-    https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md
+    https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#determining-the-parent-of-a-span
 
     Args:
         lambda_event: user-defined, so it could be anything, but this
@@ -154,11 +158,30 @@ def _determine_parent_context(
             Event as input and extracts an OTel Context from it. By default,
             the context is extracted from the HTTP headers of an API Gateway
             request.
+        disable_aws_context_propagation: By default, this instrumentation
+            will try to read the context from the `_X_AMZN_TRACE_ID` environment
+            variable set by Lambda, set this to `True` to disable this behavior.
     Returns:
         A Context with configuration found in the carrier.
     """
     parent_context = None
 
+    if not disable_aws_context_propagation:
+        xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
+
+        if xray_env_var:
+            parent_context = AwsXRayPropagator().extract(
+                {TRACE_HEADER_KEY: xray_env_var}
+            )
+
+    if (
+        parent_context
+        and get_current_span(parent_context)
+        .get_span_context()
+        .trace_flags.sampled
+    ):
+        return parent_context
+
     if event_context_extractor:
         parent_context = event_context_extractor(lambda_event)
     else:
@@ -167,33 +190,6 @@ def _determine_parent_context(
     return parent_context
 
 
-def _determine_links() -> Optional[Sequence[Link]]:
-    """Determine if a Link should be added to the Span based on the
-    environment variable `_X_AMZN_TRACE_ID`.
-
-
-    See more:
-    https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#aws-x-ray-environment-span-link
-
-    Returns:
-        A Link or None
-    """
-    links = None
-
-    xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
-
-    if xray_env_var:
-        env_context = AwsXRayPropagator().extract(
-            {TRACE_HEADER_KEY: xray_env_var}
-        )
-
-        span_context = get_current_span(env_context).get_span_context()
-        if span_context.is_valid:
-            links = [Link(span_context, {"source": "x-ray-env"})]
-
-    return links
-
-
 def _set_api_gateway_v1_proxy_attributes(
     lambda_event: Any, span: Span
 ) -> Span:
@@ -288,6 +284,7 @@ def _instrument(
     flush_timeout,
     event_context_extractor: Callable[[Any], Context],
     tracer_provider: TracerProvider = None,
+    disable_aws_context_propagation: bool = False,
     meter_provider: MeterProvider = None,
 ):
     def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
@@ -300,11 +297,11 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
         lambda_event = args[0]
 
         parent_context = _determine_parent_context(
-            lambda_event, event_context_extractor
+            lambda_event,
+            event_context_extractor,
+            disable_aws_context_propagation,
         )
 
-        links = _determine_links()
-
         span_kind = None
         try:
             if lambda_event["Records"][0]["eventSource"] in {
@@ -330,7 +327,6 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
             name=orig_handler_name,
             context=parent_context,
             kind=span_kind,
-            links=links,
         ) as span:
             if span.is_recording():
                 lambda_context = args[1]
@@ -424,6 +420,9 @@ def _instrument(self, **kwargs):
                     Event as input and extracts an OTel Context from it. By default,
                     the context is extracted from the HTTP headers of an API Gateway
                     request.
+                ``disable_aws_context_propagation``: By default, this instrumentation
+                    will try to read the context from the `_X_AMZN_TRACE_ID` environment
+                    variable set by Lambda, set this to `True` to disable this behavior.
         """
         lambda_handler = os.environ.get(ORIG_HANDLER, os.environ.get(_HANDLER))
         # pylint: disable=attribute-defined-outside-init
@@ -445,6 +444,16 @@ def _instrument(self, **kwargs):
                 flush_timeout_env,
             )
 
+        disable_aws_context_propagation = kwargs.get(
+            "disable_aws_context_propagation", False
+        ) or os.getenv(
+            OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION, "False"
+        ).strip().lower() in (
+            "true",
+            "1",
+            "t",
+        )
+
         _instrument(
             self._wrapped_module_name,
             self._wrapped_function_name,
@@ -453,6 +462,7 @@ def _instrument(self, **kwargs):
                 "event_context_extractor", _default_event_context_extractor
             ),
             tracer_provider=kwargs.get("tracer_provider"),
+            disable_aws_context_propagation=disable_aws_context_propagation,
             meter_provider=kwargs.get("meter_provider"),
         )
 
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py
@@ -27,6 +27,7 @@
     _HANDLER,
     _X_AMZN_TRACE_ID,
     OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT,
+    OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION,
     AwsLambdaInstrumentor,
 )
 from opentelemetry.propagate import get_global_textmap
@@ -137,9 +138,7 @@ def test_active_tracing(self):
         self.assertEqual(len(spans), 1)
         span = spans[0]
         self.assertEqual(span.name, os.environ[_HANDLER])
-        self.assertNotEqual(
-            span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID
-        )
+        self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
         self.assertEqual(span.kind, SpanKind.SERVER)
         self.assertSpanHasAttributes(
             span,
@@ -150,7 +149,11 @@ def test_active_tracing(self):
         )
 
         parent_context = span.parent
-        self.assertEqual(None, parent_context)
+        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()
 
@@ -162,8 +165,11 @@ class TestCase:
             context: Dict
             expected_traceid: int
             expected_parentid: int
+            xray_traceid: str
             expected_state_value: str = None
             expected_trace_state_len: int = 0
+            disable_aws_context_propagation: bool = False
+            disable_aws_context_propagation_envvar: str = ""
 
         def custom_event_context_extractor(lambda_event):
             return get_global_textmap().extract(lambda_event["foo"]["headers"])
@@ -182,9 +188,42 @@ def custom_event_context_extractor(lambda_event):
                 expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
                 expected_trace_state_len=3,
                 expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+            ),
+            TestCase(
+                name="custom_extractor_not_sampled_xray",
+                custom_extractor=custom_event_context_extractor,
+                context={
+                    "foo": {
+                        "headers": {
+                            TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                            TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                        }
+                    }
+                },
+                expected_traceid=MOCK_W3C_TRACE_ID,
+                expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
+                expected_trace_state_len=3,
+                expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+            ),
+            TestCase(
+                name="custom_extractor_sampled_xray",
+                custom_extractor=custom_event_context_extractor,
+                context={
+                    "foo": {
+                        "headers": {
+                            TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                            TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                        }
+                    }
+                },
+                expected_traceid=MOCK_XRAY_TRACE_ID,
+                expected_parentid=MOCK_XRAY_PARENT_SPAN_ID,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
             ),
             TestCase(
-                name="custom_extractor",
+                name="custom_extractor_sampled_xray_disable_aws_propagation",
                 custom_extractor=custom_event_context_extractor,
                 context={
                     "foo": {
@@ -194,24 +233,47 @@ def custom_event_context_extractor(lambda_event):
                         }
                     }
                 },
+                disable_aws_context_propagation=True,
                 expected_traceid=MOCK_W3C_TRACE_ID,
                 expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
                 expected_trace_state_len=3,
                 expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+            ),
+            TestCase(
+                name="no_custom_extractor_xray_disable_aws_propagation_via_env_var",
+                custom_extractor=None,
+                context={
+                    "headers": {
+                        TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME: f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                    }
+                },
+                disable_aws_context_propagation=False,
+                disable_aws_context_propagation_envvar="true",
+                expected_traceid=MOCK_W3C_TRACE_ID,
+                expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
+                expected_trace_state_len=3,
+                expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
             ),
         ]
         for test in tests:
             test_env_patch = mock.patch.dict(
                 "os.environ",
                 {
                     **os.environ,
+                    # NOT Active Tracing
+                    _X_AMZN_TRACE_ID: test.xray_traceid,
+                    OTEL_LAMBDA_DISABLE_AWS_CONTEXT_PROPAGATION: test.disable_aws_context_propagation_envvar,
                     # NOT using the X-Ray Propagator
                     OTEL_PROPAGATORS: "tracecontext",
                 },
             )
             test_env_patch.start()
             AwsLambdaInstrumentor().instrument(
-                event_context_extractor=test.custom_extractor
+                event_context_extractor=test.custom_extractor,
+                disable_aws_context_propagation=test.disable_aws_context_propagation,
             )
             mock_execute_lambda(test.context)
             spans = self.memory_exporter.get_finished_spans()
@@ -239,65 +301,6 @@ def custom_event_context_extractor(lambda_event):
             AwsLambdaInstrumentor().uninstrument()
             test_env_patch.stop()
 
-    def test_links_from_lambda_event(self):
-        @dataclass
-        class TestCase:
-            name: str
-            context: Dict
-            expected_link_trace_id: int
-            expected_link_attributes: dict
-            xray_traceid: str
-
-        tests = [
-            TestCase(
-                name="valid_xray_trace",
-                context={},
-                expected_link_trace_id=MOCK_XRAY_TRACE_ID,
-                expected_link_attributes={"source": "x-ray-env"},
-                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
-            ),
-            TestCase(
-                name="invalid_xray_trace",
-                context={},
-                expected_link_trace_id=None,
-                expected_link_attributes={},
-                xray_traceid="0",
-            ),
-        ]
-        for test in tests:
-            test_env_patch = mock.patch.dict(
-                "os.environ",
-                {
-                    **os.environ,
-                    # NOT Active Tracing
-                    _X_AMZN_TRACE_ID: test.xray_traceid,
-                    # NOT using the X-Ray Propagator
-                    OTEL_PROPAGATORS: "tracecontext",
-                },
-            )
-            test_env_patch.start()
-            AwsLambdaInstrumentor().instrument()
-            mock_execute_lambda(test.context)
-            spans = self.memory_exporter.get_finished_spans()
-            assert spans
-            self.assertEqual(len(spans), 1)
-            span = spans[0]
-
-            if test.expected_link_trace_id is None:
-                self.assertEqual(0, len(span.links))
-            else:
-                link = span.links[0]
-                self.assertEqual(
-                    link.context.trace_id, test.expected_link_trace_id
-                )
-                self.assertEqual(
-                    link.attributes, test.expected_link_attributes
-                )
-
-            self.memory_exporter.clear()
-            AwsLambdaInstrumentor().uninstrument()
-            test_env_patch.stop()
-
     def test_lambda_no_error_with_invalid_flush_timeout(self):
         test_env_patch = mock.patch.dict(
             "os.environ",
diff --git a/tox.ini b/tox.ini
@@ -34,7 +34,7 @@ envlist =
     ; instrumentation-aiopg intentionally excluded from pypy3
 
     ; opentelemetry-instrumentation-aws-lambda
-    py3{7,8,9,10,11}-test-instrumentation-aws-lambda
+    py3{7,8,9}-test-instrumentation-aws-lambda
 
     ; opentelemetry-instrumentation-botocore
     py3{7,8,9,10,11}-test-instrumentation-botocore
