update awslambda to use _X_AMZN_TRACE_ID as a Span Link

Author: tsloughter

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
+from typing import Any, Callable, Collection, Optional, Sequence
 from urllib.parse import urlencode
 
 from wrapt import wrap_function_wrapper
@@ -90,9 +90,11 @@ 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,
+    get_current_span,
     get_tracer,
     get_tracer_provider,
 )
@@ -149,7 +151,7 @@ def _determine_parent_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#determining-the-parent-of-a-span
+    https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md
 
     Args:
         lambda_event: user-defined, so it could be anything, but this
@@ -166,29 +168,44 @@ def _determine_parent_context(
     """
     parent_context = None
 
+    if event_context_extractor:
+        parent_context = event_context_extractor(lambda_event)
+    else:
+        parent_context = _default_event_context_extractor(lambda_event)
+
+    return parent_context
+
+def _determine_links(
+    disable_aws_context_propagation: bool = False
+) -> 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
+
+    Args:
+        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 Link or None
+    """
+    env_context = None
+    links = 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(
+            env_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:
-        parent_context = _default_event_context_extractor(lambda_event)
-
-    return parent_context
+            links = [Link(get_current_span(env_context).get_span_context())]
 
+    return links
 
 def _set_api_gateway_v1_proxy_attributes(
     lambda_event: Any, span: Span
@@ -294,6 +311,8 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
             disable_aws_context_propagation,
         )
 
+        links = _determine_links(disable_aws_context_propagation)
+
         span_kind = None
         try:
             if lambda_event["Records"][0]["eventSource"] in {
@@ -315,10 +334,13 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
 
         tracer = get_tracer(__name__, __version__, tracer_provider)
 
+
+
         with tracer.start_as_current_span(
             name=orig_handler_name,
             context=parent_context,
             kind=span_kind,
+            links=links
         ) as span:
             if span.is_recording():
                 lambda_context = args[1]

instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py

@@ -138,7 +138,7 @@ def test_active_tracing(self):
         self.assertEqual(len(spans), 1)
         span = spans[0]
         self.assertEqual(span.name, os.environ[_HANDLER])
-        self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
+        self.assertNotEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
         self.assertEqual(span.kind, SpanKind.SERVER)
         self.assertSpanHasAttributes(
             span,
@@ -149,11 +149,7 @@ def test_active_tracing(self):
         )
 
         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)
+        self.assertEqual(None, parent_context)
 
         test_env_patch.stop()
 
@@ -165,11 +161,8 @@ 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"])
@@ -187,11 +180,10 @@ def custom_event_context_extractor(lambda_event):
                 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,
+                expected_state_value=MOCK_W3C_TRACE_STATE_VALUE
             ),
             TestCase(
-                name="custom_extractor_not_sampled_xray",
+                name="custom_extractor",
                 custom_extractor=custom_event_context_extractor,
                 context={
                     "foo": {
@@ -204,76 +196,21 @@ def custom_event_context_extractor(lambda_event):
                 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_sampled_xray_disable_aws_propagation",
-                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",
-                        }
-                    }
-                },
-                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,
-            ),
+                expected_state_value=MOCK_W3C_TRACE_STATE_VALUE
+            )
         ]
         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,
-                disable_aws_context_propagation=test.disable_aws_context_propagation,
+                event_context_extractor=test.custom_extractor
             )
             mock_execute_lambda(test.context)
             spans = self.memory_exporter.get_finished_spans()
@@ -301,6 +238,64 @@ 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
+            xray_traceid: str
+            disable_aws_context_propagation: bool = False
+            disable_aws_context_propagation_envvar: str = ""
+
+        tests = [
+            TestCase(
+                name="enabled_aws_context_propagation",
+                context={},
+                expected_link_trace_id=MOCK_XRAY_TRACE_ID,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+            ),
+            TestCase(
+                name="disable_aws_context_propagation",
+                context={},
+                expected_link_trace_id=None,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                disable_aws_context_propagation=True
+            )
+
+        ]
+        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(
+                disable_aws_context_propagation=test.disable_aws_context_propagation,
+            )
+            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 == None:
+                self.assertEqual(0, len(span.links))
+            else:
+                link = span.links[0]
+                self.assertEqual(link.context.trace_id, test.expected_link_trace_id)
+
+            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",
@@ -327,6 +322,7 @@ def test_lambda_no_error_with_invalid_flush_timeout(self):
 
         test_env_patch.stop()
 
+
     def test_lambda_handles_multiple_consumers(self):
         test_env_patch = mock.patch.dict(
             "os.environ",