From 9b13817e067d370f3ddb860baef884bf460449b5 Mon Sep 17 00:00:00 2001
From: Neil Johnson <neil@fragile.org.uk>
Date: Mon, 30 Jul 2018 22:07:07 +0100
Subject: [PATCH] factor out metrics from __init__ to app/homeserver

---
 synapse/app/homeserver.py   | 20 +++++++++++++++++++-
 synapse/storage/__init__.py | 37 +++++++++++++------------------------
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/synapse/app/homeserver.py b/synapse/app/homeserver.py
index 57b815d777..a1512ccea4 100755
--- a/synapse/app/homeserver.py
+++ b/synapse/app/homeserver.py
@@ -18,6 +18,7 @@ import logging
 import os
 import sys
 
+from prometheus_client import Gauge
 from six import iteritems
 
 from twisted.application import service
@@ -299,7 +300,12 @@ class SynapseHomeServer(HomeServer):
         except IncorrectDatabaseSetup as e:
             quit_with_error(e.message)
 
-
+# Gauges to expose monthly active user control metrics
+current_mau_gauge = Gauge("synapse_admin_current_mau", "Current MAU")
+max_mau_value_gauge = Gauge("synapse_admin_max_mau_value", "MAU Limit")
+limit_usage_by_mau_gauge = Gauge(
+    "synapse_admin_limit_usage_by_mau", "MAU Limiting enabled"
+)
 def setup(config_options):
     """
     Args:
@@ -512,6 +518,18 @@ def run(hs):
     # table will decrease
     clock.looping_call(generate_user_daily_visit_stats, 5 * 60 * 1000)
 
+    def generate_monthly_active_users():
+        count = 0
+        if hs.config.limit_usage_by_mau:
+            count = hs.get_datastore().count_monthly_users()
+        logger.info("NJ count is %d" % (count,))
+        current_mau_gauge.set(float(count))
+        max_mau_value_gauge.set(float(hs.config.max_mau_value))
+        limit_usage_by_mau_gauge.set(float(hs.config.limit_usage_by_mau))
+
+    generate_monthly_active_users()
+    clock.looping_call(generate_monthly_active_users, 5 * 60 * 1000)
+
     if hs.config.report_stats:
         logger.info("Scheduling stats reporting for 3 hour intervals")
         clock.looping_call(start_phone_stats_home, 3 * 60 * 60 * 1000)
diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index 6a75bf0e52..7b8215bf08 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -19,7 +19,6 @@ import logging
 import time
 
 from dateutil import tz
-from prometheus_client import Gauge
 
 from synapse.api.constants import PresenceState
 from synapse.storage.devices import DeviceStore
@@ -61,14 +60,6 @@ from .util.id_generators import ChainedIdGenerator, IdGenerator, StreamIdGenerat
 
 logger = logging.getLogger(__name__)
 
-# Gauges to expose monthly active user control metrics
-current_mau_gauge = Gauge("synapse_admin_current_mau", "Current MAU")
-max_mau_value_gauge = Gauge("synapse_admin_max_mau_value", "MAU Limit")
-limit_usage_by_mau_gauge = Gauge(
-    "synapse_admin_limit_usage_by_mau", "MAU Limiting enabled"
-)
-
-
 class DataStore(RoomMemberStore, RoomStore,
                 RegistrationStore, StreamStore, ProfileStore,
                 PresenceStore, TransactionStore,
@@ -102,6 +93,7 @@ class DataStore(RoomMemberStore, RoomStore,
         self._clock = hs.get_clock()
         self.database_engine = hs.database_engine
 
+        self.db_conn = db_conn
         self._stream_id_gen = StreamIdGenerator(
             db_conn, "events", "stream_ordering",
             extra_tables=[("local_invites", "stream_id")]
@@ -282,22 +274,19 @@ class DataStore(RoomMemberStore, RoomStore,
         returns:
             int: count of current monthly active users
         """
-        def _count_monthly_users(txn):
-            thirty_days_ago = int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30)
-            sql = """
-                SELECT COUNT(*) FROM user_ips
-                WHERE last_seen > ?
-            """
-            txn.execute(sql, (thirty_days_ago,))
-            count, = txn.fetchone()
 
-            self._current_mau = count
-            current_mau_gauge.set(self._current_mau)
-            max_mau_value_gauge.set(self.hs.config.max_mau_value)
-            limit_usage_by_mau_gauge.set(self.hs.config.limit_usage_by_mau)
-            logger.info("calling mau stats")
-            return count
-        return self.runInteraction("count_monthly_users", _count_monthly_users)
+        thirty_days_ago = int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30)
+        sql = """
+            SELECT COALESCE(count(*), 0) FROM (
+                SELECT user_id FROM user_ips
+                WHERE last_seen > ?
+                GROUP BY user_id
+            ) u
+        """
+        txn = self.db_conn.cursor()
+        txn.execute(sql, (thirty_days_ago,))
+        count, = txn.fetchone()
+        return count
 
 
     def count_r30_users(self):