fix caching and tests

1 files changed, 43 insertions, 48 deletions

diff --git a/synapse/storage/monthly_active_users.py b/synapse/storage/monthly_active_users.py
index 2872ba4cae..d06c90cbcc 100644
--- a/synapse/storage/monthly_active_users.py
+++ b/synapse/storage/monthly_active_users.py
@@ -1,6 +1,21 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 New Vector
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from twisted.internet import defer
-from synapse.util.caches.descriptors import cachedInlineCallbacks
-from synapse.storage.engines import PostgresEngine, Sqlite3Engine
+
+from synapse.util.caches.descriptors import cached, cachedInlineCallbacks
 
 from ._base import SQLBaseStore
 
@@ -9,7 +24,7 @@ class MonthlyActiveUsersStore(SQLBaseStore):
     def __init__(self, dbconn, hs):
         super(MonthlyActiveUsersStore, self).__init__(None, hs)
         self._clock = hs.get_clock()
-        self.max_mau_value = hs.config.max_mau_value
+        self.hs = hs
 
     def reap_monthly_active_users(self):
         """
@@ -19,62 +34,41 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             Defered()
         """
         def _reap_users(txn):
+
             thirty_days_ago = (
                 int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30)
             )
 
-            if isinstance(self.database_engine, PostgresEngine):
-                sql = """
-                    DELETE FROM monthly_active_users
-                    WHERE timestamp < ?
-                    RETURNING user_id
-                    """
-                txn.execute(sql, (thirty_days_ago,))
-                res = txn.fetchall()
-                for r in res:
-                    self.is_user_monthly_active.invalidate(r)
-
-                sql = """
-                    DELETE FROM monthly_active_users
-                    ORDER BY timestamp desc
-                    LIMIT -1 OFFSET ?
-                    RETURNING user_id
-                    """
-                txn.execute(sql, (self.max_mau_value,))
-                res = txn.fetchall()
-                for r in res:
-                    self.is_user_monthly_active.invalidate(r)
-                    print r
-                self.get_monthly_active_count.invalidate()
-            elif isinstance(self.database_engine, Sqlite3Engine):
-                sql = "DELETE FROM monthly_active_users WHERE timestamp < ?"
+            sql = "DELETE FROM monthly_active_users WHERE timestamp < ?"
 
-                txn.execute(sql, (thirty_days_ago,))
-                sql = """
-                    DELETE FROM monthly_active_users
-                    ORDER BY timestamp desc
-                    LIMIT -1 OFFSET ?
-                    """
-                txn.execute(sql, (self.max_mau_value,))
+            txn.execute(sql, (thirty_days_ago,))
+            sql = """
+                DELETE FROM monthly_active_users
+                ORDER BY timestamp desc
+                LIMIT -1 OFFSET ?
+                """
+            txn.execute(sql, (self.hs.config.max_mau_value,))
 
-                # It seems poor to invalidate the whole cache, but the alternative
-                # is to select then delete which has its own problems.
-                # It seems unlikely that anyone using this feature on large datasets
-                # would be using sqlite and if they are then there will be
-                # larger perf issues than this one to encourage an upgrade to postgres.
+        res = self.runInteraction("reap_monthly_active_users", _reap_users)
+        # It seems poor to invalidate the whole cache, Postgres supports
+        # 'Returning' which would allow me to invalidate only the
+        # specific users, but sqlite has no way to do this and instead
+        # I would need to SELECT and the DELETE which without locking
+        # is racy.
+        # Have resolved to invalidate the whole cache for now and do
+        # something about it if and when the perf becomes significant
+        self.is_user_monthly_active.invalidate_all()
+        self.get_monthly_active_count.invalidate_all()
+        return res
 
-                self.is_user_monthly_active.invalidate_all()
-                self.get_monthly_active_count.invalidate_all()
-
-        return self.runInteraction("reap_monthly_active_users", _reap_users)
-
-    @cachedInlineCallbacks(num_args=0)
+    @cached(num_args=0)
     def get_monthly_active_count(self):
         """
             Generates current count of monthly active users.abs
             Return:
                 Defered(int): Number of current monthly active users
         """
+
         def _count_users(txn):
             sql = "SELECT COALESCE(count(*), 0) FROM monthly_active_users"
 
@@ -91,7 +85,7 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             Deferred(bool): True if a new entry was created, False if an
                 existing one was updated.
         """
-        return self._simple_upsert(
+        self._simple_upsert(
             desc="upsert_monthly_active_user",
             table="monthly_active_users",
             keyvalues={
@@ -102,6 +96,8 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             },
             lock=False,
         )
+        self.is_user_monthly_active.invalidate((user_id,))
+        self.get_monthly_active_count.invalidate(())
 
     @cachedInlineCallbacks(num_args=1)
     def is_user_monthly_active(self, user_id):
@@ -120,5 +116,4 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             retcol="user_id",
             desc="is_user_monthly_active",
         )
-
         defer.returnValue(bool(user_present))