diff options
| -rw-r--r-- | synapse/handlers/user_directory.py | 45 | ||||
| -rw-r--r-- | synapse/storage/_base.py | 13 | ||||
| -rw-r--r-- | synapse/storage/user_directory.py | 52 | ||||
| -rw-r--r-- | tests/handlers/test_user_directory.py | 16 |
4 files changed, 74 insertions, 52 deletions
diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py index fc45123d0c..20a026e776 100644 --- a/synapse/handlers/user_directory.py +++ b/synapse/handlers/user_directory.py @@ -60,14 +60,16 @@ class UserDirectoryHandler(object): self.update_user_directory = hs.config.update_user_directory self.search_all_users = hs.config.user_directory_search_all_users + # If we're a worker, don't sleep when doing the initial room work, as it + # won't monopolise the master's CPU. + if hs.config.worker_app: + self.INITIAL_ROOM_SLEEP_MS = 0 + self.INITIAL_USER_SLEEP_MS = 0 + # When start up for the first time we need to populate the user_directory. # This is a set of user_id's we've inserted already self.initially_handled_users = set() - self.register_background_update_handler( - "users_in_public_rooms_initial", self._populate_users_in_public_rooms - ) - # The current position in the current_state_delta stream self.pos = None @@ -81,41 +83,6 @@ class UserDirectoryHandler(object): # we start populating the user directory self.clock.call_later(0, self.notify_new_event) - @defer.inlineCallbacks - def _populate_users_in_public_rooms(self, progress, batch_size): - """ - Populate the users_in_public_rooms table with the contents of the - users_who_share_public_rooms table. - """ - - def _fetch(txn): - sql = "SELECT DISTINCT other_user_id FROM users_who_share_public_rooms" - txn.execute(sql) - return txn.fetchall() - - users = yield self.store.runInteraction( - "populate_users_in_public_rooms_fetch", _fetch - ) - - if users: - - def _fill(txn): - self._simple_upsert_many_txn( - txn, - table="users_in_public_rooms", - key_names=["user_id"], - key_values=users, - value_names=(), - value_values=None, - ) - - users = yield self.store.runInteraction( - "populate_users_in_public_rooms_fill", _fill - ) - - yield self._end_background_update("users_in_public_rooms_initial") - defer.returnValue(1) - def search_users(self, user_id, search_term, limit): """Searches for users in directory diff --git a/synapse/storage/_base.py b/synapse/storage/_base.py index a0333d5309..7e3903859b 100644 --- a/synapse/storage/_base.py +++ b/synapse/storage/_base.py @@ -767,18 +767,25 @@ class SQLBaseStore(object): """ allvalues = {} allvalues.update(keyvalues) - allvalues.update(values) allvalues.update(insertion_values) + if not values: + latter = "NOTHING" + else: + allvalues.update(values) + latter = ( + "UPDATE SET " + ", ".join(k + "=EXCLUDED." + k for k in values) + ) + sql = ( "INSERT INTO %s (%s) VALUES (%s) " - "ON CONFLICT (%s) DO UPDATE SET %s" + "ON CONFLICT (%s) DO %s" ) % ( table, ", ".join(k for k in allvalues), ", ".join("?" for _ in allvalues), ", ".join(k for k in keyvalues), - ", ".join(k + "=EXCLUDED." + k for k in values), + latter ) txn.execute(sql, list(allvalues.values())) diff --git a/synapse/storage/user_directory.py b/synapse/storage/user_directory.py index 8f40277b50..a15366a117 100644 --- a/synapse/storage/user_directory.py +++ b/synapse/storage/user_directory.py @@ -22,16 +22,57 @@ from twisted.internet import defer from synapse.api.constants import EventTypes, JoinRules from synapse.storage.engines import PostgresEngine, Sqlite3Engine +from synapse.storage.background_updates import BackgroundUpdateStore from synapse.storage.state import StateFilter from synapse.types import get_domain_from_id, get_localpart_from_id from synapse.util.caches.descriptors import cached, cachedInlineCallbacks -from ._base import SQLBaseStore - logger = logging.getLogger(__name__) -class UserDirectoryStore(SQLBaseStore): +class UserDirectoryStore(BackgroundUpdateStore): + def __init__(self, dbconn, hs): + super(UserDirectoryStore, self).__init__(dbconn, hs) + + self.register_background_update_handler( + "users_in_public_rooms_initial", self._populate_users_in_public_rooms + ) + + + @defer.inlineCallbacks + def _populate_users_in_public_rooms(self, progress, batch_size): + """ + Populate the users_in_public_rooms table with the contents of the + users_who_share_public_rooms table. + """ + + def _fetch(txn): + sql = "SELECT DISTINCT other_user_id FROM users_who_share_public_rooms" + txn.execute(sql) + return txn.fetchall() + + users = yield self.runInteraction( + "populate_users_in_public_rooms_fetch", _fetch + ) + + if users: + def _fill(txn): + self._simple_upsert_many_txn( + txn, + table="users_in_public_rooms", + key_names=["user_id"], + key_values=users, + value_names=(), + value_values=None, + ) + + users = yield self.runInteraction( + "populate_users_in_public_rooms_fill", _fill + ) + + yield self._end_background_update("users_in_public_rooms_initial") + defer.returnValue(1) + @defer.inlineCallbacks def is_room_world_readable_or_publicly_joinable(self, room_id): """Check if the room is either world_readable or publically joinable @@ -353,8 +394,7 @@ class UserDirectoryStore(SQLBaseStore): txn, "users_in_public_rooms", keyvalues={"user_id": user_id}, - values={}, - desc="add_user_as_in_public_room", + values=None, ) for user_id, other_user_id in user_id_tuples: @@ -603,7 +643,7 @@ class UserDirectoryStore(SQLBaseStore): else: join_clause = """ LEFT JOIN ( - SELECT other_user_id AS user_id FROM users_who_share_public_rooms + SELECT user_id FROM users_in_public_rooms UNION SELECT other_user_id AS user_id FROM users_who_share_private_rooms WHERE user_id = ? diff --git a/tests/handlers/test_user_directory.py b/tests/handlers/test_user_directory.py index 0e0ac0a48b..7a78451a6d 100644 --- a/tests/handlers/test_user_directory.py +++ b/tests/handlers/test_user_directory.py @@ -116,12 +116,13 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase): # Check we have populated the database correctly. shares_public = self.get_users_who_share_public_rooms() shares_private = self.get_users_who_share_private_rooms() + public_users = self.get_users_in_public_rooms() self.assertEqual(shares_public, []) self.assertEqual( self._compress_shared(shares_private), set([(u1, u2, room), (u2, u1, room)]) ) - self.assertEqual(set(public_users), set([u1, u2])) + self.assertEqual(public_users, []) # We get one search result when searching for user2 by user1. s = self.get_success(self.handler.search_users(u1, "user2", 10)) @@ -145,7 +146,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase): self.assertEqual(shares_public, []) self.assertEqual(self._compress_shared(shares_private), set()) - self.assertEqual(public_users, [u1]) + self.assertEqual(public_users, []) # User1 now gets no search results for any of the other users. s = self.get_success(self.handler.search_users(u1, "user2", 10)) @@ -165,10 +166,10 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase): def get_users_in_public_rooms(self): return self.get_success( - self.store._simple_select_list( + self.store._simple_select_onecol( "users_in_public_rooms", None, - ["user_id"], + "user_id", ) ) @@ -214,9 +215,12 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase): shares_public = self.get_users_who_share_public_rooms() shares_private = self.get_users_who_share_private_rooms() + public_users = self.get_users_in_public_rooms() + # Nothing updated yet self.assertEqual(shares_private, []) self.assertEqual(shares_public, []) + self.assertEqual(public_users, []) # Reset the handled users caches self.handler.initially_handled_users = set() @@ -233,6 +237,7 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase): shares_public = self.get_users_who_share_public_rooms() shares_private = self.get_users_who_share_private_rooms() + public_users = self.get_users_in_public_rooms() # User 1 and User 2 share public rooms self.assertEqual( @@ -245,6 +250,9 @@ class UserDirectoryTestCase(unittest.HomeserverTestCase): set([(u1, u3, private_room), (u3, u1, private_room)]), ) + # User 1 and 2 are in public rooms + self.assertEqual(set(public_users), set([u1, u2])) + def test_search_all_users(self): """ Search all users = True means that a user does not have to share a |