fix caching and tests

3 files changed, 80 insertions, 62 deletions

diff --git a/synapse/app/homeserver.py b/synapse/app/homeserver.py
index 7d4ea493bc..3a67db8b30 100755
--- a/synapse/app/homeserver.py
+++ b/synapse/app/homeserver.py
@@ -64,7 +64,6 @@ from synapse.rest.media.v0.content_repository import ContentRepoResource
 from synapse.server import HomeServer
 from synapse.storage import are_all_users_on_domain
 from synapse.storage.engines import IncorrectDatabaseSetup, create_engine
-from synapse.storage.monthly_active_users import MonthlyActiveUsersStore
 from synapse.storage.prepare_database import UpgradeDatabaseException, prepare_database
 from synapse.util.caches import CACHE_SIZE_FACTOR
 from synapse.util.httpresourcetree import create_resource_tree
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))
diff --git a/tests/storage/test_monthly_active_users.py b/tests/storage/test_monthly_active_users.py
index d8d25a6069..c32109ecc5 100644
--- a/tests/storage/test_monthly_active_users.py
+++ b/tests/storage/test_monthly_active_users.py
@@ -1,6 +1,19 @@
-from twisted.internet import defer
+# -*- 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 synapse.storage.monthly_active_users import MonthlyActiveUsersStore
+from twisted.internet import defer
 
 import tests.unittest
 import tests.utils
@@ -10,20 +23,19 @@ from tests.utils import setup_test_homeserver
 class MonthlyActiveUsersTestCase(tests.unittest.TestCase):
     def __init__(self, *args, **kwargs):
         super(MonthlyActiveUsersTestCase, self).__init__(*args, **kwargs)
-        self.mau = None
 
     @defer.inlineCallbacks
     def setUp(self):
-        hs = yield setup_test_homeserver()
-        self.mau = MonthlyActiveUsersStore(None, hs)
+        self.hs = yield setup_test_homeserver()
+        self.store = self.hs.get_datastore()
 
     @defer.inlineCallbacks
     def test_can_insert_and_count_mau(self):
-        count = yield self.mau.get_monthly_active_count()
+        count = yield self.store.get_monthly_active_count()
         self.assertEqual(0, count)
 
-        yield self.mau.upsert_monthly_active_user("@user:server")
-        count = yield self.mau.get_monthly_active_count()
+        yield self.store.upsert_monthly_active_user("@user:server")
+        count = yield self.store.get_monthly_active_count()
 
         self.assertEqual(1, count)
 
@@ -32,11 +44,23 @@ class MonthlyActiveUsersTestCase(tests.unittest.TestCase):
         user_id1 = "@user1:server"
         user_id2 = "@user2:server"
         user_id3 = "@user3:server"
-        result = yield self.mau.is_user_monthly_active(user_id1)
+        result = yield self.store.is_user_monthly_active(user_id1)
         self.assertFalse(result)
-        yield self.mau.upsert_monthly_active_user(user_id1)
-        yield self.mau.upsert_monthly_active_user(user_id2)
-        result = yield self.mau.is_user_monthly_active(user_id1)
+        yield self.store.upsert_monthly_active_user(user_id1)
+        yield self.store.upsert_monthly_active_user(user_id2)
+        result = yield self.store.is_user_monthly_active(user_id1)
         self.assertTrue(result)
-        result = yield self.mau.is_user_monthly_active(user_id3)
+        result = yield self.store.is_user_monthly_active(user_id3)
         self.assertFalse(result)
+
+    @defer.inlineCallbacks
+    def test_reap_monthly_active_users(self):
+        self.hs.config.max_mau_value = 5
+        initial_users = 10
+        for i in range(initial_users):
+            yield self.store.upsert_monthly_active_user("@user%d:server" % i)
+        count = yield self.store.get_monthly_active_count()
+        self.assertTrue(count, initial_users)
+        yield self.store.reap_monthly_active_users()
+        count = yield self.store.get_monthly_active_count()
+        self.assertTrue(count, initial_users - self.hs.config.max_mau_value)