attempt at fixing tests (done with LLM)

Author: FrenchGithubUser

synapse/handlers/stats.py

@@ -72,6 +72,10 @@ def __init__(self, hs: "HomeServer"):
         # Guard to ensure we only process deltas one at a time
         self._is_processing = False
 
+        # Initialize room count metrics to 0
+        known_rooms_gauge.set(0, {SERVER_NAME_LABEL: self.server_name})
+        locally_joined_rooms_gauge.set(0, {SERVER_NAME_LABEL: self.server_name})
+
         if self.stats_enabled and hs.config.worker.run_background_tasks:
             self.notifier.add_replication_callback(self.notify_new_event)
 

synapse/util/batching_queue.py

@@ -20,6 +20,7 @@
 #
 
 import logging
+import weakref
 from typing import (
     TYPE_CHECKING,
     Awaitable,
@@ -46,13 +47,66 @@
 V = TypeVar("V")
 R = TypeVar("R")
 
-# number_queued = meter.create_observable_gauge(
-#     "synapse_util_batching_queue_number_queued",
-#     description="The number of items waiting in the queue across all keys",
-# )
+# Global registry to track all BatchingQueue instances.
+# We use a WeakSet so that queues can be garbage collected when no longer referenced.
+_batching_queue_registry: "weakref.WeakSet[BatchingQueue]" = weakref.WeakSet()
+
+
+def _collect_number_queued(options: object) -> list[Observation]:
+    """Callback to collect number_queued metrics from all BatchingQueue instances."""
+    observations = []
+    for queue in _batching_queue_registry:
+        observations.append(
+            Observation(
+                sum(len(q) for q in queue._next_values.values()),
+                {"name": queue._name, SERVER_NAME_LABEL: queue.server_name},
+            )
+        )
+    return observations
+
+
+def _collect_number_of_keys(options: object) -> list[Observation]:
+    """Callback to collect number_of_keys metrics from all BatchingQueue instances."""
+    observations = []
+    for queue in _batching_queue_registry:
+        observations.append(
+            Observation(
+                len(queue._next_values),
+                {"name": queue._name, SERVER_NAME_LABEL: queue.server_name},
+            )
+        )
+    return observations
+
+
+def _collect_number_in_flight(options: object) -> list[Observation]:
+    """Callback to collect number_in_flight metrics from all BatchingQueue instances."""
+    observations = []
+    for queue in _batching_queue_registry:
+        observations.append(
+            Observation(
+                queue._number_in_flight,
+                {"name": queue._name, SERVER_NAME_LABEL: queue.server_name},
+            )
+        )
+    return observations
+
+
+# Global observable gauges that collect from all BatchingQueue instances
+number_queued = meter.create_observable_gauge(
+    "synapse_util_batching_queue_number_queued",
+    callbacks=[_collect_number_queued],
+    description="The number of items waiting in the queue across all keys",
+)
+
+number_of_keys = meter.create_observable_gauge(
+    "synapse_util_batching_queue_number_of_keys",
+    callbacks=[_collect_number_of_keys],
+    description="The number of distinct keys that have items queued",
+)
 
 number_in_flight = meter.create_observable_gauge(
     "synapse_util_batching_queue_number_pending",
+    callbacks=[_collect_number_in_flight],
     description="The number of items across all keys either being processed or waiting in a queue",
 )
 
@@ -107,45 +161,19 @@ def __init__(
         # The function to call with batches of values.
         self._process_batch_callback = process_batch_callback
 
-        self.number_queued = meter.create_observable_gauge(
-            "synapse_util_batching_queue_number_queued",
-            callbacks=[
-                lambda options: [
-                    Observation(
-                        sum(len(q) for q in self._next_values.values()),
-                        {"name": self._name, SERVER_NAME_LABEL: self.server_name},
-                    )
-                ]
-            ],
-            description="The number of items waiting in the queue across all keys",
-        )
+        # Counter for number of items in flight (being processed or waiting).
+        self._number_in_flight: int = 0
 
-        self.number_of_keys = meter.create_observable_gauge(
-            "synapse_util_batching_queue_number_of_keys",
-            description="The number of distinct keys that have items queued",
-            callbacks=[
-                lambda options: [
-                    Observation(
-                        len(self._next_values),
-                        {"name": self._name, SERVER_NAME_LABEL: self.server_name},
-                    )
-                ]
-            ],
-        )
-
-        self._number_in_flight_metric = meter.create_up_down_counter(
-            "synapse_util_batching_queue_number_pending",
-            description="The number of items across all keys either being processed or waiting in a queue",
-        )
+        # Register this instance with the global registry so metrics can be collected.
+        _batching_queue_registry.add(self)
 
     def shutdown(self) -> None:
         """
         Prepares the object for garbage collection by removing any handed out
         references.
         """
-        # there doesn't seem to be an otel equivalent for those
-        # number_queued.remove(self._name, self.server_name)
-        # number_of_keys.remove(self._name, self.server_name)
+        # The global registry uses WeakSet, so instances are automatically
+        # removed when garbage collected. No explicit cleanup needed.
 
     async def add_to_queue(self, value: V, key: Hashable = ()) -> R:
         """Adds the value to the queue with the given key, returning the result
@@ -167,13 +195,11 @@ async def add_to_queue(self, value: V, key: Hashable = ()) -> R:
         if key not in self._processing_keys:
             self.hs.run_as_background_process(self._name, self._process_queue, key)
 
-        self._number_in_flight_metric.add(
-            1, {"name": self._name, SERVER_NAME_LABEL: self.server_name}
-        )
-        res = await make_deferred_yieldable(d)
-        self._number_in_flight_metric.add(
-            -1, {"name": self._name, SERVER_NAME_LABEL: self.server_name}
-        )
+        self._number_in_flight += 1
+        try:
+            res = await make_deferred_yieldable(d)
+        finally:
+            self._number_in_flight -= 1
         return res
 
     async def _process_queue(self, key: Hashable) -> None:

tests/handlers/test_stats.py

@@ -57,16 +57,50 @@ def _set_metrics_to_zero(self) -> None:
         This method resets the metrics to zero before each test to ensure
         that each test starts with a clean slate.
         """
-        from opentelemetry import metrics
-        from opentelemetry.sdk.metrics import MeterProvider
-        from opentelemetry.sdk.metrics.export import InMemoryMetricReader
-
-        # Create a fresh reader and provider
-        self.reader = InMemoryMetricReader()
-        provider = MeterProvider(metric_readers=[self.reader])
-        # Set the global provider
-        # Any new metric instruments created after this will use the clean state.
-        metrics.set_meter_provider(provider)
+        from synapse.metrics import (
+            SERVER_NAME_LABEL,
+            known_rooms_gauge,
+            locally_joined_rooms_gauge,
+        )
+
+        # Reset the gauge values to 0 for this server
+        known_rooms_gauge.set(0, {SERVER_NAME_LABEL: self.hs.hostname})
+        locally_joined_rooms_gauge.set(0, {SERVER_NAME_LABEL: self.hs.hostname})
+
+    def _get_gauge_values(
+        self, metrics: list[tuple[str, dict[str, str]]]
+    ) -> list[Optional[float]]:
+        """
+        Get multiple gauge values from the Prometheus registry in a single call.
+
+        The standard REGISTRY.get_sample_value() doesn't work for OpenTelemetry
+        metrics because the OTel exporter doesn't register its metric names.
+        Additionally, the OTel collector only returns data on the first collect()
+        call, so we must collect all data once and then look up all values.
+
+        Args:
+            metrics: List of (metric_name, labels) tuples to look up.
+
+        Returns:
+            List of values in the same order as the input metrics.
+        """
+        # Collect all data from all collectors into a lookup dict
+        all_samples: dict[tuple[str, tuple[tuple[str, str], ...]], float] = {}
+        for collector in REGISTRY._collector_to_names.keys():
+            try:
+                for metric_family in collector.collect():
+                    for sample in metric_family.samples:
+                        key = (metric_family.name, tuple(sorted(sample.labels.items())))
+                        all_samples[key] = sample.value
+            except Exception:
+                continue
+
+        # Look up each requested metric
+        results: list[Optional[float]] = []
+        for metric_name, labels in metrics:
+            key = (metric_name, tuple(sorted(labels.items())))
+            results.append(all_samples.get(key))
+        return results
 
     def _add_background_updates(self) -> None:
         """
@@ -184,19 +218,18 @@ def test_create_room(self) -> None:
         When we create a room, it should have statistics already ready.
         """
         self._perform_background_initial_update()
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_known_rooms_total", labels={"server_name": self.hs.hostname}
-            ),
-            0.0,
-        )
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_locally_joined_rooms_total",
-                labels={"server_name": self.hs.hostname},
-            ),
-            0.0,
-        )
+        known_rooms, locally_joined = self._get_gauge_values(
+            [
+                ("synapse_known_rooms_total", {"server_name": self.hs.hostname}),
+                (
+                    "synapse_locally_joined_rooms_total",
+                    {"server_name": self.hs.hostname},
+                ),
+            ]
+        )
+        self.assertEqual(known_rooms, 0.0)
+        self.assertEqual(locally_joined, 0.0)
+
         u1 = self.register_user("u1", "pass")
         u1token = self.login("u1", "pass")
         r1 = self.helper.create_room_as(u1, tok=u1token)
@@ -223,19 +256,17 @@ def test_create_room(self) -> None:
         self.assertEqual(r2stats["banned_members"], 0)
 
         # There are 2 rooms created. Check the room metrics were udpated.
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_known_rooms_total", labels={"server_name": self.hs.hostname}
-            ),
-            2,
-        )
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_locally_joined_rooms_total",
-                labels={"server_name": self.hs.hostname},
-            ),
-            2,
-        )
+        known_rooms, locally_joined = self._get_gauge_values(
+            [
+                ("synapse_known_rooms_total", {"server_name": self.hs.hostname}),
+                (
+                    "synapse_locally_joined_rooms_total",
+                    {"server_name": self.hs.hostname},
+                ),
+            ]
+        )
+        self.assertEqual(known_rooms, 2)
+        self.assertEqual(locally_joined, 2)
 
     def test_updating_profile_information_does_not_increase_joined_members_count(
         self,
@@ -647,19 +678,17 @@ def test_room_metrics(self) -> None:
         """
 
         self._perform_background_initial_update()
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_known_rooms_total", labels={"server_name": self.hs.hostname}
-            ),
-            0.0,
-        )
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_locally_joined_rooms_total",
-                labels={"server_name": self.hs.hostname},
-            ),
-            0.0,
-        )
+        known_rooms, locally_joined = self._get_gauge_values(
+            [
+                ("synapse_known_rooms_total", {"server_name": self.hs.hostname}),
+                (
+                    "synapse_locally_joined_rooms_total",
+                    {"server_name": self.hs.hostname},
+                ),
+            ]
+        )
+        self.assertEqual(known_rooms, 0.0)
+        self.assertEqual(locally_joined, 0.0)
 
         u1 = self.register_user("u1", "pass")
         u1token = self.login("u1", "pass")
@@ -670,19 +699,17 @@ def test_room_metrics(self) -> None:
         self.helper.leave(r2, u1, tok=u1token)
 
         # Check the locally joined rooms metric after creating rooms
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_locally_joined_rooms_total",
-                labels={"server_name": self.hs.hostname},
-            ),
-            1,
-        )
-        self.assertEqual(
-            REGISTRY.get_sample_value(
-                "synapse_known_rooms_total", labels={"server_name": self.hs.hostname}
-            ),
-            2,
-        )
+        known_rooms, locally_joined = self._get_gauge_values(
+            [
+                ("synapse_known_rooms_total", {"server_name": self.hs.hostname}),
+                (
+                    "synapse_locally_joined_rooms_total",
+                    {"server_name": self.hs.hostname},
+                ),
+            ]
+        )
+        self.assertEqual(locally_joined, 1)
+        self.assertEqual(known_rooms, 2)
 
         # Check the stats for both rooms
         r1stats = self._get_current_stats("room", r1)

tests/util/test_batching_queue.py

@@ -28,6 +28,8 @@
 from synapse.util.batching_queue import (
     BatchingQueue,
     number_in_flight,
+    number_of_keys,
+    number_queued,
 )
 
 from tests.unittest import HomeserverTestCase
@@ -61,27 +63,33 @@ async def _process_queue(self, values: list[str]) -> str:
 
     def _get_sample_with_name(self, metric: ObservableGauge, name: str) -> float:
         """For a prometheus metric get the value of the sample that has a
-        matching "name" label.
+        matching "name" label and matching metric name.
         """
-        print(vars(metric))
+        # The metric.name attribute gives us the OTel instrument name
+        metric_name = metric.name
+
         for metric_family in REGISTRY.collect():
+            # Check if this metric family corresponds to our metric
+            # (the family name should match or contain the metric name)
+            if metric_family.name != metric_name:
+                continue
             for sample in metric_family.samples:
-                if sample.labels.get("name") == name:  # and sample.name == metric.name:
+                if sample.labels.get("name") == name:
                     return sample.value
 
-        self.fail("Found no matching sample")
+        self.fail(f"Found no matching sample for metric={metric_name}, name={name}")
 
     def _assert_metrics(self, queued: int, keys: int, in_flight: int) -> None:
         """Assert that the metrics are correct"""
 
-        sample = self._get_sample_with_name(self.queue.number_queued, self.queue._name)
+        sample = self._get_sample_with_name(number_queued, self.queue._name)
         self.assertEqual(
             sample,
             queued,
             "number_queued",
         )
 
-        sample = self._get_sample_with_name(self.queue.number_of_keys, self.queue._name)
+        sample = self._get_sample_with_name(number_of_keys, self.queue._name)
         self.assertEqual(sample, keys, "number_of_keys")
 
         sample = self._get_sample_with_name(number_in_flight, self.queue._name)