Add companion endpoint for backpagination of thread subscriptions

Author: reivilibre

synapse/handlers/sliding_sync/extensions.py

@@ -932,7 +932,7 @@ async def get_thread_subscriptions_extension_response(
 
         to_stream_id = to_token.thread_subscriptions_key
 
-        updates = await self.store.get_updated_thread_subscriptions_for_user(
+        updates = await self.store.get_latest_updated_thread_subscriptions_for_user(
             user_id=sync_config.user.to_string(),
             from_id=from_stream_id,
             to_id=to_stream_id,

synapse/rest/client/thread_subscriptions.py

@@ -1,19 +1,37 @@
 from http import HTTPStatus
-from typing import Tuple
+from typing import Dict, Tuple
+
+import attr
+from typing_extensions import TypeAlias
 
 from synapse._pydantic_compat import StrictBool
 from synapse.api.errors import Codes, NotFoundError, SynapseError
 from synapse.http.server import HttpServer
 from synapse.http.servlet import (
     RestServlet,
     parse_and_validate_json_object_from_request,
+    parse_integer,
+    parse_string,
 )
 from synapse.http.site import SynapseRequest
 from synapse.rest.client._base import client_patterns
 from synapse.server import HomeServer
-from synapse.types import JsonDict, RoomID
+from synapse.types import (
+    JsonDict,
+    RoomID,
+    SlidingSyncStreamToken,
+    ThreadSubscriptionsToken,
+)
+from synapse.types.handlers.sliding_sync import SlidingSyncResult
 from synapse.types.rest import RequestBodyModel
 
+_ThreadSubscription: TypeAlias = (
+    SlidingSyncResult.Extensions.ThreadSubscriptionsExtension.ThreadSubscription
+)
+_ThreadUnsubscription: TypeAlias = (
+    SlidingSyncResult.Extensions.ThreadSubscriptionsExtension.ThreadUnsubscription
+)
+
 
 class ThreadSubscriptionsRestServlet(RestServlet):
     PATTERNS = client_patterns(
@@ -93,6 +111,129 @@ async def on_DELETE(
         return HTTPStatus.OK, {}
 
 
+class ThreadSubscriptionsPaginationRestServlet(RestServlet):
+    PATTERNS = client_patterns(
+        "/io.element.msc4308/thread_subscriptions$",
+        unstable=True,
+        releases=(),
+    )
+    CATEGORY = "Thread Subscriptions requests (unstable)"
+
+    # Maximum number of thread subscriptions to return in one request.
+    MAX_LIMIT = 512
+
+    def __init__(self, hs: "HomeServer"):
+        self.auth = hs.get_auth()
+        self.is_mine = hs.is_mine
+        self.store = hs.get_datastores().main
+
+    async def on_GET(self, request: SynapseRequest) -> Tuple[int, JsonDict]:
+        requester = await self.auth.get_user_by_req(request)
+
+        limit = min(
+            parse_integer(request, "limit", default=100, negative=False),
+            ThreadSubscriptionsPaginationRestServlet.MAX_LIMIT,
+        )
+        from_end_opt = parse_string(request, "from", required=False)
+        to_start_opt = parse_string(request, "to", required=False)
+        _direction = parse_string(request, "dir", required=True, allowed_values=("b",))
+
+        if limit <= 0:
+            raise SynapseError(
+                HTTPStatus.BAD_REQUEST,
+                "limit must be greater than 0",
+                errcode=Codes.INVALID_PARAM,
+            )
+
+        if from_end_opt is not None:
+            try:
+                # because of backwards pagination, the `from` token is actually the
+                # bound closest to the end of the stream
+                end_stream_id = ThreadSubscriptionsToken.from_string(
+                    from_end_opt
+                ).stream_id
+            except ValueError:
+                raise SynapseError(
+                    HTTPStatus.BAD_REQUEST,
+                    "`from` is not a valid token",
+                    errcode=Codes.INVALID_PARAM,
+                )
+        else:
+            end_stream_id = self.store.get_max_thread_subscriptions_stream_id()
+
+        if to_start_opt is not None:
+            # because of backwards pagination, the `to` token is actually the
+            # bound closest to the start of the stream
+            try:
+                start_stream_id = ThreadSubscriptionsToken.from_string(
+                    to_start_opt
+                ).stream_id
+            except ValueError:
+                # we also accept sliding sync `pos` tokens on this parameter
+                try:
+                    sliding_sync_pos = await SlidingSyncStreamToken.from_string(
+                        self.store, to_start_opt
+                    )
+                    start_stream_id = (
+                        sliding_sync_pos.stream_token.thread_subscriptions_key
+                    )
+                except ValueError:
+                    raise SynapseError(
+                        HTTPStatus.BAD_REQUEST,
+                        "`to` is not a valid token",
+                        errcode=Codes.INVALID_PARAM,
+                    )
+        else:
+            # the start of time is ID 1; the lower bound is exclusive though
+            start_stream_id = 0
+
+        subscriptions = (
+            await self.store.get_latest_updated_thread_subscriptions_for_user(
+                requester.user.to_string(),
+                from_id=start_stream_id,
+                to_id=end_stream_id,
+                limit=limit,
+            )
+        )
+
+        subscribed_threads: Dict[str, Dict[str, JsonDict]] = {}
+        unsubscribed_threads: Dict[str, Dict[str, JsonDict]] = {}
+        for stream_id, room_id, thread_root_id, subscribed, automatic in subscriptions:
+            if subscribed:
+                subscribed_threads.setdefault(room_id, {})[thread_root_id] = (
+                    attr.asdict(
+                        _ThreadSubscription(
+                            automatic=automatic,
+                            bump_stamp=stream_id,
+                        )
+                    )
+                )
+            else:
+                unsubscribed_threads.setdefault(room_id, {})[thread_root_id] = (
+                    attr.asdict(_ThreadUnsubscription(bump_stamp=stream_id))
+                )
+
+        result: JsonDict = {}
+        if subscribed_threads:
+            result["subscribed"] = subscribed_threads
+        if unsubscribed_threads:
+            result["unsubscribed"] = unsubscribed_threads
+
+        if len(subscriptions) == limit:
+            # We hit the limit, so there might be more entries to return.
+            # Generate a new token that has moved backwards, ready for the next
+            # request.
+            min_returned_stream_id, _, _, _, _ = subscriptions[0]
+            result["end"] = ThreadSubscriptionsToken(
+                # We subtract one because the 'later in the stream' bound is inclusive,
+                # and we already saw the element at index 0.
+                stream_id=min_returned_stream_id - 1
+            ).to_string()
+
+        return HTTPStatus.OK, result
+
+
 def register_servlets(hs: "HomeServer", http_server: HttpServer) -> None:
     if hs.config.experimental.msc4306_enabled:
         ThreadSubscriptionsRestServlet(hs).register(http_server)
+        ThreadSubscriptionsPaginationRestServlet(hs).register(http_server)

synapse/storage/databases/main/thread_subscriptions.py

@@ -350,30 +350,38 @@ def get_updated_thread_subscriptions_txn(
             get_updated_thread_subscriptions_txn,
         )
 
-    async def get_updated_thread_subscriptions_for_user(
+    async def get_latest_updated_thread_subscriptions_for_user(
         self, user_id: str, from_id: int, to_id: int, limit: int
     ) -> List[Tuple[int, str, str, bool, Optional[bool]]]:
-        """Get updates to thread subscriptions for a specific user.
+        """Get the latest updates to thread subscriptions for a specific user.
 
         Args:
             user_id: The ID of the user
             from_id: The starting stream ID (exclusive)
             to_id: The ending stream ID (inclusive)
             limit: The maximum number of rows to return
+                If there are too many rows to return, rows from the start (closer to `from_id`)
+                will be omitted.
 
         Returns:
             A list of (stream_id, room_id, thread_root_event_id, subscribed, automatic) tuples.
+            The row with lowest `stream_id` is the first row.
         """
 
         def get_updated_thread_subscriptions_for_user_txn(
             txn: LoggingTransaction,
         ) -> List[Tuple[int, str, str, bool, Optional[bool]]]:
             sql = """
+                WITH the_updates AS (
+                    SELECT stream_id, room_id, event_id, subscribed, automatic
+                    FROM thread_subscriptions
+                    WHERE user_id = ? AND ? < stream_id AND stream_id <= ?
+                    ORDER BY stream_id DESC
+                    LIMIT ?
+                )
                 SELECT stream_id, room_id, event_id, subscribed, automatic
-                FROM thread_subscriptions
-                WHERE user_id = ? AND ? < stream_id AND stream_id <= ?
+                FROM the_updates
                 ORDER BY stream_id ASC
-                LIMIT ?
             """
 
             txn.execute(sql, (user_id, from_id, to_id, limit))

tests/rest/client/sliding_sync/test_extension_thread_subscriptions.py

@@ -353,6 +353,123 @@ def test_limit_parameter(self) -> None:
             len(thread_subscriptions["subscribed"][room_id]), 3, thread_subscriptions
         )
 
+    def test_limit_and_companion_backpagination(self) -> None:
+        """
+        Create 1 thread subscription, do a sync, create 4 more,
+        then sync with a limit of 2 and fill in the gap
+        using the companion /thread_subscriptions endpoint.
+        """
+
+        thread_root_ids: List[str] = []
+
+        def make_subscription() -> None:
+            thread_root_resp = self.helper.send(
+                room_id, body="Some thread root", tok=user1_tok
+            )
+            thread_root_ids.append(thread_root_resp["event_id"])
+            self._subscribe_to_thread(
+                user1_id, room_id, thread_root_ids[-1], automatic=False
+            )
+
+        user1_id = self.register_user("user1", "pass")
+        user1_tok = self.login(user1_id, "pass")
+        room_id = self.helper.create_room_as(user1_id, tok=user1_tok)
+
+        # get the baseline stream_id of the thread_subscriptions stream
+        # before we write any data.
+        # Required because the initial value differs between SQLite and Postgres.
+        base = self.store.get_max_thread_subscriptions_stream_id()
+
+        # Make our first subscription
+        make_subscription()
+
+        # Sync for the first time
+        sync_body = {
+            "lists": {},
+            "extensions": {EXT_NAME: {"enabled": True, "limit": 2}},
+        }
+
+        sync_resp, first_sync_pos = self.do_sync(sync_body, tok=user1_tok)
+
+        thread_subscriptions = sync_resp["extensions"][EXT_NAME]
+        self.assertEqual(
+            thread_subscriptions["subscribed"],
+            {
+                room_id: {
+                    thread_root_ids[0]: {"automatic": False, "bump_stamp": base + 1},
+                }
+            },
+        )
+
+        # Get our pos for the next sync
+        first_sync_pos = sync_resp["pos"]
+
+        # Create 4 more thread subsrciptions and subscribe to each
+        for _ in range(5):
+            make_subscription()
+
+        # Now sync again. Our limit is 2,
+        # so we should get the latest 2 subscriptions,
+        # with a gap of 3 more subscriptions in the middle
+        sync_resp, _pos = self.do_sync(sync_body, tok=user1_tok, since=first_sync_pos)
+
+        thread_subscriptions = sync_resp["extensions"][EXT_NAME]
+        self.assertEqual(
+            thread_subscriptions["subscribed"],
+            {
+                room_id: {
+                    thread_root_ids[4]: {"automatic": False, "bump_stamp": base + 5},
+                    thread_root_ids[5]: {"automatic": False, "bump_stamp": base + 6},
+                }
+            },
+        )
+        # 1st backpagination: expecting a page with 2 subscriptions
+        page, end_tok = self._do_backpaginate(
+            from_tok=thread_subscriptions["prev_batch"],
+            to_tok=first_sync_pos,
+            limit=2,
+            access_token=user1_tok,
+        )
+        self.assertIsNotNone(end_tok, "backpagination should continue")
+        self.assertEqual(
+            page["subscribed"],
+            {
+                room_id: {
+                    thread_root_ids[2]: {"automatic": False, "bump_stamp": base + 3},
+                    thread_root_ids[3]: {"automatic": False, "bump_stamp": base + 4},
+                }
+            },
+        )
+
+        # 2nd backpagination: expecting a page with only 1 subscription
+        # and no other token for further backpagination
+        assert end_tok is not None
+        page, end_tok = self._do_backpaginate(
+            from_tok=end_tok, to_tok=first_sync_pos, limit=2, access_token=user1_tok
+        )
+        self.assertIsNone(end_tok, "backpagination should have finished")
+        self.assertEqual(
+            page["subscribed"],
+            {
+                room_id: {
+                    thread_root_ids[1]: {"automatic": False, "bump_stamp": base + 2},
+                }
+            },
+        )
+
+    def _do_backpaginate(
+        self, *, from_tok: str, to_tok: str, limit: int, access_token: str
+    ) -> Tuple[JsonDict, Optional[str]]:
+        channel = self.make_request(
+            "GET",
+            "/_matrix/client/unstable/io.element.msc4308/thread_subscriptions"
+            f"?from={from_tok}&to={to_tok}&limit={limit}&dir=b",
+            access_token=access_token,
+        )
+
+        self.assertEqual(channel.code, HTTPStatus.OK, channel.json_body)
+        body = channel.json_body
+        return body, cast(Optional[str], body.get("end"))
 
     def _subscribe_to_thread(
         self, user_id: str, room_id: str, thread_root_id: str, automatic: bool

tests/storage/test_thread_subscriptions.py

@@ -182,7 +182,7 @@ def test_purge_thread_subscriptions_for_user(self) -> None:
         self._subscribe(self.other_thread_root_id, automatic=False)
 
         subscriptions = self.get_success(
-            self.store.get_updated_thread_subscriptions_for_user(
+            self.store.get_latest_updated_thread_subscriptions_for_user(
                 self.user_id,
                 from_id=0,
                 to_id=50,
@@ -205,7 +205,7 @@ def test_purge_thread_subscriptions_for_user(self) -> None:
 
         # Check user has no subscriptions
         subscriptions = self.get_success(
-            self.store.get_updated_thread_subscriptions_for_user(
+            self.store.get_latest_updated_thread_subscriptions_for_user(
                 self.user_id,
                 from_id=0,
                 to_id=50,
@@ -255,7 +255,7 @@ def test_get_updated_thread_subscriptions_for_user(self) -> None:
 
         # Get updates for main user
         updates = self.get_success(
-            self.store.get_updated_thread_subscriptions_for_user(
+            self.store.get_latest_updated_thread_subscriptions_for_user(
                 self.user_id, 0, stream_id2, 10
             )
         )
@@ -265,7 +265,7 @@ def test_get_updated_thread_subscriptions_for_user(self) -> None:
 
         # Get updates for other user
         updates = self.get_success(
-            self.store.get_updated_thread_subscriptions_for_user(
+            self.store.get_latest_updated_thread_subscriptions_for_user(
                 other_user_id, 0, max(stream_id1, stream_id2), 10
             )
         )