fix(aws-lambda): support ALB multiValueHeaders

Author: jbfenton

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

@@ -157,20 +157,40 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
     Returns:
         A Context with configuration found in the event.
     """
-    headers = None
-    try:
-        headers = lambda_event["headers"]
-    except (TypeError, KeyError):
+    headers = _extract_http_headers(lambda_event)
+    if not headers:
         logger.debug(
             "Extracting context from Lambda Event failed: either enable X-Ray active tracing or configure API Gateway to trigger this Lambda function as a pure proxy. Otherwise, generated spans will have an invalid (empty) parent context."
         )
-    if not isinstance(headers, dict):
-        headers = {}
     return get_global_textmap().extract(
         CIDict(headers),
     )
 
 
+def _extract_http_headers(lambda_event: Any) -> dict[str, Any]:
+    try:
+        headers = lambda_event["headers"]
+    except (TypeError, KeyError):
+        headers = None
+
+    if isinstance(headers, dict):
+        return headers
+
+    try:
+        multi_value_headers = lambda_event["multiValueHeaders"]
+    except (TypeError, KeyError):
+        return {}
+
+    if not isinstance(multi_value_headers, dict):
+        return {}
+
+    normalized_headers = {}
+    for key, values in multi_value_headers.items():
+        if isinstance(values, list) and values:
+            normalized_headers[key] = values[0]
+    return normalized_headers
+
+
 def _determine_parent_context(
     lambda_event: Any,
     event_context_extractor: Callable[[Any], Context],
@@ -207,8 +227,9 @@ def _set_api_gateway_v1_proxy_attributes(
     """
     span.set_attribute(HTTP_METHOD, lambda_event.get("httpMethod"))
 
-    if lambda_event.get("headers"):
-        headers = CIDict(lambda_event["headers"])
+    headers = _extract_http_headers(lambda_event)
+    if headers:
+        headers = CIDict(headers)
         if "User-Agent" in headers:
             span.set_attribute(
                 HTTP_USER_AGENT,

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

@@ -3,6 +3,7 @@
 
 import logging
 import os
+from copy import deepcopy
 from dataclasses import dataclass
 from importlib import import_module, reload
 from typing import Any, Callable, Dict
@@ -741,9 +742,48 @@ def test_alb_multi_value_header_event_sets_attributes(self):
             {
                 FAAS_TRIGGER: "http",
                 HTTP_METHOD: "GET",
+                HTTP_SCHEME: "https",
+                HTTP_USER_AGENT: "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6)",
+                NET_HOST_NAME: "lambda-846800462-us-east-2.elb.amazonaws.com",
             },
         )
 
+    def test_alb_multi_value_header_event_extracts_parent_context(self):
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                OTEL_PROPAGATORS: "tracecontext",
+            },
+        )
+        test_env_patch.start()
+        reload(propagate)
+
+        AwsLambdaInstrumentor().instrument()
+
+        event = deepcopy(MOCK_LAMBDA_ALB_MULTI_VALUE_HEADER_EVENT)
+        event["multiValueHeaders"][
+            TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME
+        ] = [MOCK_W3C_TRACE_CONTEXT_SAMPLED]
+
+        mock_execute_lambda(event)
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+
+        span, *_ = spans
+        self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+
+        parent_context = span.parent
+        self.assertEqual(
+            parent_context.trace_id, span.get_span_context().trace_id
+        )
+        self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+        self.assertTrue(parent_context.is_remote)
+
+        test_env_patch.stop()
+
     def test_dynamo_db_event_sets_attributes(self):
         AwsLambdaInstrumentor().instrument()