Fixing w3c baggage support in opentelemetry-instrumentation-aws-lambda

Author: dtorok

Diff for: instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py

@@ -76,6 +76,7 @@ def custom_event_context_extractor(lambda_event):
 
 from wrapt import wrap_function_wrapper
 
+from opentelemetry import context as context_api
 from opentelemetry.context.context import Context
 from opentelemetry.instrumentation.aws_lambda.package import _instruments
 from opentelemetry.instrumentation.aws_lambda.version import __version__
@@ -303,66 +304,76 @@ def _instrumented_lambda_handler_call(  # noqa pylint: disable=too-many-branches
             schema_url="https://opentelemetry.io/schemas/1.11.0",
         )
 
-        with tracer.start_as_current_span(
-            name=orig_handler_name,
-            context=parent_context,
-            kind=span_kind,
-        ) as span:
-            if span.is_recording():
-                lambda_context = args[1]
-                # NOTE: The specs mention an exception here, allowing the
-                # `SpanAttributes.CLOUD_RESOURCE_ID` attribute to be set as a span
-                # attribute instead of a resource attribute.
-                #
-                # See more:
-                # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/faas/aws-lambda.md#resource-detector
-                span.set_attribute(
-                    SpanAttributes.CLOUD_RESOURCE_ID,
-                    lambda_context.invoked_function_arn,
-                )
-                span.set_attribute(
-                    SpanAttributes.FAAS_INVOCATION_ID,
-                    lambda_context.aws_request_id,
-                )
-
-                # NOTE: `cloud.account.id` can be parsed from the ARN as the fifth item when splitting on `:`
-                #
-                # See more:
-                # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/faas/aws-lambda.md#all-triggers
-                account_id = lambda_context.invoked_function_arn.split(":")[4]
-                span.set_attribute(
-                    ResourceAttributes.CLOUD_ACCOUNT_ID,
-                    account_id,
-                )
+        token = context_api.attach(parent_context)
+        try:
+            with tracer.start_as_current_span(
+                name=orig_handler_name,
+                context=parent_context,
+                kind=span_kind,
+            ) as span:
+                if span.is_recording():
+                    lambda_context = args[1]
+                    # NOTE: The specs mention an exception here, allowing the
+                    # `SpanAttributes.CLOUD_RESOURCE_ID` attribute to be set as a span
+                    # attribute instead of a resource attribute.
+                    #
+                    # See more:
+                    # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/faas/aws-lambda.md#resource-detector
+                    span.set_attribute(
+                        SpanAttributes.CLOUD_RESOURCE_ID,
+                        lambda_context.invoked_function_arn,
+                    )
+                    span.set_attribute(
+                        SpanAttributes.FAAS_INVOCATION_ID,
+                        lambda_context.aws_request_id,
+                    )
 
-            exception = None
-            result = None
-            try:
-                result = call_wrapped(*args, **kwargs)
-            except Exception as exc:  # pylint: disable=W0703
-                exception = exc
-                span.set_status(Status(StatusCode.ERROR))
-                span.record_exception(exception)
-
-            # 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"
-            ):
-                span.set_attribute(SpanAttributes.FAAS_TRIGGER, "http")
-
-                if lambda_event.get("version") == "2.0":
-                    _set_api_gateway_v2_proxy_attributes(lambda_event, span)
-                else:
-                    _set_api_gateway_v1_proxy_attributes(lambda_event, span)
-
-                if isinstance(result, dict) and result.get("statusCode"):
+                    # NOTE: `cloud.account.id` can be parsed from the ARN as the fifth item when splitting on `:`
+                    #
+                    # See more:
+                    # https://github.com/open-telemetry/semantic-conventions/blob/main/docs/faas/aws-lambda.md#all-triggers
+                    account_id = lambda_context.invoked_function_arn.split(
+                        ":"
+                    )[4]
                     span.set_attribute(
-                        SpanAttributes.HTTP_STATUS_CODE,
-                        result.get("statusCode"),
+                        ResourceAttributes.CLOUD_ACCOUNT_ID,
+                        account_id,
                     )
 
+                exception = None
+                result = None
+                try:
+                    result = call_wrapped(*args, **kwargs)
+                except Exception as exc:  # pylint: disable=W0703
+                    exception = exc
+                    span.set_status(Status(StatusCode.ERROR))
+                    span.record_exception(exception)
+
+                # 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"
+                ):
+                    span.set_attribute(SpanAttributes.FAAS_TRIGGER, "http")
+
+                    if lambda_event.get("version") == "2.0":
+                        _set_api_gateway_v2_proxy_attributes(
+                            lambda_event, span
+                        )
+                    else:
+                        _set_api_gateway_v1_proxy_attributes(
+                            lambda_event, span
+                        )
+
+                    if isinstance(result, dict) and result.get("statusCode"):
+                        span.set_attribute(
+                            SpanAttributes.HTTP_STATUS_CODE,
+                            result.get("statusCode"),
+                        )
+        finally:
+            context_api.detach(token)
+
         now = time.time()
         _tracer_provider = tracer_provider or get_tracer_provider()
         if hasattr(_tracer_provider, "force_flush"):

Diff for: instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/lambda_function.py

@@ -12,9 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+
+from opentelemetry import baggage as baggage_api
+
 
 def handler(event, context):
-    return "200 ok"
+    baggage_content = dict(baggage_api.get_all().items())
+    return json.dumps({"baggage_content": baggage_content})
 
 
 def rest_api_handler(event, context):

Diff for: instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py

@@ -11,14 +11,15 @@
 # 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 json
 import os
 from dataclasses import dataclass
 from importlib import import_module, reload
 from typing import Any, Callable, Dict
 from unittest import mock
 
 from opentelemetry import propagate
+from opentelemetry.baggage.propagation import W3CBaggagePropagator
 from opentelemetry.environment_variables import OTEL_PROPAGATORS
 from opentelemetry.instrumentation.aws_lambda import (
     _HANDLER,
@@ -79,6 +80,9 @@ def __init__(self, aws_request_id, invoked_function_arn):
 MOCK_W3C_TRACE_STATE_KEY = "vendor_specific_key"
 MOCK_W3C_TRACE_STATE_VALUE = "test_value"
 
+MOCK_W3C_BAGGAGE_KEY = "baggage_key"
+MOCK_W3C_BAGGAGE_VALUE = "baggage_value"
+
 
 def mock_execute_lambda(event=None):
     """Mocks the AWS Lambda execution.
@@ -97,7 +101,7 @@ def mock_execute_lambda(event=None):
 
     module_name, handler_name = os.environ[_HANDLER].rsplit(".", 1)
     handler_module = import_module(module_name.replace("/", "."))
-    getattr(handler_module, handler_name)(event, MOCK_LAMBDA_CONTEXT)
+    return getattr(handler_module, handler_name)(event, MOCK_LAMBDA_CONTEXT)
 
 
 class TestAwsLambdaInstrumentor(TestBase):
@@ -181,6 +185,9 @@ class TestCase:
             expected_state_value: str = None
             expected_trace_state_len: int = 0
             propagators: str = "tracecontext"
+            expected_baggage: str = None
+            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"])
@@ -266,6 +273,23 @@ def custom_event_context_extractor(lambda_event):
                 expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
                 xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
             ),
+            TestCase(
+                name="baggae_propagation",
+                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",
+                        W3CBaggagePropagator._BAGGAGE_HEADER_NAME: f"{MOCK_W3C_BAGGAGE_KEY}={MOCK_W3C_BAGGAGE_VALUE}",
+                    }
+                },
+                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_baggage=MOCK_W3C_BAGGAGE_VALUE,
+            ),
         ]
         for test in tests:
 
@@ -284,7 +308,9 @@ def custom_event_context_extractor(lambda_event):
             AwsLambdaInstrumentor().instrument(
                 event_context_extractor=test.custom_extractor,
             )
-            mock_execute_lambda(test.context)
+            result = mock_execute_lambda(test.context)
+            result = json.loads(result)
+
             spans = self.memory_exporter.get_finished_spans()
             assert spans
             self.assertEqual(len(spans), 1)
@@ -305,6 +331,10 @@ def custom_event_context_extractor(lambda_event):
                 parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
                 test.expected_state_value,
             )
+            self.assertEqual(
+                result["baggage_content"].get(MOCK_W3C_BAGGAGE_KEY),
+                test.expected_baggage,
+            )
             self.assertTrue(parent_context.is_remote)
             self.memory_exporter.clear()
             AwsLambdaInstrumentor().uninstrument()