Add migration and storage methods for users_to_send_full_presence_to table

2 files changed, 109 insertions, 1 deletions

diff --git a/synapse/storage/databases/main/presence.py b/synapse/storage/databases/main/presence.py
index c207d917b1..6e4aa2629a 100644
--- a/synapse/storage/databases/main/presence.py
+++ b/synapse/storage/databases/main/presence.py
@@ -12,13 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, Tuple
+from typing import Dict, Iterable, List, Tuple
 
 from synapse.api.presence import UserPresenceState
 from synapse.storage._base import SQLBaseStore, make_in_list_sql_clause
 from synapse.util.caches.descriptors import cached, cachedList
+from synapse.util.constants import HOUR_IN_MS
 from synapse.util.iterutils import batch_iter
 
+# How long to keep entries in the users_to_send_full_presence_to db table
+USERS_TO_SEND_FULL_PRESENCE_TO_ENTRY_LIFETIME_MS = 14 * 24 * HOUR_IN_MS
+
 
 class PresenceStore(SQLBaseStore):
     async def update_presence(self, presence_states):
@@ -156,6 +160,79 @@ class PresenceStore(SQLBaseStore):
 
         return {row["user_id"]: UserPresenceState(**row) for row in rows}
 
+    async def is_user_in_users_to_send_full_presence_to(self, user_id: str) -> bool:
+        """Check whether the given user is one we need to send full presence to.
+
+        Args:
+            user_id: The ID of the user to check.
+
+        Returns:
+            True if the user should have full presence sent to them, False otherwise.
+        """
+
+        def _is_user_in_users_to_send_full_presence_to_txn(txn):
+            sql = """
+                SELECT 1 FROM users_to_send_full_presence_to
+                WHERE user_id = ?
+            """
+            txn.execute(sql, (user_id,))
+            return bool(txn.fetchone())
+
+        return await self.db_pool.runInteraction(
+            "is_user_in_users_to_send_full_presence_to",
+            _is_user_in_users_to_send_full_presence_to_txn,
+        )
+
+    async def add_users_to_send_full_presence_to(self, user_ids: Iterable[str]):
+        """Adds to the list of users who should receive a full snapshot of presence
+        upon their next sync.
+
+        Entries are kept for at least USERS_TO_SEND_FULL_PRESENCE_TO_ENTRY_LIFETIME_MS,
+        and are culled whenever this method is called.
+
+        Args:
+            user_ids: An iterable of user IDs.
+        """
+        time_now = self.hs.get_clock().time_msec()
+
+        def _add_users_to_send_full_presence_to_txn(txn):
+            # Add user entries to the table (or update their added_ms if they already exist)
+            self.db_pool.simple_upsert_many_txn(
+                txn,
+                table="users_to_send_full_presence_to",
+                key_names=("user_id",),
+                key_values=((user_id,) for user_id in user_ids),
+                value_names=("added_ms",),
+                value_values=((time_now,) for _ in user_ids),
+            )
+
+            # Delete entries in the table that have expired
+            sql = """
+                DELETE FROM users_to_send_full_presence_to
+                WHERE added_ms < ?
+            """
+            txn.execute(
+                sql, (time_now - USERS_TO_SEND_FULL_PRESENCE_TO_ENTRY_LIFETIME_MS)
+            )
+
+        await self.db_pool.runInteraction(
+            "add_users_to_send_full_presence_to",
+            _add_users_to_send_full_presence_to_txn,
+        )
+
+    async def remove_user_from_users_to_send_full_presence_to(self, user_id: str):
+        """Removes a user from those to send full presence to. This should only be done
+        once we have sent full presence to them following a successful sync.
+
+        Args:
+            user_id: The ID of the user to remove from the table.
+        """
+        await self.db_pool.simple_delete(
+            table="users_to_send_full_presence_to",
+            keyvalues={"user_id": user_id},
+            desc="remove_user_from_users_to_send_full_presence_to",
+        )
+
     async def get_presence_for_all_users(
         self,
         include_offline: bool = True,
diff --git a/synapse/storage/databases/main/schema/delta/59/12users_to_send_full_presence_to.sql b/synapse/storage/databases/main/schema/delta/59/12users_to_send_full_presence_to.sql
new file mode 100644
index 0000000000..34ec1c0268
--- /dev/null
+++ b/synapse/storage/databases/main/schema/delta/59/12users_to_send_full_presence_to.sql
@@ -0,0 +1,31 @@
+/* Copyright 2021 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.
+ * 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.
+ */
+
+-- Add a table that keeps track of a list of users who should, upon their next
+-- sync request, receive presence for all currently online users that they are
+-- "interested" in.
+
+-- The motivation for a DB table over an in-memory list is so that this list
+-- can be added to and retrieved from by any worker. Specifically, we don't
+-- want to duplicate work across multiple sync workers.
+
+CREATE TABLE IF NOT EXISTS users_to_send_full_presence_to(
+    user_id TEXT PRIMARY KEY,
+    added_ms BIGINT NOT NULL
+);
+
+-- For checking expired entries
+CREATE INDEX IF NOT EXISTS users_to_send_full_presence_to_added_ms
+    ON users_to_send_full_presence_to(added_ms);
\ No newline at end of file