PLAT-38547 Account-data compare-and-swap

synapse/api/errors.py

@@ -152,6 +152,9 @@ class Codes(str, Enum):
     # Part of MSC4326
     UNKNOWN_DEVICE = "ORG.MATRIX.MSC4326.M_UNKNOWN_DEVICE"
 
+    # Beeper: account data compare-and-swap via com.beeper.expect_revision_id
+    EXPECTED_REVISION_ID_MISMATCH = "COM.BEEPER.REVISION_ID_MISMATCH"
+
 
 class CodeMessageException(RuntimeError):
     """An exception with integer code, a message string attributes and optional headers.

synapse/handlers/account_data.py

@@ -104,7 +104,8 @@ async def _notify_modules(
                 logger.exception("Failed to run module callback %s: %s", callback, e)
 
     async def add_account_data_to_room(
-        self, user_id: str, room_id: str, account_data_type: str, content: JsonDict
+        self, user_id: str, room_id: str, account_data_type: str, content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> int:
         """Add some account_data to a room for a user.
 
@@ -113,13 +114,15 @@ async def add_account_data_to_room(
             room_id: The room to add a tag for.
             account_data_type: The type of account_data to add.
             content: A json object to associate with the tag.
+            expected_revision_id: If set, only write if the stored content's
+                `com.beeper.revision_id` matches (compare-and-swap).
 
         Returns:
             The maximum stream ID.
         """
         if self._instance_name in self._account_data_writers:
             max_stream_id = await self._store.add_account_data_to_room(
-                user_id, room_id, account_data_type, content
+                user_id, room_id, account_data_type, content, expected_revision_id
             )
 
             self._notifier.on_new_event(
@@ -136,6 +139,7 @@ async def add_account_data_to_room(
                 room_id=room_id,
                 account_data_type=account_data_type,
                 content=content,
+                expected_revision_id=expected_revision_id,
             )
             return response["max_stream_id"]
 
@@ -181,22 +185,25 @@ async def remove_account_data_for_room(
             return response["max_stream_id"]
 
     async def add_account_data_for_user(
-        self, user_id: str, account_data_type: str, content: JsonDict
+        self, user_id: str, account_data_type: str, content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> int:
         """Add some global account_data for a user.
 
         Args:
             user_id: The user to add some account data for.
             account_data_type: The type of account_data to add.
             content: The content json dictionary.
+            expected_revision_id: If set, only write if the stored content's
+                `com.beeper.revision_id` matches (compare-and-swap).
 
         Returns:
             The maximum stream ID.
         """
 
         if self._instance_name in self._account_data_writers:
             max_stream_id = await self._store.add_account_data_for_user(
-                user_id, account_data_type, content
+                user_id, account_data_type, content, expected_revision_id
             )
 
             self._notifier.on_new_event(
@@ -212,6 +219,7 @@ async def add_account_data_for_user(
                 user_id=user_id,
                 account_data_type=account_data_type,
                 content=content,
+                expected_revision_id=expected_revision_id,
             )
             return response["max_stream_id"]
 

synapse/replication/http/account_data.py

@@ -58,19 +58,23 @@ def __init__(self, hs: "HomeServer"):
 
     @staticmethod
     async def _serialize_payload(  # type: ignore[override]
-        user_id: str, account_data_type: str, content: JsonDict
+        user_id: str, account_data_type: str, content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> JsonDict:
-        payload = {
+        payload: JsonDict = {
             "content": content,
         }
+        if expected_revision_id is not None:
+            payload["expected_revision_id"] = expected_revision_id
 
         return payload
 
     async def _handle_request(  # type: ignore[override]
         self, request: Request, content: JsonDict, user_id: str, account_data_type: str
     ) -> tuple[int, JsonDict]:
         max_stream_id = await self.handler.add_account_data_for_user(
-            user_id, account_data_type, content["content"]
+            user_id, account_data_type, content["content"],
+            expected_revision_id=content.get("expected_revision_id"),
         )
 
         return 200, {"max_stream_id": max_stream_id}
@@ -138,11 +142,14 @@ def __init__(self, hs: "HomeServer"):
 
     @staticmethod
     async def _serialize_payload(  # type: ignore[override]
-        user_id: str, room_id: str, account_data_type: str, content: JsonDict
+        user_id: str, room_id: str, account_data_type: str, content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> JsonDict:
-        payload = {
+        payload: JsonDict = {
             "content": content,
         }
+        if expected_revision_id is not None:
+            payload["expected_revision_id"] = expected_revision_id
 
         return payload
 
@@ -155,7 +162,8 @@ async def _handle_request(  # type: ignore[override]
         account_data_type: str,
     ) -> tuple[int, JsonDict]:
         max_stream_id = await self.handler.add_account_data_to_room(
-            user_id, room_id, account_data_type, content["content"]
+            user_id, room_id, account_data_type, content["content"],
+            expected_revision_id=content.get("expected_revision_id"),
         )
 
         return 200, {"max_stream_id": max_stream_id}

synapse/rest/client/account_data.py

@@ -26,7 +26,11 @@
 from synapse.api.constants import AccountDataTypes, ReceiptTypes
 from synapse.api.errors import AuthError, Codes, NotFoundError, SynapseError
 from synapse.http.server import HttpServer
-from synapse.http.servlet import RestServlet, parse_json_object_from_request
+from synapse.http.servlet import (
+    RestServlet,
+    parse_json_object_from_request,
+    parse_string,
+)
 from synapse.http.site import SynapseRequest
 from synapse.rest.client.read_marker import ReadMarkerRestServlet
 from synapse.types import JsonDict, JsonMapping, RoomID
@@ -86,6 +90,7 @@ async def on_PUT(
         _check_can_set_account_data_type(account_data_type)
 
         body = parse_json_object_from_request(request)
+        expected_revision_id = parse_string(request, "com.beeper.expect_revision_id")
 
         # If experimental support for MSC3391 is enabled, then providing an empty dict
         # as the value for an account data type should be functionally equivalent to
@@ -97,7 +102,9 @@ async def on_PUT(
                 )
                 return 200, {}
 
-        await self.handler.add_account_data_for_user(user_id, account_data_type, body)
+        await self.handler.add_account_data_for_user(
+            user_id, account_data_type, body, expected_revision_id
+        )
 
         return 200, {}
 
@@ -209,6 +216,7 @@ async def on_PUT(
         _check_can_set_account_data_type(account_data_type)
 
         body = parse_json_object_from_request(request)
+        expected_revision_id = parse_string(request, "com.beeper.expect_revision_id")
 
         # If experimental support for MSC3391 is enabled, then providing an empty dict
         # as the value for an account data type should be functionally equivalent to
@@ -221,7 +229,7 @@ async def on_PUT(
                 return 200, {}
 
         await self.handler.add_account_data_to_room(
-            user_id, room_id, account_data_type, body
+            user_id, room_id, account_data_type, body, expected_revision_id
         )
 
         return 200, {}

synapse/storage/databases/main/account_data.py

@@ -41,6 +41,7 @@
 )
 from synapse.storage.databases.main.cache import CacheInvalidationWorkerStore
 from synapse.storage.databases.main.push_rule import PushRulesWorkerStore
+from synapse.storage.engines import PostgresEngine
 from synapse.storage.invite_rule import (
     AllowAllInviteRulesConfig,
     InviteRulesConfig,
@@ -640,7 +641,8 @@ def process_replication_position(
         super().process_replication_position(stream_name, instance_name, token)
 
     async def add_account_data_to_room(
-        self, user_id: str, room_id: str, account_data_type: str, content: JsonDict
+        self, user_id: str, room_id: str, account_data_type: str, content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> int:
         """Add some account_data to a room for a user.
 
@@ -649,6 +651,8 @@ async def add_account_data_to_room(
             room_id: The room to add a tag for.
             account_data_type: The type of account_data to add.
             content: A json object to associate with the tag.
+            expected_revision_id: If set, only write if the stored content's
+                `com.beeper.revision_id` matches (compare-and-swap).
 
         Returns:
             The maximum stream ID.
@@ -657,16 +661,26 @@ async def add_account_data_to_room(
 
         content_json = json_encoder.encode(content)
 
-        async with self._account_data_id_gen.get_next() as next_id:
-            await self.db_pool.simple_upsert(
-                desc="add_room_account_data",
+        def _add_account_data_to_room_txn(
+            txn: LoggingTransaction, next_id: int
+        ) -> None:
+            self._upsert_account_data_txn(
+                txn,
                 table="room_account_data",
                 keyvalues={
                     "user_id": user_id,
                     "room_id": room_id,
                     "account_data_type": account_data_type,
                 },
                 values={"stream_id": next_id, "content": content_json},
+                expected_revision_id=expected_revision_id,
+            )
+
+        async with self._account_data_id_gen.get_next() as next_id:
+            await self.db_pool.runInteraction(
+                "add_room_account_data",
+                _add_account_data_to_room_txn,
+                next_id,
             )
 
             self._account_data_stream_cache.entity_has_changed(user_id, next_id)
@@ -742,14 +756,17 @@ def _remove_account_data_for_room_txn(
         return self._account_data_id_gen.get_current_token()
 
     async def add_account_data_for_user(
-        self, user_id: str, account_data_type: str, content: JsonDict
+        self, user_id: str, account_data_type: str, content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> int:
         """Add some global account_data for a user.
 
         Args:
             user_id: The user to add a tag for.
             account_data_type: The type of account_data to add.
             content: A json object to associate with the tag.
+            expected_revision_id: If set, only write if the stored content's
+                `com.beeper.revision_id` matches (compare-and-swap).
 
         Returns:
             The maximum stream ID.
@@ -764,6 +781,7 @@ async def add_account_data_for_user(
                 user_id,
                 account_data_type,
                 content,
+                expected_revision_id,
             )
 
             self._account_data_stream_cache.entity_has_changed(user_id, next_id)
@@ -774,13 +792,92 @@ async def add_account_data_for_user(
 
         return self._account_data_id_gen.get_current_token()
 
+    def _upsert_account_data_txn(
+        self,
+        txn: LoggingTransaction,
+        table: str,
+        keyvalues: dict[str, str],
+        values: dict[str, Any],
+        expected_revision_id: str | None,
+    ) -> None:
+        """Beeper: upsert account data, enforcing the compare-and-swap
+        condition when an expected revision ID is given.
+
+        Raises a 409 SynapseError if the stored content has a (string)
+        `com.beeper.revision_id` that differs from the expected one. A missing
+        row, missing field, or non-string field matches any expected value.
+        """
+        if expected_revision_id is None:
+            self.db_pool.simple_upsert_txn(txn, table, keyvalues, values)
+            return
+
+        select_sql = "SELECT content FROM %s WHERE %s" % (
+            table,
+            " AND ".join("%s = ?" % k for k in keyvalues),
+        )
+        if isinstance(self.database_engine, PostgresEngine):
+            # Lock the row so concurrent CAS writes serialize. Under Synapse's
+            # default REPEATABLE READ isolation, a row modified by a concurrent
+            # transaction instead raises a serialization failure, which
+            # runInteraction retries with a fresh snapshot. (SQLite serializes
+            # writes anyway, so a plain SELECT suffices there.)
+            select_sql += " FOR UPDATE"
+        txn.execute(select_sql, list(keyvalues.values()))
+        row = txn.fetchone()
+
+        if row is None:
+            # There is no row to lock, so a check-then-write would let two
+            # concurrent first writes both pass the check. Make the INSERT
+            # itself the atomic point instead: it only succeeds if no
+            # concurrent write landed first (in-flight inserts on the same
+            # key serialize via speculative insertion), so on success the
+            # no-data match genuinely held at write time. On conflict, fall
+            # through to compare against the winning row.
+            if self.db_pool.simple_upsert_txn_native_upsert(
+                txn, table, keyvalues, values={}, insertion_values=values
+            ):
+                return
+            txn.execute(select_sql, list(keyvalues.values()))
+            row = txn.fetchone()
+            if row is None:
+                # Under REPEATABLE READ the conflicting row was committed
+                # after our snapshot, so the re-read cannot see it. The
+                # upsert below then hits that invisible row and raises a
+                # serialization failure (40001), which runInteraction
+                # retries from scratch with a fresh snapshot that does see
+                # the row and compares against it. (Under READ COMMITTED
+                # the re-read would have found the row directly.)
+                self.db_pool.simple_upsert_txn(txn, table, keyvalues, values)
+                return
+
+        stored_content = db_to_json(row[0])
+        stored_revision_id = None
+        if isinstance(stored_content, dict):
+            rev = stored_content.get("com.beeper.revision_id")
+            if isinstance(rev, str):
+                stored_revision_id = rev
+
+        if (
+            stored_revision_id is not None
+            and stored_revision_id != expected_revision_id
+        ):
+            raise SynapseError(
+                409,
+                "Account data revision ID mismatch",
+                Codes.EXPECTED_REVISION_ID_MISMATCH,
+                additional_fields={"com.beeper.current_content": stored_content},
+            )
+
+        self.db_pool.simple_upsert_txn(txn, table, keyvalues, values)
+
     def _add_account_data_for_user(
         self,
         txn: LoggingTransaction,
         next_id: int,
         user_id: str,
         account_data_type: str,
         content: JsonDict,
+        expected_revision_id: str | None = None,
     ) -> None:
         content_json = json_encoder.encode(content)
 
@@ -796,11 +893,12 @@ def _add_account_data_for_user(
                     )
                 }
 
-        self.db_pool.simple_upsert_txn(
+        self._upsert_account_data_txn(
             txn,
             table="account_data",
             keyvalues={"user_id": user_id, "account_data_type": account_data_type},
             values={"stream_id": next_id, "content": content_json},
+            expected_revision_id=expected_revision_id,
         )
 
         # Ignored users get denormalized into a separate table as an optimisation.