Servers-known-about statistic (#5981)

2 files changed, 90 insertions, 0 deletions

diff --git a/synapse/config/metrics.py b/synapse/config/metrics.py
index 3698441963..653b990e67 100644
--- a/synapse/config/metrics.py
+++ b/synapse/config/metrics.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 # Copyright 2015, 2016 OpenMarket Ltd
+# Copyright 2019 The Matrix.org Foundation C.I.C.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import attr
+
 from ._base import Config, ConfigError
 
 MISSING_SENTRY = """Missing sentry-sdk library. This is required to enable sentry
@@ -20,6 +23,18 @@ MISSING_SENTRY = """Missing sentry-sdk library. This is required to enable sentr
     """
 
 
+@attr.s
+class MetricsFlags(object):
+    known_servers = attr.ib(default=False, validator=attr.validators.instance_of(bool))
+
+    @classmethod
+    def all_off(cls):
+        """
+        Instantiate the flags with all options set to off.
+        """
+        return cls(**{x.name: False for x in attr.fields(cls)})
+
+
 class MetricsConfig(Config):
     def read_config(self, config, **kwargs):
         self.enable_metrics = config.get("enable_metrics", False)
@@ -27,6 +42,12 @@ class MetricsConfig(Config):
         self.metrics_port = config.get("metrics_port")
         self.metrics_bind_host = config.get("metrics_bind_host", "127.0.0.1")
 
+        if self.enable_metrics:
+            _metrics_config = config.get("metrics_flags") or {}
+            self.metrics_flags = MetricsFlags(**_metrics_config)
+        else:
+            self.metrics_flags = MetricsFlags.all_off()
+
         self.sentry_enabled = "sentry" in config
         if self.sentry_enabled:
             try:
@@ -58,6 +79,16 @@ class MetricsConfig(Config):
         #sentry:
         #    dsn: "..."
 
+        # Flags to enable Prometheus metrics which are not suitable to be
+        # enabled by default, either for performance reasons or limited use.
+        #
+        metrics_flags:
+            # Publish synapse_federation_known_servers, a g auge of the number of
+            # servers this homeserver knows about, including itself. May cause
+            # performance problems on large homeservers.
+            #
+            #known_servers: true
+
         # Whether or not to report anonymized homeserver usage statistics.
         """
 
diff --git a/synapse/storage/roommember.py b/synapse/storage/roommember.py
index f8b682ebd9..4df8ebdacd 100644
--- a/synapse/storage/roommember.py
+++ b/synapse/storage/roommember.py
@@ -24,8 +24,10 @@ from canonicaljson import json
 from twisted.internet import defer
 
 from synapse.api.constants import EventTypes, Membership
+from synapse.metrics import LaterGauge
 from synapse.metrics.background_process_metrics import run_as_background_process
 from synapse.storage._base import LoggingTransaction
+from synapse.storage.engines import Sqlite3Engine
 from synapse.storage.events_worker import EventsWorkerStore
 from synapse.types import get_domain_from_id
 from synapse.util.async_helpers import Linearizer
@@ -74,6 +76,63 @@ class RoomMemberWorkerStore(EventsWorkerStore):
         self._check_safe_current_state_events_membership_updated_txn(txn)
         txn.close()
 
+        if self.hs.config.metrics_flags.known_servers:
+            self._known_servers_count = 1
+            self.hs.get_clock().looping_call(
+                run_as_background_process,
+                60 * 1000,
+                "_count_known_servers",
+                self._count_known_servers,
+            )
+            self.hs.get_clock().call_later(
+                1000,
+                run_as_background_process,
+                "_count_known_servers",
+                self._count_known_servers,
+            )
+            LaterGauge(
+                "synapse_federation_known_servers",
+                "",
+                [],
+                lambda: self._known_servers_count,
+            )
+
+    @defer.inlineCallbacks
+    def _count_known_servers(self):
+        """
+        Count the servers that this server knows about.
+
+        The statistic is stored on the class for the
+        `synapse_federation_known_servers` LaterGauge to collect.
+        """
+
+        def _transact(txn):
+            if isinstance(self.database_engine, Sqlite3Engine):
+                query = """
+                    SELECT COUNT(DISTINCT substr(out.user_id, pos+1))
+                    FROM (
+                        SELECT rm.user_id as user_id, instr(rm.user_id, ':')
+                            AS pos FROM room_memberships as rm
+                        INNER JOIN current_state_events as c ON rm.event_id = c.event_id
+                        WHERE c.type = 'm.room.member'
+                    ) as out
+                """
+            else:
+                query = """
+                    SELECT COUNT(DISTINCT split_part(state_key, ':', 2))
+                    FROM current_state_events
+                    WHERE type = 'm.room.member' AND membership = 'join';
+                """
+            txn.execute(query)
+            return list(txn)[0][0]
+
+        count = yield self.runInteraction("get_known_servers", _transact)
+
+        # We always know about ourselves, even if we have nothing in
+        # room_memberships (for example, the server is new).
+        self._known_servers_count = max([count, 1])
+        return self._known_servers_count
+
     def _check_safe_current_state_events_membership_updated_txn(self, txn):
         """Checks if it is safe to assume the new current_state_events
         membership column is up to date