feat: add room count metrics

Author: itsoyou

synapse/handlers/stats.py

@@ -32,7 +32,12 @@
 )
 
 from synapse.api.constants import EventContentFields, EventTypes, Membership
-from synapse.metrics import SERVER_NAME_LABEL, event_processing_positions
+from synapse.metrics import (
+    SERVER_NAME_LABEL,
+    event_processing_positions,
+    known_rooms_gauge,
+    locally_joined_rooms_gauge,
+)
 from synapse.metrics.background_process_metrics import run_as_background_process
 from synapse.storage.databases.main.state_deltas import StateDelta
 from synapse.types import JsonDict
@@ -153,6 +158,19 @@ async def _unsafe_process(self) -> None:
 
             self.pos = max_pos
 
+            (
+                known_room_count,
+                locally_joined_room_count,
+            ) = await self.store.get_room_stats()
+
+            # Update room count metrics
+            known_rooms_gauge.labels(**{SERVER_NAME_LABEL: self.server_name}).set(
+                known_room_count
+            )
+            locally_joined_rooms_gauge.labels(
+                **{SERVER_NAME_LABEL: self.server_name}
+            ).set(locally_joined_room_count)
+
     async def _handle_deltas(
         self, deltas: Iterable[StateDelta]
     ) -> Tuple[Dict[str, CounterType[str]], Dict[str, CounterType[str]]]:

synapse/metrics/__init__.py

@@ -645,6 +645,19 @@ def collect(self) -> Iterable[Metric]:
     labelnames=["name", SERVER_NAME_LABEL],
 )
 
+# Gauges for room counts
+known_rooms_gauge = Gauge(
+    "synapse_known_rooms_total",
+    "Total number of rooms",
+    labelnames=[SERVER_NAME_LABEL],
+)
+
+locally_joined_rooms_gauge = Gauge(
+    "synapse_locally_joined_rooms_total",
+    "Total number of locally joined rooms",
+    labelnames=[SERVER_NAME_LABEL],
+)
+
 
 def register_threadpool(*, name: str, server_name: str, threadpool: ThreadPool) -> None:
     """

synapse/storage/databases/main/stats.py

@@ -324,7 +324,6 @@ async def bulk_update_stats_delta(
     ) -> None:
         """Bulk update stats tables for a given stream_id and updates the stats
         incremental position.
-
         Args:
             ts: Current timestamp in ms
             updates: The updates to commit as a mapping of
@@ -358,6 +357,24 @@ def _bulk_update_stats_delta_txn(txn: LoggingTransaction) -> None:
             "bulk_update_stats_delta", _bulk_update_stats_delta_txn
         )
 
+    async def get_room_stats(self) -> Tuple[int, int]:
+        """
+        Retrieve the total number of rooms and locally joined rooms.
+        """
+
+        def _get_room_stats_txn(txn: LoggingTransaction) -> Tuple[int, int]:
+            sql = """
+                SELECT
+                    count(*) AS total,
+                    count(CASE WHEN local_users_in_room > 0 THEN 1 END) AS locally_joined
+                FROM room_stats_current;
+                """
+            txn.execute(sql)
+            row = cast(Tuple[int, int], txn.fetchone())
+            return row[0], row[1]
+
+        return await self.db_pool.runInteraction("get_room_stats", _get_room_stats_txn)
+
     async def update_stats_delta(
         self,
         ts: int,

tests/handlers/test_stats.py

@@ -20,6 +20,8 @@
 
 from typing import Any, Dict, List, Optional, Tuple, cast
 
+from prometheus_client import REGISTRY, Gauge
+
 from twisted.internet.testing import MemoryReactor
 
 from synapse.rest import admin
@@ -46,6 +48,21 @@ class StatsRoomTests(unittest.HomeserverTestCase):
     def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
         self.store = hs.get_datastores().main
         self.handler = self.hs.get_stats_handler()
+        self._set_metrics_to_zero()
+
+    def _set_metrics_to_zero(self) -> None:
+        """
+        There is an issue when the tests are run in parallel where the metrics
+        are not reset between tests, leading to incorrect values.
+        This method resets the metrics to zero before each test to ensure
+        that each test starts with a clean slate.
+        """
+        metrics = ["synapse_known_rooms_total", "synapse_locally_joined_rooms_total"]
+        for metric_name in metrics:
+            gauge = REGISTRY._names_to_collectors.get(metric_name)
+            if gauge is not None and isinstance(gauge, Gauge):
+                for labels in gauge._metrics:
+                    gauge.labels(*labels).set(0)
 
     def _add_background_updates(self) -> None:
         """
@@ -162,9 +179,20 @@ 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,
+        )
         u1 = self.register_user("u1", "pass")
         u1token = self.login("u1", "pass")
         r1 = self.helper.create_room_as(u1, tok=u1token)
@@ -190,6 +218,21 @@ def test_create_room(self) -> None:
         self.assertEqual(r2stats["invited_members"], 0)
         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,
+        )
+
     def test_updating_profile_information_does_not_increase_joined_members_count(
         self,
     ) -> None:
@@ -592,3 +635,57 @@ def test_incomplete_stats(self) -> None:
 
         self.assertEqual(u1stats_complete["joined_rooms"], 1)
         self.assertEqual(u2stats_complete["joined_rooms"], 1)
+
+    def test_room_metrics(self) -> None:
+        """
+        Test that the `synapse_locally_joined_rooms_total` and
+        `synapse_known_rooms_total` metrics are updated correctly.
+        """
+
+        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,
+        )
+
+        u1 = self.register_user("u1", "pass")
+        u1token = self.login("u1", "pass")
+
+        # create 2 rooms, one with a local user and one without it.
+        r1 = self.helper.create_room_as(u1, tok=u1token)
+        r2 = self.helper.create_room_as(u1, tok=u1token)
+        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,
+        )
+
+        # Check the stats for both rooms
+        r1stats = self._get_current_stats("room", r1)
+        r2stats = self._get_current_stats("room", r2)
+
+        assert r1stats is not None
+        assert r2stats is not None
+
+        self.assertEqual(r1stats["joined_members"], 1)
+        self.assertEqual(r2stats["joined_members"], 0)