Feat: Use the lambda runtime as parent context as opt-in

In a previous PR open-telemetry#1657 the lambda behaviour was changed so that it
extracs the context from the headers of the lambda event instead of from
the lambda runtime, using the _X_AMZN_TRACE_ID env var.

This behaviour is undesireble as it is a breaking change to existing
users.

This PR will add the previous behaviour through an opt-in flag so that
users can be gracefully migrated or these two "modes" of operation can
exist without conflict.

Signed-off-by: Raphael Silva <[email protected]>

Author: rapphil

instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py

@@ -107,6 +107,7 @@ def custom_event_context_extractor(lambda_event):
 OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT = (
     "OTEL_INSTRUMENTATION_AWS_LAMBDA_FLUSH_TIMEOUT"
 )
+OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION = "OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION"
 
 
 def _default_event_context_extractor(lambda_event: Any) -> Context:
@@ -140,7 +141,9 @@ 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,
+    use_aws_context_propagation: bool,
+    event_context_extractor: Callable[[Any], Context],
 ) -> Context:
     """Determine the parent context for the current Lambda invocation.
 
@@ -158,15 +161,26 @@ def _determine_parent_context(
         A Context with configuration found in the carrier.
     """
     parent_context = None
-
-    if event_context_extractor:
+    if use_aws_context_propagation:
+        parent_context = _get_x_ray_context()
+    elif event_context_extractor:
         parent_context = event_context_extractor(lambda_event)
     else:
         parent_context = _default_event_context_extractor(lambda_event)
 
     return parent_context
 
 
+def _get_x_ray_context() -> Optional[Context]:
+    """Determine teh context propagated through the lambda runtime"""
+    xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
+    if xray_env_var:
+        env_context = AwsXRayPropagator().extract({TRACE_HEADER_KEY: xray_env_var})
+        return env_context
+
+    return None
+
+
 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`.
@@ -180,31 +194,23 @@ def _determine_links() -> Optional[Sequence[Link]]:
     """
     links = None
 
-    xray_env_var = os.environ.get(_X_AMZN_TRACE_ID)
+    x_ray_context = _get_x_ray_context()
 
-    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 x_ray_context:
+        span_context = get_current_span(x_ray_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:
+def _set_api_gateway_v1_proxy_attributes(lambda_event: Any, span: Span) -> Span:
     """Sets HTTP attributes for REST APIs and v1 HTTP APIs
 
     More info:
     https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format
     """
-    span.set_attribute(
-        SpanAttributes.HTTP_METHOD, lambda_event.get("httpMethod")
-    )
+    span.set_attribute(SpanAttributes.HTTP_METHOD, lambda_event.get("httpMethod"))
 
     if lambda_event.get("headers"):
         if "User-Agent" in lambda_event["headers"]:
@@ -231,16 +237,12 @@ def _set_api_gateway_v1_proxy_attributes(
                 f"{lambda_event['resource']}?{urlencode(lambda_event['queryStringParameters'])}",
             )
         else:
-            span.set_attribute(
-                SpanAttributes.HTTP_TARGET, lambda_event["resource"]
-            )
+            span.set_attribute(SpanAttributes.HTTP_TARGET, lambda_event["resource"])
 
     return span
 
 
-def _set_api_gateway_v2_proxy_attributes(
-    lambda_event: Any, span: Span
-) -> Span:
+def _set_api_gateway_v2_proxy_attributes(lambda_event: Any, span: Span) -> Span:
     """Sets HTTP attributes for v2 HTTP APIs
 
     More info:
@@ -289,21 +291,26 @@ def _instrument(
     event_context_extractor: Callable[[Any], Context],
     tracer_provider: TracerProvider = None,
     meter_provider: MeterProvider = None,
+    use_aws_context_propagation: bool = False,
 ):
     def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
         call_wrapped, instance, args, kwargs
     ):
-        orig_handler_name = ".".join(
-            [wrapped_module_name, wrapped_function_name]
-        )
+        orig_handler_name = ".".join([wrapped_module_name, wrapped_function_name])
 
         lambda_event = args[0]
 
+        # We are not fully complying with the specification here to be backwards
+        # compatible with the old version of the specification.
+        # the ``use_aws_context_propagation`` flag allow us
+        # to opt-in into the previous behavior
         parent_context = _determine_parent_context(
-            lambda_event, event_context_extractor
+            lambda_event, use_aws_context_propagation, event_context_extractor
         )
 
-        links = _determine_links()
+        links = None
+        if not use_aws_context_propagation:
+            links = _determine_links()
 
         span_kind = None
         try:
@@ -354,9 +361,7 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
             # If the request came from an API Gateway, extract http attributes from the event
             # https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#api-gateway
             # https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-server-semantic-conventions
-            if isinstance(lambda_event, dict) and lambda_event.get(
-                "requestContext"
-            ):
+            if isinstance(lambda_event, dict) and lambda_event.get("requestContext"):
                 span.set_attribute(SpanAttributes.FAAS_TRIGGER, "http")
 
                 if lambda_event.get("version") == "2.0":
@@ -424,6 +429,13 @@ 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.
+                ``use_aws_context_propagation``: whether to use the AWS context propagation
+                    to populate the parent context. When set to true, the spans
+                    from the lambda runtime will not be added as span link to the
+                    span of the lambda invocation.
+                    Defaults to False.
+        """
+        _instrument(**kwargs)
         """
         lambda_handler = os.environ.get(ORIG_HANDLER, os.environ.get(_HANDLER))
         # pylint: disable=attribute-defined-outside-init
@@ -454,6 +466,11 @@ def _instrument(self, **kwargs):
             ),
             tracer_provider=kwargs.get("tracer_provider"),
             meter_provider=kwargs.get("meter_provider"),
+            use_aws_context_propagation=kwargs.get(
+                "use_aws_context_propagation",
+                os.environ.get(OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION, "False").lower()
+                in ("true", "1", "t"),
+            ),
         )
 
     def _uninstrument(self, **kwargs):

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_USE_AWS_CONTEXT_PROPAGATION,
     AwsLambdaInstrumentor,
 )
 from opentelemetry.propagate import get_global_textmap
@@ -59,9 +60,7 @@ def __init__(self, aws_request_id, invoked_function_arn):
 MOCK_XRAY_PARENT_SPAN_ID = 0x3328B8445A6DBAD2
 MOCK_XRAY_TRACE_CONTEXT_COMMON = f"Root={TRACE_ID_VERSION}-{MOCK_XRAY_TRACE_ID_STR[:TRACE_ID_FIRST_PART_LENGTH]}-{MOCK_XRAY_TRACE_ID_STR[TRACE_ID_FIRST_PART_LENGTH:]};Parent={MOCK_XRAY_PARENT_SPAN_ID:x}"
 MOCK_XRAY_TRACE_CONTEXT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1"
-MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = (
-    f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0"
-)
+MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0"
 
 # See more:
 # https://www.w3.org/TR/trace-context/#examples-of-http-traceparent-headers
@@ -116,43 +115,70 @@ def tearDown(self):
         AwsLambdaInstrumentor().uninstrument()
 
     def test_active_tracing(self):
-        test_env_patch = mock.patch.dict(
-            "os.environ",
-            {
-                **os.environ,
-                # Using Active tracing
-                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
-            },
-        )
-        test_env_patch.start()
+        @dataclass
+        class Testcase:
+            name: str
+            use_aws_context_propagation: str
+            expected_trace_id: str
 
-        AwsLambdaInstrumentor().instrument()
+        tests = [
+            Testcase(
+                name="Use aws context propgation",
+                use_aws_context_propagation="true",
+                expected_trace_id=MOCK_XRAY_TRACE_ID,
+            ),
+            Testcase(
+                name="Do not use aws context propgation",
+                use_aws_context_propagation="false",
+                expected_trace_id=None,
+            ),
+        ]
 
-        mock_execute_lambda()
+        for test in tests:
+            test_env_patch = mock.patch.dict(
+                "os.environ",
+                {
+                    **os.environ,
+                    # Using Active tracing
+                    _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+                    OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION: test.use_aws_context_propagation,
+                },
+            )
+            test_env_patch.start()
 
-        spans = self.memory_exporter.get_finished_spans()
+            AwsLambdaInstrumentor().instrument()
 
-        assert spans
+            mock_execute_lambda()
 
-        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.kind, SpanKind.SERVER)
-        self.assertSpanHasAttributes(
-            span,
-            {
-                ResourceAttributes.FAAS_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
-                SpanAttributes.FAAS_EXECUTION: MOCK_LAMBDA_CONTEXT.aws_request_id,
-            },
-        )
+            spans = self.memory_exporter.get_finished_spans()
 
-        parent_context = span.parent
-        self.assertEqual(None, parent_context)
+            assert spans
 
-        test_env_patch.stop()
+            self.assertEqual(len(spans), 1)
+            span = spans[0]
+            self.assertEqual(span.name, os.environ[_HANDLER])
+            parent_context = span.parent
+            if test.expected_trace_id is None:
+                self.assertNotEqual(
+                    span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID
+                )
+                self.assertEqual(None, parent_context)
+            else:
+                self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
+                self.assertEqual(
+                    parent_context.trace_id, span.get_span_context().trace_id
+                )
+            self.assertEqual(span.kind, SpanKind.SERVER)
+            self.assertSpanHasAttributes(
+                span,
+                {
+                    ResourceAttributes.FAAS_ID: MOCK_LAMBDA_CONTEXT.invoked_function_arn,
+                    SpanAttributes.FAAS_EXECUTION: MOCK_LAMBDA_CONTEXT.aws_request_id,
+                },
+            )
+            self.memory_exporter.clear()
+            AwsLambdaInstrumentor().uninstrument()
+            test_env_patch.stop()
 
     def test_parent_context_from_lambda_event(self):
         @dataclass
@@ -218,14 +244,10 @@ def custom_event_context_extractor(lambda_event):
             assert spans
             self.assertEqual(len(spans), 1)
             span = spans[0]
-            self.assertEqual(
-                span.get_span_context().trace_id, test.expected_traceid
-            )
+            self.assertEqual(span.get_span_context().trace_id, test.expected_traceid)
 
             parent_context = span.parent
-            self.assertEqual(
-                parent_context.trace_id, span.get_span_context().trace_id
-            )
+            self.assertEqual(parent_context.trace_id, span.get_span_context().trace_id)
             self.assertEqual(parent_context.span_id, test.expected_parentid)
             self.assertEqual(
                 len(parent_context.trace_state), test.expected_trace_state_len
@@ -247,6 +269,7 @@ class TestCase:
             expected_link_trace_id: int
             expected_link_attributes: dict
             xray_traceid: str
+            use_xray_propagator: str
 
         tests = [
             TestCase(
@@ -255,13 +278,23 @@ class TestCase:
                 expected_link_trace_id=MOCK_XRAY_TRACE_ID,
                 expected_link_attributes={"source": "x-ray-env"},
                 xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+                use_xray_propagator="false",
             ),
             TestCase(
                 name="invalid_xray_trace",
                 context={},
                 expected_link_trace_id=None,
                 expected_link_attributes={},
                 xray_traceid="0",
+                use_xray_propagator="false",
+            ),
+            TestCase(
+                name="use_xra",
+                context={},
+                expected_link_trace_id=None,
+                expected_link_attributes={},
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+                use_xray_propagator="true",
             ),
         ]
         for test in tests:
@@ -273,6 +306,7 @@ class TestCase:
                     _X_AMZN_TRACE_ID: test.xray_traceid,
                     # NOT using the X-Ray Propagator
                     OTEL_PROPAGATORS: "tracecontext",
+                    OTEL_LAMBDA_USE_AWS_CONTEXT_PROPAGATION: test.use_xray_propagator,
                 },
             )
             test_env_patch.start()
@@ -287,12 +321,8 @@ class TestCase:
                 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.assertEqual(link.context.trace_id, test.expected_link_trace_id)
+                self.assertEqual(link.attributes, test.expected_link_attributes)
 
             self.memory_exporter.clear()
             AwsLambdaInstrumentor().uninstrument()