From d240796dedcfae1f6929c1501e7e335df417cfaf Mon Sep 17 00:00:00 2001
From: David Baker <dave@matrix.org>
Date: Tue, 31 May 2016 17:20:07 +0100
Subject: [PATCH] Basic, un-cached support for secondary_directory_servers

---
 synapse/federation/federation_client.py | 21 ++++++++++++++++
 synapse/federation/transport/client.py  | 12 +++++++++
 synapse/federation/transport/server.py  |  2 +-
 synapse/handlers/room.py                | 33 ++++++++++++++++++++++++-
 synapse/rest/client/v1/room.py          |  3 ++-
 5 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/synapse/federation/federation_client.py b/synapse/federation/federation_client.py
index 37ee469fa..ba8d71c05 100644
--- a/synapse/federation/federation_client.py
+++ b/synapse/federation/federation_client.py
@@ -24,6 +24,7 @@ from synapse.api.errors import (
     CodeMessageException, HttpResponseException, SynapseError,
 )
 from synapse.util import unwrapFirstError
+from synapse.util.async import concurrently_execute
 from synapse.util.caches.expiringcache import ExpiringCache
 from synapse.util.logutils import log_function
 from synapse.events import FrozenEvent
@@ -550,6 +551,26 @@ class FederationClient(FederationBase):
 
         raise RuntimeError("Failed to send to any server.")
 
+    @defer.inlineCallbacks
+    def get_public_rooms(self, destinations):
+        results_by_server = {}
+
+        @defer.inlineCallbacks
+        def _get_result(s):
+            if s == self.server_name:
+                defer.returnValue()
+
+            try:
+                result = yield self.transport_layer.get_public_rooms(s)
+                results_by_server[s] = result
+            except:
+                logger.exception("Error getting room list from server %r", s)
+
+
+        yield concurrently_execute(_get_result, destinations, 3)
+
+        defer.returnValue(results_by_server)
+
     @defer.inlineCallbacks
     def query_auth(self, destination, room_id, event_id, local_auth):
         """
diff --git a/synapse/federation/transport/client.py b/synapse/federation/transport/client.py
index cd2841c4d..ebb698e27 100644
--- a/synapse/federation/transport/client.py
+++ b/synapse/federation/transport/client.py
@@ -224,6 +224,18 @@ class TransportLayerClient(object):
 
         defer.returnValue(response)
 
+    @defer.inlineCallbacks
+    @log_function
+    def get_public_rooms(self, remote_server):
+        path = PREFIX + "/publicRooms"
+
+        response = yield self.client.get_json(
+            destination=remote_server,
+            path=path,
+        )
+
+        defer.returnValue(response)
+
     @defer.inlineCallbacks
     @log_function
     def exchange_third_party_invite(self, destination, room_id, event_dict):
diff --git a/synapse/federation/transport/server.py b/synapse/federation/transport/server.py
index f23c02efd..da9e7a326 100644
--- a/synapse/federation/transport/server.py
+++ b/synapse/federation/transport/server.py
@@ -527,7 +527,7 @@ class PublicRoomList(BaseFederationServlet):
 
     @defer.inlineCallbacks
     def on_GET(self, request):
-        data = yield self.room_list_handler.get_public_room_list()
+        data = yield self.room_list_handler.get_local_public_room_list()
         defer.returnValue((200, data))
 
     # Avoid doing remote HS authorization checks which are done by default by
diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index 3d63b3c51..b0aa9fb51 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -345,7 +345,7 @@ class RoomListHandler(BaseHandler):
         super(RoomListHandler, self).__init__(hs)
         self.response_cache = ResponseCache()
 
-    def get_public_room_list(self):
+    def get_local_public_room_list(self):
         result = self.response_cache.get(())
         if not result:
             result = self.response_cache.set((), self._get_public_room_list())
@@ -427,6 +427,37 @@ class RoomListHandler(BaseHandler):
         # FIXME (erikj): START is no longer a valid value
         defer.returnValue({"start": "START", "end": "END", "chunk": results})
 
+    @defer.inlineCallbacks
+    def get_aggregated_public_room_list(self):
+        """
+        Get the public room list from this server and the servers
+        specified in the secondary_directory_servers config option.
+        XXX: Pagination...
+        """
+        federated_by_server = yield self.hs.get_replication_layer().get_public_rooms(
+            self.hs.config.secondary_directory_servers
+        )
+        public_rooms = yield self.get_local_public_room_list()
+
+        # keep track of which room IDs we've seen so we can de-dup
+        room_ids = set()
+
+        # tag all the ones in our list with our server name.
+        # Also add the them to the de-deping set
+        for room in public_rooms['chunk']:
+            room["server_name"] = self.hs.hostname
+            room_ids.add(room["room_id"])
+
+        # Now add the results from federation
+        for server_name, server_result in federated_by_server.items():
+            for room in server_result["chunk"]:
+                if room["room_id"] not in room_ids:
+                    room["server_name"] = server_name
+                    public_rooms["chunk"].append(room)
+                    room_ids.add(room["room_id"])
+
+        defer.returnValue(public_rooms)
+
 
 class RoomContextHandler(BaseHandler):
     @defer.inlineCallbacks
diff --git a/synapse/rest/client/v1/room.py b/synapse/rest/client/v1/room.py
index 2d22bbdaa..db52a1fc3 100644
--- a/synapse/rest/client/v1/room.py
+++ b/synapse/rest/client/v1/room.py
@@ -280,7 +280,8 @@ class PublicRoomListRestServlet(ClientV1RestServlet):
     @defer.inlineCallbacks
     def on_GET(self, request):
         handler = self.hs.get_room_list_handler()
-        data = yield handler.get_public_room_list()
+        data = yield handler.get_aggregated_public_room_list()
+
         defer.returnValue((200, data))