Use stable identifiers for faster joins (#14832)

* Use new query param when requesting a partial join

* Read new query param when serving partial join

* Provide new field names when serving partial joins

* Read new field names from partial join response

* Changelog

6 files changed, 89 insertions, 24 deletions

diff --git a/changelog.d/14832.misc b/changelog.d/14832.misc
new file mode 100644
index 0000000000..61e7401e43
--- /dev/null
+++ b/changelog.d/14832.misc
@@ -0,0 +1 @@
+Faster joins: use stable identifiers from [MSC3706](https://github.com/matrix-org/matrix-spec-proposals/pull/3706).
diff --git a/synapse/federation/federation_server.py b/synapse/federation/federation_server.py
index bb20af6e91..c65dbf87fb 100644
--- a/synapse/federation/federation_server.py
+++ b/synapse/federation/federation_server.py
@@ -725,10 +725,12 @@ class FederationServer(FederationBase):
             "state": [p.get_pdu_json(time_now) for p in state_events],
             "auth_chain": [p.get_pdu_json(time_now) for p in auth_chain_events],
             "org.matrix.msc3706.partial_state": caller_supports_partial_state,
+            "members_omitted": caller_supports_partial_state,
         }
 
         if servers_in_room is not None:
             resp["org.matrix.msc3706.servers_in_room"] = list(servers_in_room)
+            resp["servers_in_room"] = list(servers_in_room)
 
         return resp
 
diff --git a/synapse/federation/transport/client.py b/synapse/federation/transport/client.py
index 77f1f39cac..c8471d4cf7 100644
--- a/synapse/federation/transport/client.py
+++ b/synapse/federation/transport/client.py
@@ -357,6 +357,7 @@ class TransportLayerClient:
         if self._faster_joins_enabled:
             # lazy-load state on join
             query_params["org.matrix.msc3706.partial_state"] = "true"
+            query_params["omit_members"] = "true"
 
         return await self.client.put_json(
             destination=destination,
@@ -909,6 +910,14 @@ class SendJoinParser(ByteParser[SendJoinResponse]):
                     use_float="True",
                 )
             )
+            # The stable field name comes last, so it "wins" if the fields disagree
+            self._coros.append(
+                ijson.items_coro(
+                    _partial_state_parser(self._response),
+                    "members_omitted",
+                    use_float="True",
+                )
+            )
 
             self._coros.append(
                 ijson.items_coro(
@@ -918,6 +927,15 @@ class SendJoinParser(ByteParser[SendJoinResponse]):
                 )
             )
 
+            # Again, stable field name comes last
+            self._coros.append(
+                ijson.items_coro(
+                    _servers_in_room_parser(self._response),
+                    "servers_in_room",
+                    use_float="True",
+                )
+            )
+
     def write(self, data: bytes) -> int:
         for c in self._coros:
             c.send(data)
diff --git a/synapse/federation/transport/server/federation.py b/synapse/federation/transport/server/federation.py
index 53e77b4bb6..c0a700905b 100644
--- a/synapse/federation/transport/server/federation.py
+++ b/synapse/federation/transport/server/federation.py
@@ -437,9 +437,16 @@ class FederationV2SendJoinServlet(BaseFederationServerServlet):
 
         partial_state = False
         if self._msc3706_enabled:
-            partial_state = parse_boolean_from_args(
-                query, "org.matrix.msc3706.partial_state", default=False
-            )
+            # The stable query parameter wins, if it disagrees with the unstable
+            # parameter for some reason.
+            stable_param = parse_boolean_from_args(query, "omit_members", default=None)
+            if stable_param is not None:
+                partial_state = stable_param
+            else:
+                partial_state = parse_boolean_from_args(
+                    query, "org.matrix.msc3706.partial_state", default=False
+                )
+
         result = await self.handler.on_send_join_request(
             origin, content, room_id, caller_supports_partial_state=partial_state
         )
diff --git a/tests/federation/test_federation_server.py b/tests/federation/test_federation_server.py
index 177e5b5afc..27770304be 100644
--- a/tests/federation/test_federation_server.py
+++ b/tests/federation/test_federation_server.py
@@ -224,7 +224,7 @@ class SendJoinFederationTests(unittest.FederatingHomeserverTestCase):
         )
         channel = self.make_signed_federation_request(
             "PUT",
-            f"/_matrix/federation/v2/send_join/{self._room_id}/x?org.matrix.msc3706.partial_state=true",
+            f"/_matrix/federation/v2/send_join/{self._room_id}/x?omit_members=true",
             content=join_event_dict,
         )
         self.assertEqual(channel.code, HTTPStatus.OK, channel.json_body)
diff --git a/tests/federation/transport/test_client.py b/tests/federation/transport/test_client.py
index b84c74fc0e..c90635e0a0 100644
--- a/tests/federation/transport/test_client.py
+++ b/tests/federation/transport/test_client.py
@@ -13,12 +13,14 @@
 # limitations under the License.
 
 import json
+from typing import List, Optional
 from unittest.mock import Mock
 
 import ijson.common
 
 from synapse.api.room_versions import RoomVersions
 from synapse.federation.transport.client import SendJoinParser
+from synapse.types import JsonDict
 from synapse.util import ExceptionBundle
 
 from tests.unittest import TestCase
@@ -71,33 +73,68 @@ class SendJoinParserTestCase(TestCase):
 
     def test_partial_state(self) -> None:
         """Check that the partial_state flag is correctly parsed"""
-        parser = SendJoinParser(RoomVersions.V1, False)
-        response = {
-            "org.matrix.msc3706.partial_state": True,
-        }
 
-        serialised_response = json.dumps(response).encode()
+        def parse(response: JsonDict) -> bool:
+            parser = SendJoinParser(RoomVersions.V1, False)
+            serialised_response = json.dumps(response).encode()
 
-        # Send data to the parser
-        parser.write(serialised_response)
+            # Send data to the parser
+            parser.write(serialised_response)
 
-        # Retrieve and check the parsed SendJoinResponse
-        parsed_response = parser.finish()
-        self.assertTrue(parsed_response.partial_state)
+            # Retrieve and check the parsed SendJoinResponse
+            parsed_response = parser.finish()
+            return parsed_response.partial_state
 
-    def test_servers_in_room(self) -> None:
-        """Check that the servers_in_room field is correctly parsed"""
-        parser = SendJoinParser(RoomVersions.V1, False)
-        response = {"org.matrix.msc3706.servers_in_room": ["hs1", "hs2"]}
+        self.assertTrue(parse({"members_omitted": True}))
+        self.assertTrue(parse({"org.matrix.msc3706.partial_state": True}))
 
-        serialised_response = json.dumps(response).encode()
+        self.assertFalse(parse({"members_omitted": False}))
+        self.assertFalse(parse({"org.matrix.msc3706.partial_state": False}))
 
-        # Send data to the parser
-        parser.write(serialised_response)
+        # If there's a conflict, the stable field wins.
+        self.assertTrue(
+            parse({"members_omitted": True, "org.matrix.msc3706.partial_state": False})
+        )
+        self.assertFalse(
+            parse({"members_omitted": False, "org.matrix.msc3706.partial_state": True})
+        )
 
-        # Retrieve and check the parsed SendJoinResponse
-        parsed_response = parser.finish()
-        self.assertEqual(parsed_response.servers_in_room, ["hs1", "hs2"])
+    def test_servers_in_room(self) -> None:
+        """Check that the servers_in_room field is correctly parsed"""
+
+        def parse(response: JsonDict) -> Optional[List[str]]:
+            parser = SendJoinParser(RoomVersions.V1, False)
+            serialised_response = json.dumps(response).encode()
+
+            # Send data to the parser
+            parser.write(serialised_response)
+
+            # Retrieve and check the parsed SendJoinResponse
+            parsed_response = parser.finish()
+            return parsed_response.servers_in_room
+
+        self.assertEqual(
+            parse({"org.matrix.msc3706.servers_in_room": ["hs1", "hs2"]}),
+            ["hs1", "hs2"],
+        )
+        self.assertEqual(parse({"servers_in_room": ["example.com"]}), ["example.com"])
+
+        # If both are provided, the stable identifier should win
+        self.assertEqual(
+            parse(
+                {
+                    "org.matrix.msc3706.servers_in_room": ["old"],
+                    "servers_in_room": ["new"],
+                }
+            ),
+            ["new"],
+        )
+
+        # And lastly, we should be able to tell if neither field was present.
+        self.assertEqual(
+            parse({}),
+            None,
+        )
 
     def test_errors_closing_coroutines(self) -> None:
         """Check we close all coroutines, even if closing the first raises an Exception.