Enable global propagator for AWS instrumentation (#2599)

* Enable global propagator for AWS instrumentation Fixes #2598 * Add entry point test case * Update instrumentation/opentelemetry-instrumentation-aws-lambda/pyproject.toml Co-authored-by: Riccardo Magliocchetti <[email protected]> * Add test for propagator * Fix entry point name * Update instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py Co-authored-by: Emídio Neto <[email protected]> * Try with ubuntu latest * Try with 24.04 * Fix propagator key * Fix lint * Revert ununtuns --------- Co-authored-by: Riccardo Magliocchetti <[email protected]> Co-authored-by: Emídio Neto <[email protected]>
open-telemetry · Jul 24, 2024 · a47810c · a47810c
1 parent e799a74
commit a47810c
Show file tree

Hide file tree

Showing 9 changed files with 66 additions and 86 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `opentelemetry-instrumentation-flask` Add `http.route` and `http.target` to metric attributes
   ([#2621](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2621))
+- `opentelemetry-instrumentation-aws-lambda` Enable global propagator for AWS instrumentation
+  ([#2708](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2708))
 - `opentelemetry-instrumentation-sklearn` Deprecated the sklearn instrumentation
   ([#2708](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2708))
 - `opentelemetry-instrumentation-pyramid` Record exceptions raised when serving a request

diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/pyproject.toml b/instrumentation/opentelemetry-instrumentation-aws-lambda/pyproject.toml
@@ -33,6 +33,9 @@ dependencies = [
 [project.optional-dependencies]
 instruments = []
 
+[project.entry-points.opentelemetry_instrumentor]
+aws-lambda = "opentelemetry.instrumentation.aws_lambda:AwsLambdaInstrumentor"
+
 [project.urls]
 Homepage = "https://github.com/open-telemetry/opentelemetry-python-contrib/tree/main/instrumentation/opentelemetry-instrumentation-aws-lambda"
 

diff --git a/...metry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py b/...metry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py
@@ -83,10 +83,6 @@ def custom_event_context_extractor(lambda_event):
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.metrics import MeterProvider, get_meter_provider
 from opentelemetry.propagate import get_global_textmap
-from opentelemetry.propagators.aws.aws_xray_propagator import (
-    TRACE_HEADER_KEY,
-    AwsXRayPropagator,
-)
 from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import (
@@ -96,7 +92,6 @@ def custom_event_context_extractor(lambda_event):
     get_tracer,
     get_tracer_provider,
 )
-from opentelemetry.trace.propagation import get_current_span
 from opentelemetry.trace.status import Status, StatusCode
 
 logger = logging.getLogger(__name__)
@@ -107,9 +102,6 @@ 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:
@@ -145,7 +137,6 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
 def _determine_parent_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.
 
@@ -159,36 +150,14 @@ 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 is None:
+        return _default_event_context_extractor(lambda_event)
 
-    if event_context_extractor:
-        parent_context = event_context_extractor(lambda_event)
-    else:
-        parent_context = _default_event_context_extractor(lambda_event)
-
-    return parent_context
+    return event_context_extractor(lambda_event)
 
 
 def _set_api_gateway_v1_proxy_attributes(
@@ -286,14 +255,15 @@ def _instrument(
     flush_timeout,
     event_context_extractor: Callable[[Any], Context],
     tracer_provider: TracerProvider = None,
-    disable_aws_context_propagation: bool = False,
     meter_provider: MeterProvider = None,
 ):
+
     # pylint: disable=too-many-locals
     # pylint: disable=too-many-statements
     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]
         )
@@ -303,7 +273,6 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
         parent_context = _determine_parent_context(
             lambda_event,
             event_context_extractor,
-            disable_aws_context_propagation,
         )
 
         try:
@@ -451,9 +420,6 @@ 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
@@ -475,16 +441,6 @@ 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,
@@ -493,7 +449,6 @@ 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/...ntation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation.py b/...ntation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation.py
diff --git a/.../opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py b/.../opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py
@@ -11,18 +11,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import os
 from dataclasses import dataclass
-from importlib import import_module
+from importlib import import_module, reload
 from typing import Any, Callable, Dict
 from unittest import mock
 
+from opentelemetry import propagate
 from opentelemetry.environment_variables import OTEL_PROPAGATORS
 from opentelemetry.instrumentation.aws_lambda import (
     _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
@@ -37,6 +38,7 @@
 from opentelemetry.trace.propagation.tracecontext import (
     TraceContextTextMapPropagator,
 )
+from opentelemetry.util._importlib_metadata import entry_points
 
 from .mocks.api_gateway_http_api_event import (
     MOCK_LAMBDA_API_GATEWAY_HTTP_API_EVENT,
@@ -56,6 +58,7 @@ def __init__(self, aws_request_id, invoked_function_arn):
 )
 
 MOCK_XRAY_TRACE_ID = 0x5FB7331105E8BB83207FA31D4D9CDB4C
+
 MOCK_XRAY_TRACE_ID_STR = f"{MOCK_XRAY_TRACE_ID:x}"
 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}"
@@ -81,6 +84,7 @@ def mock_execute_lambda(event=None):
     """Mocks the AWS Lambda execution.
 
     NOTE: We don't use `moto`'s `mock_lambda` because we are not instrumenting
+
     calls to AWS Lambda using the AWS SDK. Instead, we are instrumenting AWS
     Lambda itself.
 
@@ -122,10 +126,13 @@ def test_active_tracing(self):
             {
                 **os.environ,
                 # Using Active tracing
+                OTEL_PROPAGATORS: "xray-lambda",
                 _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
             },
         )
+
         test_env_patch.start()
+        reload(propagate)
 
         AwsLambdaInstrumentor().instrument()
 
@@ -173,8 +180,7 @@ class TestCase:
             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 = ""
+            propagators: str = "tracecontext"
 
         def custom_event_context_extractor(lambda_event):
             return get_global_textmap().extract(lambda_event["foo"]["headers"])
@@ -226,9 +232,10 @@ def custom_event_context_extractor(lambda_event):
                 expected_traceid=MOCK_XRAY_TRACE_ID,
                 expected_parentid=MOCK_XRAY_PARENT_SPAN_ID,
                 xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
+                propagators="xray-lambda",
             ),
             TestCase(
-                name="custom_extractor_sampled_xray_disable_aws_propagation",
+                name="custom_extractor_sampled_xray",
                 custom_extractor=custom_event_context_extractor,
                 context={
                     "foo": {
@@ -238,24 +245,21 @@ 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",
+                name="no_custom_extractor_xray",
                 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,
@@ -264,21 +268,21 @@ def custom_event_context_extractor(lambda_event):
             ),
         ]
         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",
+                    OTEL_PROPAGATORS: test.propagators,
                 },
             )
             test_env_patch.start()
+            reload(propagate)
+
             AwsLambdaInstrumentor().instrument(
                 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()
@@ -374,6 +378,7 @@ def test_lambda_handles_invalid_event_source(self):
             },
         )
         test_env_patch.start()
+        reload(propagate)
 
         AwsLambdaInstrumentor().instrument()
 
@@ -513,3 +518,15 @@ def test_no_op_tracer_provider(self):
         spans = self.memory_exporter.get_finished_spans()
         assert spans is not None
         self.assertEqual(len(spans), 0)
+
+    def test_load_entry_point(self):
+        self.assertIs(
+            next(
+                iter(
+                    entry_points(
+                        group="opentelemetry_instrumentor", name="aws-lambda"
+                    )
+                )
+            ).load(),
+            AwsLambdaInstrumentor,
+        )
diff --git a/propagator/opentelemetry-propagator-aws-xray/pyproject.toml b/propagator/opentelemetry-propagator-aws-xray/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
 
 [project.entry-points.opentelemetry_propagator]
 xray = "opentelemetry.propagators.aws:AwsXRayPropagator"
-xray_lambda = "opentelemetry.propagators.aws:AwsXRayLambdaPropagator"
+xray-lambda = "opentelemetry.propagators.aws:AwsXRayLambdaPropagator"
 
 [project.urls]
 Homepage = "https://github.com/open-telemetry/opentelemetry-python-contrib/tree/main/propagator/opentelemetry-propagator-aws-xray"

diff --git a/propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/__init__.py b/propagator/opentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/__init__.py
@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from opentelemetry.propagators.aws.aws_xray_propagator import AwsXRayPropagator
+from opentelemetry.propagators.aws.aws_xray_propagator import (
+    AwsXRayLambdaPropagator,
+    AwsXRayPropagator,
+)
 
-__all__ = ["AwsXRayPropagator"]
+__all__ = ["AwsXRayPropagator", "AwsXRayLambdaPropagator"]
diff --git a/...pentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py b/...pentelemetry-propagator-aws-xray/src/opentelemetry/propagators/aws/aws_xray_propagator.py
@@ -328,9 +328,9 @@ def fields(self):
         return {TRACE_HEADER_KEY}
 
 
-class AwsXrayLambdaPropagator(AwsXRayPropagator):
+class AwsXRayLambdaPropagator(AwsXRayPropagator):
     """Implementation of the AWS X-Ray Trace Header propagation protocol but
-    with special handling for Lambda's ``_X_AMZN_TRACE_ID` environment
+    with special handling for Lambda's ``_X_AMZN_TRACE_ID`` environment
     variable.
     """