Merge branch 'main' into chore/celery-housekeeping

Author: aabmass

.changelog/4216.fixed

@@ -0,0 +1 @@
+`opentelemetry-instrumentation-aws-lambda`: fix improper handling of header casing

.changelog/4504.fixed

@@ -0,0 +1 @@
+`opentelemetry-instrumentation-celery`: clear completed task ids from `task_id_to_start_time`

.github/workflows/changelog.yml

@@ -30,7 +30,7 @@ jobs:
 
       - name: Ensure no direct changes to CHANGELOG.md
         run: |
-          if [[ $(git diff --name-only FETCH_HEAD -- '**/CHANGELOG.md') ]]
+          if [[ $(git diff --name-only FETCH_HEAD -- 'CHANGELOG.md' '**/CHANGELOG.md') ]]
           then
             echo "CHANGELOG.md files should not be directly modified."
             echo "Please add a changelog fragment file to the appropriate .changelog/ directory instead."

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

@@ -73,6 +73,7 @@ def custom_event_context_extractor(lambda_event):
 from opentelemetry.context.context import Context
 from opentelemetry.instrumentation.aws_lambda.package import _instruments
 from opentelemetry.instrumentation.aws_lambda.version import __version__
+from opentelemetry.instrumentation.cidict import CIDict
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.metrics import MeterProvider, get_meter_provider
@@ -165,7 +166,9 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
         )
     if not isinstance(headers, dict):
         headers = {}
-    return get_global_textmap().extract(headers)
+    return get_global_textmap().extract(
+        CIDict(headers),
+    )
 
 
 def _determine_parent_context(
@@ -205,20 +208,21 @@ def _set_api_gateway_v1_proxy_attributes(
     span.set_attribute(HTTP_METHOD, lambda_event.get("httpMethod"))
 
     if lambda_event.get("headers"):
-        if "User-Agent" in lambda_event["headers"]:
+        headers = CIDict(lambda_event["headers"])
+        if "User-Agent" in headers:
             span.set_attribute(
                 HTTP_USER_AGENT,
-                lambda_event["headers"]["User-Agent"],
+                headers["User-Agent"],
             )
-        if "X-Forwarded-Proto" in lambda_event["headers"]:
+        if "X-Forwarded-Proto" in headers:
             span.set_attribute(
                 HTTP_SCHEME,
-                lambda_event["headers"]["X-Forwarded-Proto"],
+                headers["X-Forwarded-Proto"],
             )
-        if "Host" in lambda_event["headers"]:
+        if "Host" in headers:
             span.set_attribute(
                 NET_HOST_NAME,
-                lambda_event["headers"]["Host"],
+                headers["Host"],
             )
     if "resource" in lambda_event:
         span.set_attribute(HTTP_ROUTE, lambda_event["resource"])

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

@@ -314,6 +314,36 @@ def custom_event_context_extractor(lambda_event):
                 expected_baggage=MOCK_W3C_BAGGAGE_VALUE,
                 propagators="tracecontext,baggage",
             ),
+            TestCase(
+                name="case_insensitive_headers_uppercase",
+                custom_extractor=None,
+                context={
+                    "headers": {
+                        TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME.upper(): MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME.upper(): f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                    }
+                },
+                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="case_insensitive_headers_mixedcase",
+                custom_extractor=None,
+                context={
+                    "headers": {
+                        "TraceParent": MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        "tRaCeStAtE": f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                    }
+                },
+                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,
+            ),
         ]
         for test in tests:
             with self.subTest(test_name=test.name):
@@ -389,6 +419,57 @@ def test_lambda_no_error_with_invalid_flush_timeout(self):
 
         test_env_patch.stop()
 
+    def test_api_gateway_v1_attributes_case_insensitivity(self):
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda(
+            {
+                "httpMethod": "GET",
+                "headers": {
+                    "user-agent": "lowercase-agent",
+                    "host": "lowercase-host",
+                    "x-forwarded-proto": "http",
+                },
+                "resource": "/test",
+                "requestContext": {
+                    "version": "1.0",
+                },
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(
+            span.attributes.get(HTTP_USER_AGENT), "lowercase-agent"
+        )
+        self.assertEqual(span.attributes.get(NET_HOST_NAME), "lowercase-host")
+        self.assertEqual(span.attributes.get(HTTP_SCHEME), "http")
+
+        self.memory_exporter.clear()
+
+        mock_execute_lambda(
+            {
+                "httpMethod": "GET",
+                "headers": {
+                    "uSeR-aGeNt": "mixed-agent",
+                    "hOsT": "mixed-host",
+                    "X-fOrWaRdEd-PrOtO": "https",
+                },
+                "resource": "/test",
+                "requestContext": {
+                    "version": "1.0",
+                },
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.attributes.get(HTTP_USER_AGENT), "mixed-agent")
+        self.assertEqual(span.attributes.get(NET_HOST_NAME), "mixed-host")
+        self.assertEqual(span.attributes.get(HTTP_SCHEME), "https")
+
     def test_lambda_handles_multiple_consumers(self):
         test_env_patch = mock.patch.dict(
             "os.environ",

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py

@@ -114,8 +114,12 @@ def keys(self, carrier: Request) -> list[str]:
 
 
 class CeleryInstrumentor(BaseInstrumentor):
-    metrics = None
-    task_id_to_start_time = {}
+    def __init__(self):
+        super().__init__()
+        if not hasattr(self, "metrics"):
+            self.metrics = None
+        if not hasattr(self, "task_id_to_start_time"):
+            self.task_id_to_start_time = {}
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -138,6 +142,7 @@ def _instrument(self, **kwargs):
             schema_url="https://opentelemetry.io/schemas/1.11.0",
         )
 
+        self.task_id_to_start_time = {}
         self.create_celery_metrics(meter)
 
         signals.task_prerun.connect(self._trace_prerun, weak=False)
@@ -158,6 +163,7 @@ def _uninstrument(self, **kwargs):
         signals.after_task_publish.disconnect(self._trace_after_publish)
         signals.task_failure.disconnect(self._trace_failure)
         signals.task_retry.disconnect(self._trace_retry)
+        self.task_id_to_start_time = {}
 
     def _trace_prerun(self, *args, **kwargs):
         task = utils.retrieve_task(kwargs)
@@ -212,6 +218,7 @@ def _trace_postrun(self, *args, **kwargs):
         self.update_task_duration_time(task_id)
         labels = {"task": task.name, "worker": task.request.hostname}
         self._record_histograms(task_id, labels)
+        self.task_id_to_start_time.pop(task_id, None)
         # if the process sending the task is not instrumented
         # there's no incoming context and no token to detach
         if token is not None:

instrumentation/opentelemetry-instrumentation-celery/tests/test_tasks.py

@@ -92,6 +92,24 @@ def test_task(self):
         self.assertEqual(consumer.parent.span_id, producer.context.span_id)
         self.assertEqual(consumer.context.trace_id, producer.context.trace_id)
 
+    def test_task_clears_start_time_cache(self):
+        """Test that the `task_id_to_start_time` cache is cleared after a task finishes,
+        to prevent memory leaks."""
+        instrumentor = CeleryInstrumentor()
+        instrumentor.instrument()
+
+        result = task_add.delay(1, 2)
+
+        timeout = time.time() + 60 * 1  # 1 minutes from now
+        while not result.ready():
+            if time.time() > timeout:
+                break
+            time.sleep(0.05)
+
+        self.assertTrue(result.ready())
+        self.assertEqual(result.result, 3)
+        self.assertEqual(instrumentor.task_id_to_start_time, {})
+
     def test_task_raises(self):
         CeleryInstrumentor().instrument()
 

opentelemetry-instrumentation/src/opentelemetry/instrumentation/cidict.py

@@ -0,0 +1,81 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import (
+    Any,
+    Iterable,
+    Iterator,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
+
+KT = TypeVar("KT")
+VT = TypeVar("VT")
+
+
+class CIDict(MutableMapping[KT, VT]):
+    def __init__(
+        self,
+        data: Optional[Union[Mapping[KT, VT], Iterable[Tuple[KT, VT]]]] = None,
+    ) -> None:
+        self._data: dict[KT, Tuple[KT, VT]] = {}
+        if data is None:
+            data = {}
+        self.update(data)
+
+    @staticmethod
+    def _normalize_key(key: KT) -> KT:
+        if isinstance(key, str):
+            return key.lower()  # type: ignore
+        return key
+
+    def _get_entry(self, key: KT) -> Tuple[KT, VT]:
+        normalized_key = self._normalize_key(key)
+        if normalized_key in self._data:
+            return self._data[normalized_key]
+        raise KeyError(repr(key))
+
+    def original_key(self, key: KT) -> KT:
+        return self._get_entry(key)[0]
+
+    def normalized_items(self) -> Iterable[Tuple[KT, VT]]:
+        return ((key, value[1]) for key, value in self._data.items())
+
+    def __setitem__(self, key: KT, value: VT, /) -> None:
+        self._data[self._normalize_key(key)] = (key, value)
+
+    def __delitem__(self, key: KT, /) -> None:
+        try:
+            del self._data[self._normalize_key(key)]
+        except KeyError:
+            raise KeyError(repr(key)) from None
+
+    def __getitem__(self, key: KT, /) -> VT:
+        return self._get_entry(key)[1]
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __iter__(self) -> Iterator[KT]:
+        return (key for key, _ in self._data.values())
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({dict(self.items())!r})"
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, CIDict):
+            return dict(self.normalized_items()) == dict(
+                other.normalized_items()
+            )
+        if not isinstance(other, Mapping):
+            return False
+        ciother: CIDict[Any, Any] = CIDict(other)
+        return dict(self.normalized_items()) == dict(
+            ciother.normalized_items()
+        )