Add caches for whether a room has been forgotten by a user

1 files changed, 7 insertions, 4 deletions

diff --git a/synapse/storage/roommember.py b/synapse/storage/roommember.py
index 69398b7c8e..e1777d7afa 100644
--- a/synapse/storage/roommember.py
+++ b/synapse/storage/roommember.py
@@ -18,7 +18,7 @@ from twisted.internet import defer
 from collections import namedtuple
 
 from ._base import SQLBaseStore
-from synapse.util.caches.descriptors import cached
+from synapse.util.caches.descriptors import cached, cachedInlineCallbacks
 
 from synapse.api.constants import Membership
 from synapse.types import UserID
@@ -270,6 +270,7 @@ class RoomMemberStore(SQLBaseStore):
 
         defer.returnValue(ret)
 
+    @defer.inlineCallbacks
     def forget(self, user_id, room_id):
         """Indicate that user_id wishes to discard history for room_id."""
         def f(txn):
@@ -284,9 +285,11 @@ class RoomMemberStore(SQLBaseStore):
                 "  room_id = ?"
             )
             txn.execute(sql, (user_id, room_id))
-        self.runInteraction("forget_membership", f)
+        yield self.runInteraction("forget_membership", f)
+        self.was_forgotten_at.invalidate_all()
+        self.did_forget.invalidate((user_id, room_id))
 
-    @defer.inlineCallbacks
+    @cachedInlineCallbacks(num_args=2)
     def did_forget(self, user_id, room_id):
         """Returns whether user_id has elected to discard history for room_id.
 
@@ -310,7 +313,7 @@ class RoomMemberStore(SQLBaseStore):
         count = yield self.runInteraction("did_forget_membership", f)
         defer.returnValue(count == 0)
 
-    @defer.inlineCallbacks
+    @cachedInlineCallbacks(num_args=3)
     def was_forgotten_at(self, user_id, room_id, event_id):
         """Returns whether user_id has elected to discard history for room_id at event_id.