class SlavedDeviceInboxStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceInboxStore, self).__init__(db_conn, hs) self._device_inbox_id_gen = SlavedIdTracker( db_conn, "device_max_stream_id", "stream_id", ) self._device_inbox_stream_cache = StreamChangeCache( "DeviceInboxStreamChangeCache", self._device_inbox_id_gen.get_current_token()) get_to_device_stream_token = DataStore.get_to_device_stream_token.__func__ get_new_messages_for_device = DataStore.get_new_messages_for_device.__func__ delete_messages_for_device = DataStore.delete_messages_for_device.__func__ def stream_positions(self): result = super(SlavedDeviceInboxStore, self).stream_positions() result["to_device"] = self._device_inbox_id_gen.get_current_token() return result def process_replication(self, result): stream = result.get("to_device") if stream: self._device_inbox_id_gen.advance(int(stream["position"])) for row in stream["rows"]: stream_id = row[0] user_id = row[1] self._device_inbox_stream_cache.entity_has_changed( user_id, stream_id) return super(SlavedDeviceInboxStore, self).process_replication(result)
def test_has_any_entity_changed(self): """ StreamChangeCache.has_any_entity_changed will return True if any entities have been changed since the provided stream position, and False if they have not. If the cache has entries and the provided stream position is before it, it will return True, otherwise False if the cache has no entries. """ cache = StreamChangeCache("#test", 1) # With no entities, it returns False for the past, present, and future. self.assertFalse(cache.has_any_entity_changed(0)) self.assertFalse(cache.has_any_entity_changed(1)) self.assertFalse(cache.has_any_entity_changed(2)) # We add an entity cache.entity_has_changed("*****@*****.**", 2) # With an entity, it returns True for the past, the stream start # position, and False for the stream position the entity was changed # on and ones after it. self.assertTrue(cache.has_any_entity_changed(0)) self.assertTrue(cache.has_any_entity_changed(1)) self.assertFalse(cache.has_any_entity_changed(2)) self.assertFalse(cache.has_any_entity_changed(3))
class SlavedGroupServerStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedGroupServerStore, self).__init__(db_conn, hs) self.hs = hs self._group_updates_id_gen = SlavedIdTracker( db_conn, "local_group_updates", "stream_id", ) self._group_updates_stream_cache = StreamChangeCache( "_group_updates_stream_cache", self._group_updates_id_gen.get_current_token(), ) get_groups_changes_for_user = DataStore.get_groups_changes_for_user.__func__ get_group_stream_token = DataStore.get_group_stream_token.__func__ get_all_groups_for_user = DataStore.get_all_groups_for_user.__func__ def stream_positions(self): result = super(SlavedGroupServerStore, self).stream_positions() result["groups"] = self._group_updates_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "groups": self._group_updates_id_gen.advance(token) for row in rows: self._group_updates_stream_cache.entity_has_changed( row.user_id, token ) return super(SlavedGroupServerStore, self).process_replication_rows( stream_name, token, rows )
class SlavedGroupServerStore(GroupServerWorkerStore, BaseSlavedStore): def __init__(self, database: DatabasePool, db_conn, hs): super(SlavedGroupServerStore, self).__init__(database, db_conn, hs) self.hs = hs self._group_updates_id_gen = SlavedIdTracker(db_conn, "local_group_updates", "stream_id") self._group_updates_stream_cache = StreamChangeCache( "_group_updates_stream_cache", self._group_updates_id_gen.get_current_token(), ) def get_group_stream_token(self): return self._group_updates_id_gen.get_current_token() def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == GroupServerStream.NAME: self._group_updates_id_gen.advance(instance_name, token) for row in rows: self._group_updates_stream_cache.entity_has_changed( row.user_id, token) return super().process_replication_rows(stream_name, instance_name, token, rows)
class SlavedGroupServerStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedGroupServerStore, self).__init__(db_conn, hs) self.hs = hs self._group_updates_id_gen = SlavedIdTracker( db_conn, "local_group_updates", "stream_id", ) self._group_updates_stream_cache = StreamChangeCache( "_group_updates_stream_cache", self._group_updates_id_gen.get_current_token(), ) get_groups_changes_for_user = __func__( DataStore.get_groups_changes_for_user) get_group_stream_token = __func__(DataStore.get_group_stream_token) get_all_groups_for_user = __func__(DataStore.get_all_groups_for_user) def stream_positions(self): result = super(SlavedGroupServerStore, self).stream_positions() result["groups"] = self._group_updates_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "groups": self._group_updates_id_gen.advance(token) for row in rows: self._group_updates_stream_cache.entity_has_changed( row.user_id, token) return super(SlavedGroupServerStore, self).process_replication_rows(stream_name, token, rows)
class SlavedDeviceInboxStore(DeviceInboxWorkerStore, BaseSlavedStore): def __init__(self, database: Database, db_conn, hs): super(SlavedDeviceInboxStore, self).__init__(database, db_conn, hs) self._device_inbox_id_gen = SlavedIdTracker( db_conn, "device_max_stream_id", "stream_id" ) self._device_inbox_stream_cache = StreamChangeCache( "DeviceInboxStreamChangeCache", self._device_inbox_id_gen.get_current_token(), ) self._device_federation_outbox_stream_cache = StreamChangeCache( "DeviceFederationOutboxStreamChangeCache", self._device_inbox_id_gen.get_current_token(), ) self._last_device_delete_cache = ExpiringCache( cache_name="last_device_delete_cache", clock=self._clock, max_len=10000, expiry_ms=30 * 60 * 1000, ) def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == ToDeviceStream.NAME: self._device_inbox_id_gen.advance(token) for row in rows: if row.entity.startswith("@"): self._device_inbox_stream_cache.entity_has_changed( row.entity, token ) else: self._device_federation_outbox_stream_cache.entity_has_changed( row.entity, token ) return super().process_replication_rows(stream_name, instance_name, token, rows)
def test_has_entity_changed(self): """ StreamChangeCache.entity_has_changed will mark entities as changed, and has_entity_changed will observe the changed entities. """ cache = StreamChangeCache("#test", 3) cache.entity_has_changed("*****@*****.**", 6) cache.entity_has_changed("*****@*****.**", 7) # If it's been changed after that stream position, return True self.assertTrue(cache.has_entity_changed("*****@*****.**", 4)) self.assertTrue(cache.has_entity_changed("*****@*****.**", 4)) # If it's been changed at that stream position, return False self.assertFalse(cache.has_entity_changed("*****@*****.**", 6)) # If there's no changes after that stream position, return False self.assertFalse(cache.has_entity_changed("*****@*****.**", 7)) # If the entity does not exist, return False. self.assertFalse(cache.has_entity_changed("*****@*****.**", 7)) # If we request before the stream cache's earliest known position, # return True, whether it's a known entity or not. self.assertTrue(cache.has_entity_changed("*****@*****.**", 0)) self.assertTrue(cache.has_entity_changed("*****@*****.**", 0))
class SlavedPresenceStore(BaseSlavedStore): def __init__(self, database: DatabasePool, db_conn, hs): super(SlavedPresenceStore, self).__init__(database, db_conn, hs) self._presence_id_gen = SlavedIdTracker(db_conn, "presence_stream", "stream_id") self._presence_on_startup = self._get_active_presence( db_conn) # type: ignore self.presence_stream_cache = StreamChangeCache( "PresenceStreamChangeCache", self._presence_id_gen.get_current_token()) _get_active_presence = DataStore._get_active_presence take_presence_startup_info = DataStore.take_presence_startup_info _get_presence_for_user = PresenceStore.__dict__["_get_presence_for_user"] get_presence_for_users = PresenceStore.__dict__["get_presence_for_users"] def get_current_presence_token(self): return self._presence_id_gen.get_current_token() def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == PresenceStream.NAME: self._presence_id_gen.advance(instance_name, token) for row in rows: self.presence_stream_cache.entity_has_changed( row.user_id, token) self._get_presence_for_user.invalidate((row.user_id, )) return super().process_replication_rows(stream_name, instance_name, token, rows)
class SlavedReceiptsStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedReceiptsStore, self).__init__(db_conn, hs) self._receipts_id_gen = SlavedIdTracker( db_conn, "receipts_linearized", "stream_id" ) self._receipts_stream_cache = StreamChangeCache( "ReceiptsRoomChangeCache", self._receipts_id_gen.get_current_token() ) get_receipts_for_user = ReceiptsStore.__dict__["get_receipts_for_user"] get_linearized_receipts_for_room = ( ReceiptsStore.__dict__["get_linearized_receipts_for_room"] ) _get_linearized_receipts_for_rooms = ( ReceiptsStore.__dict__["_get_linearized_receipts_for_rooms"] ) get_last_receipt_event_id_for_user = ( ReceiptsStore.__dict__["get_last_receipt_event_id_for_user"] ) get_max_receipt_stream_id = DataStore.get_max_receipt_stream_id.__func__ get_all_updated_receipts = DataStore.get_all_updated_receipts.__func__ get_linearized_receipts_for_rooms = ( DataStore.get_linearized_receipts_for_rooms.__func__ ) def stream_positions(self): result = super(SlavedReceiptsStore, self).stream_positions() result["receipts"] = self._receipts_id_gen.get_current_token() return result def invalidate_caches_for_receipt(self, room_id, receipt_type, user_id): self.get_receipts_for_user.invalidate((user_id, receipt_type)) self.get_linearized_receipts_for_room.invalidate_many((room_id,)) self.get_last_receipt_event_id_for_user.invalidate( (user_id, room_id, receipt_type) ) def process_replication_rows(self, stream_name, token, rows): if stream_name == "receipts": self._receipts_id_gen.advance(token) for row in rows: self.invalidate_caches_for_receipt( row.room_id, row.receipt_type, row.user_id ) self._receipts_stream_cache.entity_has_changed(row.room_id, token) return super(SlavedReceiptsStore, self).process_replication_rows( stream_name, token, rows )
class SlavedReceiptsStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedReceiptsStore, self).__init__(db_conn, hs) self._receipts_id_gen = SlavedIdTracker( db_conn, "receipts_linearized", "stream_id" ) self._receipts_stream_cache = StreamChangeCache( "ReceiptsRoomChangeCache", self._receipts_id_gen.get_current_token() ) get_receipts_for_user = ReceiptsStore.__dict__["get_receipts_for_user"] get_linearized_receipts_for_room = ( ReceiptsStore.__dict__["get_linearized_receipts_for_room"] ) _get_linearized_receipts_for_rooms = ( ReceiptsStore.__dict__["_get_linearized_receipts_for_rooms"] ) get_last_receipt_event_id_for_user = ( ReceiptsStore.__dict__["get_last_receipt_event_id_for_user"] ) get_max_receipt_stream_id = DataStore.get_max_receipt_stream_id.__func__ get_all_updated_receipts = DataStore.get_all_updated_receipts.__func__ get_linearized_receipts_for_rooms = ( DataStore.get_linearized_receipts_for_rooms.__func__ ) def stream_positions(self): result = super(SlavedReceiptsStore, self).stream_positions() result["receipts"] = self._receipts_id_gen.get_current_token() return result def process_replication(self, result): stream = result.get("receipts") if stream: self._receipts_id_gen.advance(int(stream["position"])) for row in stream["rows"]: position, room_id, receipt_type, user_id = row[:4] self.invalidate_caches_for_receipt(room_id, receipt_type, user_id) self._receipts_stream_cache.entity_has_changed(room_id, position) return super(SlavedReceiptsStore, self).process_replication(result) def invalidate_caches_for_receipt(self, room_id, receipt_type, user_id): self.get_receipts_for_user.invalidate((user_id, receipt_type)) self.get_linearized_receipts_for_room.invalidate_many((room_id,)) self.get_last_receipt_event_id_for_user.invalidate( (user_id, room_id, receipt_type) )
class SlavedDeviceStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceStore, self).__init__(db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker( db_conn, "device_lists_stream", "stream_id", ) device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max, ) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max, ) get_device_stream_token = DataStore.get_device_stream_token.__func__ get_user_whose_devices_changed = DataStore.get_user_whose_devices_changed.__func__ get_devices_by_remote = DataStore.get_devices_by_remote.__func__ _get_devices_by_remote_txn = DataStore._get_devices_by_remote_txn.__func__ _get_e2e_device_keys_txn = DataStore._get_e2e_device_keys_txn.__func__ mark_as_sent_devices_by_remote = DataStore.mark_as_sent_devices_by_remote.__func__ _mark_as_sent_devices_by_remote_txn = ( DataStore._mark_as_sent_devices_by_remote_txn.__func__) def stream_positions(self): result = super(SlavedDeviceStore, self).stream_positions() result["device_lists"] = self._device_list_id_gen.get_current_token() return result def process_replication(self, result): stream = result.get("device_lists") if stream: self._device_list_id_gen.advance(int(stream["position"])) for row in stream["rows"]: stream_id = row[0] user_id = row[1] destination = row[2] self._device_list_stream_cache.entity_has_changed( user_id, stream_id) if destination: self._device_list_federation_stream_cache.entity_has_changed( destination, stream_id) return super(SlavedDeviceStore, self).process_replication(result)
class SlavedDeviceStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceStore, self).__init__(db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker( db_conn, "device_lists_stream", "stream_id", ) device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max, ) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max, ) get_device_stream_token = __func__(DataStore.get_device_stream_token) get_user_whose_devices_changed = __func__( DataStore.get_user_whose_devices_changed) get_devices_by_remote = __func__(DataStore.get_devices_by_remote) _get_devices_by_remote_txn = __func__(DataStore._get_devices_by_remote_txn) _get_e2e_device_keys_txn = __func__(DataStore._get_e2e_device_keys_txn) mark_as_sent_devices_by_remote = __func__( DataStore.mark_as_sent_devices_by_remote) _mark_as_sent_devices_by_remote_txn = (__func__( DataStore._mark_as_sent_devices_by_remote_txn)) count_e2e_one_time_keys = EndToEndKeyStore.__dict__[ "count_e2e_one_time_keys"] def stream_positions(self): result = super(SlavedDeviceStore, self).stream_positions() result["device_lists"] = self._device_list_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "device_lists": self._device_list_id_gen.advance(token) for row in rows: self._device_list_stream_cache.entity_has_changed( row.user_id, token) if row.destination: self._device_list_federation_stream_cache.entity_has_changed( row.destination, token) return super(SlavedDeviceStore, self).process_replication_rows(stream_name, token, rows)
def test_entity_has_changed_pops_off_start(self): """ StreamChangeCache.entity_has_changed will respect the max size and purge the oldest items upon reaching that max size. """ cache = StreamChangeCache("#test", 1, max_size=2) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) # The cache is at the max size, 2 self.assertEqual(len(cache._cache), 2) # The oldest item has been popped off self.assertTrue("*****@*****.**" not in cache._entity_to_key) self.assertEqual( cache.get_all_entities_changed(2), ["*****@*****.**", "*****@*****.**"], ) self.assertIsNone(cache.get_all_entities_changed(1)) # If we update an existing entity, it keeps the two existing entities cache.entity_has_changed("*****@*****.**", 5) self.assertEqual( {"*****@*****.**", "*****@*****.**"}, set(cache._entity_to_key) ) self.assertEqual( cache.get_all_entities_changed(2), ["*****@*****.**", "*****@*****.**"], ) self.assertIsNone(cache.get_all_entities_changed(1))
class SlavedDeviceInboxStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceInboxStore, self).__init__(db_conn, hs) self._device_inbox_id_gen = SlavedIdTracker( db_conn, "device_max_stream_id", "stream_id", ) self._device_inbox_stream_cache = StreamChangeCache( "DeviceInboxStreamChangeCache", self._device_inbox_id_gen.get_current_token()) self._device_federation_outbox_stream_cache = StreamChangeCache( "DeviceFederationOutboxStreamChangeCache", self._device_inbox_id_gen.get_current_token()) self._last_device_delete_cache = ExpiringCache( cache_name="last_device_delete_cache", clock=self._clock, max_len=10000, expiry_ms=30 * 60 * 1000, ) get_to_device_stream_token = DataStore.get_to_device_stream_token.__func__ get_new_messages_for_device = DataStore.get_new_messages_for_device.__func__ get_new_device_msgs_for_remote = DataStore.get_new_device_msgs_for_remote.__func__ delete_messages_for_device = DataStore.delete_messages_for_device.__func__ delete_device_msgs_for_remote = DataStore.delete_device_msgs_for_remote.__func__ def stream_positions(self): result = super(SlavedDeviceInboxStore, self).stream_positions() result["to_device"] = self._device_inbox_id_gen.get_current_token() return result def process_replication(self, result): stream = result.get("to_device") if stream: self._device_inbox_id_gen.advance(int(stream["position"])) for row in stream["rows"]: stream_id = row[0] entity = row[1] if entity.startswith("@"): self._device_inbox_stream_cache.entity_has_changed( entity, stream_id) else: self._device_federation_outbox_stream_cache.entity_has_changed( entity, stream_id) return super(SlavedDeviceInboxStore, self).process_replication(result)
class SlavedDeviceStore(EndToEndKeyWorkerStore, DeviceWorkerStore, BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceStore, self).__init__(db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker( db_conn, "device_lists_stream", "stream_id", ) device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max, ) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max, ) def stream_positions(self): result = super(SlavedDeviceStore, self).stream_positions() result["device_lists"] = self._device_list_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "device_lists": self._device_list_id_gen.advance(token) for row in rows: self._invalidate_caches_for_devices( token, row.user_id, row.destination, ) return super(SlavedDeviceStore, self).process_replication_rows(stream_name, token, rows) def _invalidate_caches_for_devices(self, token, user_id, destination): self._device_list_stream_cache.entity_has_changed(user_id, token) if destination: self._device_list_federation_stream_cache.entity_has_changed( destination, token) self._get_cached_devices_for_user.invalidate((user_id, )) self._get_cached_user_device.invalidate_many((user_id, )) self.get_device_list_last_stream_id_for_remote.invalidate((user_id, ))
class SlavedDeviceInboxStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceInboxStore, self).__init__(db_conn, hs) self._device_inbox_id_gen = SlavedIdTracker( db_conn, "device_max_stream_id", "stream_id", ) self._device_inbox_stream_cache = StreamChangeCache( "DeviceInboxStreamChangeCache", self._device_inbox_id_gen.get_current_token() ) self._device_federation_outbox_stream_cache = StreamChangeCache( "DeviceFederationOutboxStreamChangeCache", self._device_inbox_id_gen.get_current_token() ) self._last_device_delete_cache = ExpiringCache( cache_name="last_device_delete_cache", clock=self._clock, max_len=10000, expiry_ms=30 * 60 * 1000, ) get_to_device_stream_token = DataStore.get_to_device_stream_token.__func__ get_new_messages_for_device = DataStore.get_new_messages_for_device.__func__ get_new_device_msgs_for_remote = DataStore.get_new_device_msgs_for_remote.__func__ delete_messages_for_device = DataStore.delete_messages_for_device.__func__ delete_device_msgs_for_remote = DataStore.delete_device_msgs_for_remote.__func__ def stream_positions(self): result = super(SlavedDeviceInboxStore, self).stream_positions() result["to_device"] = self._device_inbox_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "to_device": self._device_inbox_id_gen.advance(token) for row in rows: if row.entity.startswith("@"): self._device_inbox_stream_cache.entity_has_changed( row.entity, token ) else: self._device_federation_outbox_stream_cache.entity_has_changed( row.entity, token ) return super(SlavedDeviceInboxStore, self).process_replication_rows( stream_name, token, rows )
class SlavedDeviceStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceStore, self).__init__(db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker( db_conn, "device_lists_stream", "stream_id", ) device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max, ) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max, ) get_device_stream_token = __func__(DataStore.get_device_stream_token) get_user_whose_devices_changed = __func__(DataStore.get_user_whose_devices_changed) get_devices_by_remote = __func__(DataStore.get_devices_by_remote) _get_devices_by_remote_txn = __func__(DataStore._get_devices_by_remote_txn) _get_e2e_device_keys_txn = __func__(DataStore._get_e2e_device_keys_txn) mark_as_sent_devices_by_remote = __func__(DataStore.mark_as_sent_devices_by_remote) _mark_as_sent_devices_by_remote_txn = ( __func__(DataStore._mark_as_sent_devices_by_remote_txn) ) count_e2e_one_time_keys = EndToEndKeyStore.__dict__["count_e2e_one_time_keys"] def stream_positions(self): result = super(SlavedDeviceStore, self).stream_positions() result["device_lists"] = self._device_list_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "device_lists": self._device_list_id_gen.advance(token) for row in rows: self._device_list_stream_cache.entity_has_changed( row.user_id, token ) if row.destination: self._device_list_federation_stream_cache.entity_has_changed( row.destination, token ) return super(SlavedDeviceStore, self).process_replication_rows( stream_name, token, rows )
class SlavedDeviceInboxStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceInboxStore, self).__init__(db_conn, hs) self._device_inbox_id_gen = SlavedIdTracker( db_conn, "device_max_stream_id", "stream_id", ) self._device_inbox_stream_cache = StreamChangeCache( "DeviceInboxStreamChangeCache", self._device_inbox_id_gen.get_current_token()) self._device_federation_outbox_stream_cache = StreamChangeCache( "DeviceFederationOutboxStreamChangeCache", self._device_inbox_id_gen.get_current_token()) self._last_device_delete_cache = ExpiringCache( cache_name="last_device_delete_cache", clock=self._clock, max_len=10000, expiry_ms=30 * 60 * 1000, ) get_to_device_stream_token = __func__(DataStore.get_to_device_stream_token) get_new_messages_for_device = __func__( DataStore.get_new_messages_for_device) get_new_device_msgs_for_remote = __func__( DataStore.get_new_device_msgs_for_remote) delete_messages_for_device = __func__(DataStore.delete_messages_for_device) delete_device_msgs_for_remote = __func__( DataStore.delete_device_msgs_for_remote) def stream_positions(self): result = super(SlavedDeviceInboxStore, self).stream_positions() result["to_device"] = self._device_inbox_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "to_device": self._device_inbox_id_gen.advance(token) for row in rows: if row.entity.startswith("@"): self._device_inbox_stream_cache.entity_has_changed( row.entity, token) else: self._device_federation_outbox_stream_cache.entity_has_changed( row.entity, token) return super(SlavedDeviceInboxStore, self).process_replication_rows(stream_name, token, rows)
class SlavedDeviceStore(EndToEndKeyWorkerStore, DeviceWorkerStore, BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceStore, self).__init__(db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker( db_conn, "device_lists_stream", "stream_id", ) device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max, ) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max, ) def stream_positions(self): result = super(SlavedDeviceStore, self).stream_positions() result["device_lists"] = self._device_list_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "device_lists": self._device_list_id_gen.advance(token) for row in rows: self._invalidate_caches_for_devices( token, row.user_id, row.destination, ) return super(SlavedDeviceStore, self).process_replication_rows( stream_name, token, rows ) def _invalidate_caches_for_devices(self, token, user_id, destination): self._device_list_stream_cache.entity_has_changed( user_id, token ) if destination: self._device_list_federation_stream_cache.entity_has_changed( destination, token ) self._get_cached_devices_for_user.invalidate((user_id,)) self._get_cached_user_device.invalidate_many((user_id,)) self.get_device_list_last_stream_id_for_remote.invalidate((user_id,))
class SlavedPushRuleStore(SlavedEventStore): def __init__(self, db_conn, hs): super(SlavedPushRuleStore, self).__init__(db_conn, hs) self._push_rules_stream_id_gen = SlavedIdTracker( db_conn, "push_rules_stream", "stream_id", ) self.push_rules_stream_cache = StreamChangeCache( "PushRulesStreamChangeCache", self._push_rules_stream_id_gen.get_current_token(), ) get_push_rules_for_user = PushRuleStore.__dict__["get_push_rules_for_user"] get_push_rules_enabled_for_user = ( PushRuleStore.__dict__["get_push_rules_enabled_for_user"] ) have_push_rules_changed_for_user = ( DataStore.have_push_rules_changed_for_user.__func__ ) def get_push_rules_stream_token(self): return ( self._push_rules_stream_id_gen.get_current_token(), self._stream_id_gen.get_current_token(), ) def stream_positions(self): result = super(SlavedPushRuleStore, self).stream_positions() result["push_rules"] = self._push_rules_stream_id_gen.get_current_token() return result def process_replication(self, result): stream = result.get("push_rules") if stream: for row in stream["rows"]: position = row[0] user_id = row[2] self.get_push_rules_for_user.invalidate((user_id,)) self.get_push_rules_enabled_for_user.invalidate((user_id,)) self.push_rules_stream_cache.entity_has_changed( user_id, position ) self._push_rules_stream_id_gen.advance(int(stream["position"])) return super(SlavedPushRuleStore, self).process_replication(result)
class SlavedPushRuleStore(SlavedEventStore): def __init__(self, db_conn, hs): super(SlavedPushRuleStore, self).__init__(db_conn, hs) self._push_rules_stream_id_gen = SlavedIdTracker( db_conn, "push_rules_stream", "stream_id", ) self.push_rules_stream_cache = StreamChangeCache( "PushRulesStreamChangeCache", self._push_rules_stream_id_gen.get_current_token(), ) get_push_rules_for_user = PushRuleStore.__dict__["get_push_rules_for_user"] get_push_rules_enabled_for_user = ( PushRuleStore.__dict__["get_push_rules_enabled_for_user"]) have_push_rules_changed_for_user = ( DataStore.have_push_rules_changed_for_user.__func__) def get_push_rules_stream_token(self): return ( self._push_rules_stream_id_gen.get_current_token(), self._stream_id_gen.get_current_token(), ) def stream_positions(self): result = super(SlavedPushRuleStore, self).stream_positions() result[ "push_rules"] = self._push_rules_stream_id_gen.get_current_token() return result def process_replication(self, result): stream = result.get("push_rules") if stream: for row in stream["rows"]: position = row[0] user_id = row[2] self.get_push_rules_for_user.invalidate((user_id, )) self.get_push_rules_enabled_for_user.invalidate((user_id, )) self.push_rules_stream_cache.entity_has_changed( user_id, position) self._push_rules_stream_id_gen.advance(int(stream["position"])) return super(SlavedPushRuleStore, self).process_replication(result)
class UserDirectorySlaveStore( SlavedEventStore, SlavedApplicationServiceStore, SlavedRegistrationStore, SlavedClientIpStore, UserDirectoryStore, BaseSlavedStore, ): def __init__(self, db_conn, hs): super(UserDirectorySlaveStore, self).__init__(db_conn, hs) events_max = self._stream_id_gen.get_current_token() curr_state_delta_prefill, min_curr_state_delta_id = self._get_cache_dict( db_conn, "current_state_delta_stream", entity_column="room_id", stream_column="stream_id", max_value=events_max, # As we share the stream id with events token limit=1000, ) self._curr_state_delta_stream_cache = StreamChangeCache( "_curr_state_delta_stream_cache", min_curr_state_delta_id, prefilled_cache=curr_state_delta_prefill, ) def stream_positions(self): result = super(UserDirectorySlaveStore, self).stream_positions() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == EventsStream.NAME: self._stream_id_gen.advance(token) for row in rows: if row.type != EventsStreamCurrentStateRow.TypeId: continue self._curr_state_delta_stream_cache.entity_has_changed( row.data.room_id, token ) return super(UserDirectorySlaveStore, self).process_replication_rows( stream_name, token, rows )
class SlavedPushRuleStore(SlavedEventStore): def __init__(self, db_conn, hs): super(SlavedPushRuleStore, self).__init__(db_conn, hs) self._push_rules_stream_id_gen = SlavedIdTracker( db_conn, "push_rules_stream", "stream_id", ) self.push_rules_stream_cache = StreamChangeCache( "PushRulesStreamChangeCache", self._push_rules_stream_id_gen.get_current_token(), ) get_push_rules_for_user = PushRuleStore.__dict__["get_push_rules_for_user"] get_push_rules_enabled_for_user = ( PushRuleStore.__dict__["get_push_rules_enabled_for_user"]) have_push_rules_changed_for_user = ( DataStore.have_push_rules_changed_for_user.__func__) def get_push_rules_stream_token(self): return ( self._push_rules_stream_id_gen.get_current_token(), self._stream_id_gen.get_current_token(), ) def stream_positions(self): result = super(SlavedPushRuleStore, self).stream_positions() result[ "push_rules"] = self._push_rules_stream_id_gen.get_current_token() return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "push_rules": self._push_rules_stream_id_gen.advance(token) for row in rows: self.get_push_rules_for_user.invalidate((row.user_id, )) self.get_push_rules_enabled_for_user.invalidate( (row.user_id, )) self.push_rules_stream_cache.entity_has_changed( row.user_id, token) return super(SlavedPushRuleStore, self).process_replication_rows(stream_name, token, rows)
class UserDirectorySlaveStore( SlavedEventStore, SlavedApplicationServiceStore, SlavedRegistrationStore, SlavedClientIpStore, UserDirectoryStore, BaseSlavedStore, ): def __init__(self, db_conn, hs): super(UserDirectorySlaveStore, self).__init__(db_conn, hs) events_max = self._stream_id_gen.get_current_token() curr_state_delta_prefill, min_curr_state_delta_id = self._get_cache_dict( db_conn, "current_state_delta_stream", entity_column="room_id", stream_column="stream_id", max_value=events_max, # As we share the stream id with events token limit=1000, ) self._curr_state_delta_stream_cache = StreamChangeCache( "_curr_state_delta_stream_cache", min_curr_state_delta_id, prefilled_cache=curr_state_delta_prefill, ) self._current_state_delta_pos = events_max def stream_positions(self): result = super(UserDirectorySlaveStore, self).stream_positions() result["current_state_deltas"] = self._current_state_delta_pos return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "current_state_deltas": self._current_state_delta_pos = token for row in rows: self._curr_state_delta_stream_cache.entity_has_changed( row.room_id, token ) return super(UserDirectorySlaveStore, self).process_replication_rows( stream_name, token, rows )
def test_get_entities_changed(self): """ StreamChangeCache.get_entities_changed will return the entities in the given list that have changed since the provided stream ID. If the stream position is earlier than the earliest known position, it will return all of the entities queried for. """ cache = StreamChangeCache("#test", 1) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) # Query all the entries, but mid-way through the stream. We should only # get the ones after that point. self.assertEqual( cache.get_entities_changed( ["*****@*****.**", "*****@*****.**", "*****@*****.**"], stream_pos=2 ), set(["*****@*****.**", "*****@*****.**"]), ) # Query all the entries mid-way through the stream, but include one # that doesn't exist in it. We should get back the one that doesn't # exist, too. self.assertEqual( cache.get_entities_changed( [ "*****@*****.**", "*****@*****.**", "*****@*****.**", "*****@*****.**", ], stream_pos=2, ), set(["*****@*****.**", "*****@*****.**", "*****@*****.**"]), ) # Query all the entries, but before the first known point. We will get # all the entries we queried for, including ones that don't exist. self.assertEqual( cache.get_entities_changed( [ "*****@*****.**", "*****@*****.**", "*****@*****.**", "*****@*****.**", ], stream_pos=0, ), set( [ "*****@*****.**", "*****@*****.**", "*****@*****.**", "*****@*****.**", ] ), )
class SlavedPresenceStore(BaseSlavedStore): def __init__(self, database: Database, db_conn, hs): super(SlavedPresenceStore, self).__init__(database, db_conn, hs) self._presence_id_gen = SlavedIdTracker(db_conn, "presence_stream", "stream_id") self._presence_on_startup = self._get_active_presence(db_conn) self.presence_stream_cache = StreamChangeCache( "PresenceStreamChangeCache", self._presence_id_gen.get_current_token()) _get_active_presence = __func__(DataStore._get_active_presence) take_presence_startup_info = __func__(DataStore.take_presence_startup_info) _get_presence_for_user = PresenceStore.__dict__["_get_presence_for_user"] get_presence_for_users = PresenceStore.__dict__["get_presence_for_users"] def get_current_presence_token(self): return self._presence_id_gen.get_current_token() def stream_positions(self): result = super(SlavedPresenceStore, self).stream_positions() if self.hs.config.use_presence: position = self._presence_id_gen.get_current_token() result["presence"] = position return result def process_replication_rows(self, stream_name, token, rows): if stream_name == "presence": self._presence_id_gen.advance(token) for row in rows: self.presence_stream_cache.entity_has_changed( row.user_id, token) self._get_presence_for_user.invalidate((row.user_id, )) return super(SlavedPresenceStore, self).process_replication_rows(stream_name, token, rows)
class SlavedDeviceStore(EndToEndKeyWorkerStore, DeviceWorkerStore, BaseSlavedStore): def __init__(self, database: DatabasePool, db_conn, hs): super().__init__(database, db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker( db_conn, "device_lists_stream", "stream_id", extra_tables=[ ("user_signature_stream", "stream_id"), ("device_lists_outbound_pokes", "stream_id"), ], ) device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max) self._user_signature_stream_cache = StreamChangeCache( "UserSignatureStreamChangeCache", device_list_max) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max) def get_device_stream_token(self) -> int: return self._device_list_id_gen.get_current_token() def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == DeviceListsStream.NAME: self._device_list_id_gen.advance(instance_name, token) self._invalidate_caches_for_devices(token, rows) elif stream_name == UserSignatureStream.NAME: self._device_list_id_gen.advance(instance_name, token) for row in rows: self._user_signature_stream_cache.entity_has_changed( row.user_id, token) return super().process_replication_rows(stream_name, instance_name, token, rows) def _invalidate_caches_for_devices(self, token, rows): for row in rows: # The entities are either user IDs (starting with '@') whose devices # have changed, or remote servers that we need to tell about # changes. if row.entity.startswith("@"): self._device_list_stream_cache.entity_has_changed( row.entity, token) self.get_cached_devices_for_user.invalidate((row.entity, )) self._get_cached_user_device.invalidate_many((row.entity, )) self.get_device_list_last_stream_id_for_remote.invalidate( (row.entity, )) else: self._device_list_federation_stream_cache.entity_has_changed( row.entity, token)
class SlavedDeviceStore(EndToEndKeyWorkerStore, DeviceWorkerStore, BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedDeviceStore, self).__init__(db_conn, hs) self.hs = hs self._device_list_id_gen = SlavedIdTracker(db_conn, "device_lists_stream", "stream_id") device_list_max = self._device_list_id_gen.get_current_token() self._device_list_stream_cache = StreamChangeCache( "DeviceListStreamChangeCache", device_list_max) self._user_signature_stream_cache = StreamChangeCache( "UserSignatureStreamChangeCache", device_list_max) self._device_list_federation_stream_cache = StreamChangeCache( "DeviceListFederationStreamChangeCache", device_list_max) def stream_positions(self): result = super(SlavedDeviceStore, self).stream_positions() # The user signature stream uses the same stream ID generator as the # device list stream, so set them both to the device list ID # generator's current token. current_token = self._device_list_id_gen.get_current_token() result[DeviceListsStream.NAME] = current_token result[UserSignatureStream.NAME] = current_token return result def process_replication_rows(self, stream_name, token, rows): if stream_name == DeviceListsStream.NAME: self._device_list_id_gen.advance(token) for row in rows: self._invalidate_caches_for_devices(token, row.user_id, row.destination) elif stream_name == UserSignatureStream.NAME: for row in rows: self._user_signature_stream_cache.entity_has_changed( row.user_id, token) return super(SlavedDeviceStore, self).process_replication_rows(stream_name, token, rows) def _invalidate_caches_for_devices(self, token, user_id, destination): self._device_list_stream_cache.entity_has_changed(user_id, token) if destination: self._device_list_federation_stream_cache.entity_has_changed( destination, token) self._get_cached_devices_for_user.invalidate((user_id, )) self._get_cached_user_device.invalidate_many((user_id, )) self.get_device_list_last_stream_id_for_remote.invalidate((user_id, ))
def test_has_entity_changed_pops_off_start(self): """ StreamChangeCache.entity_has_changed will respect the max size and purge the oldest items upon reaching that max size. """ cache = StreamChangeCache("#test", 1, max_size=2) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) # The cache is at the max size, 2 self.assertEqual(len(cache._cache), 2) # The oldest item has been popped off self.assertTrue("*****@*****.**" not in cache._entity_to_key) # If we update an existing entity, it keeps the two existing entities cache.entity_has_changed("*****@*****.**", 5) self.assertEqual( set(["*****@*****.**", "*****@*****.**"]), set(cache._entity_to_key) )
def test_max_pos(self): """ StreamChangeCache.get_max_pos_of_last_change will return the most recent point where the entity could have changed. If the entity is not known, the stream start is provided instead. """ cache = StreamChangeCache("#test", 1) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) # Known entities will return the point where they were changed. self.assertEqual(cache.get_max_pos_of_last_change("*****@*****.**"), 2) self.assertEqual(cache.get_max_pos_of_last_change("*****@*****.**"), 3) self.assertEqual(cache.get_max_pos_of_last_change("*****@*****.**"), 4) # Unknown entities will return the stream start position. self.assertEqual(cache.get_max_pos_of_last_change("*****@*****.**"), 1)
def test_get_all_entities_changed(self): """ StreamChangeCache.get_all_entities_changed will return all changed entities since the given position. If the position is before the start of the known stream, it returns None instead. """ cache = StreamChangeCache("#test", 1) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) self.assertEqual( cache.get_all_entities_changed(1), ["*****@*****.**", "*****@*****.**", "*****@*****.**"], ) self.assertEqual( cache.get_all_entities_changed(2), ["*****@*****.**", "*****@*****.**"] ) self.assertEqual(cache.get_all_entities_changed(3), ["*****@*****.**"]) self.assertEqual(cache.get_all_entities_changed(0), None)
def test_get_all_entities_changed(self): """ StreamChangeCache.get_all_entities_changed will return all changed entities since the given position. If the position is before the start of the known stream, it returns None instead. """ cache = StreamChangeCache("#test", 1) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) self.assertEqual( cache.get_all_entities_changed(1), ["*****@*****.**", "*****@*****.**", "*****@*****.**"], ) self.assertEqual(cache.get_all_entities_changed(2), ["*****@*****.**", "*****@*****.**"]) self.assertEqual(cache.get_all_entities_changed(3), ["*****@*****.**"]) self.assertEqual(cache.get_all_entities_changed(0), None)
class AccountDataWorkerStore(SQLBaseStore): """This is an abstract base class where subclasses must implement `get_max_account_data_stream_id` which can be called in the initializer. """ def __init__(self, database: DatabasePool, db_conn, hs): self._instance_name = hs.get_instance_name() if isinstance(database.engine, PostgresEngine): self._can_write_to_account_data = ( self._instance_name in hs.config.worker.writers.account_data) self._account_data_id_gen = MultiWriterIdGenerator( db_conn=db_conn, db=database, stream_name="account_data", instance_name=self._instance_name, tables=[ ("room_account_data", "instance_name", "stream_id"), ("room_tags_revisions", "instance_name", "stream_id"), ("account_data", "instance_name", "stream_id"), ], sequence_name="account_data_sequence", writers=hs.config.worker.writers.account_data, ) else: self._can_write_to_account_data = True # We shouldn't be running in worker mode with SQLite, but its useful # to support it for unit tests. # # If this process is the writer than we need to use # `StreamIdGenerator`, otherwise we use `SlavedIdTracker` which gets # updated over replication. (Multiple writers are not supported for # SQLite). if hs.get_instance_name() in hs.config.worker.writers.account_data: self._account_data_id_gen = StreamIdGenerator( db_conn, "room_account_data", "stream_id", extra_tables=[("room_tags_revisions", "stream_id")], ) else: self._account_data_id_gen = SlavedIdTracker( db_conn, "room_account_data", "stream_id", extra_tables=[("room_tags_revisions", "stream_id")], ) account_max = self.get_max_account_data_stream_id() self._account_data_stream_cache = StreamChangeCache( "AccountDataAndTagsChangeCache", account_max) super().__init__(database, db_conn, hs) def get_max_account_data_stream_id(self) -> int: """Get the current max stream ID for account data stream Returns: int """ return self._account_data_id_gen.get_current_token() @cached() async def get_account_data_for_user( self, user_id: str ) -> Tuple[Dict[str, JsonDict], Dict[str, Dict[str, JsonDict]]]: """Get all the client account_data for a user. Args: user_id: The user to get the account_data for. Returns: A 2-tuple of a dict of global account_data and a dict mapping from room_id string to per room account_data dicts. """ def get_account_data_for_user_txn(txn): rows = self.db_pool.simple_select_list_txn( txn, "account_data", {"user_id": user_id}, ["account_data_type", "content"], ) global_account_data = { row["account_data_type"]: db_to_json(row["content"]) for row in rows } rows = self.db_pool.simple_select_list_txn( txn, "room_account_data", {"user_id": user_id}, ["room_id", "account_data_type", "content"], ) by_room = {} for row in rows: room_data = by_room.setdefault(row["room_id"], {}) room_data[row["account_data_type"]] = db_to_json( row["content"]) return global_account_data, by_room return await self.db_pool.runInteraction( "get_account_data_for_user", get_account_data_for_user_txn) @cached(num_args=2, max_entries=5000) async def get_global_account_data_by_type_for_user( self, data_type: str, user_id: str) -> Optional[JsonDict]: """ Returns: The account data. """ result = await self.db_pool.simple_select_one_onecol( table="account_data", keyvalues={ "user_id": user_id, "account_data_type": data_type }, retcol="content", desc="get_global_account_data_by_type_for_user", allow_none=True, ) if result: return db_to_json(result) else: return None @cached(num_args=2) async def get_account_data_for_room(self, user_id: str, room_id: str) -> Dict[str, JsonDict]: """Get all the client account_data for a user for a room. Args: user_id: The user to get the account_data for. room_id: The room to get the account_data for. Returns: A dict of the room account_data """ def get_account_data_for_room_txn(txn): rows = self.db_pool.simple_select_list_txn( txn, "room_account_data", { "user_id": user_id, "room_id": room_id }, ["account_data_type", "content"], ) return { row["account_data_type"]: db_to_json(row["content"]) for row in rows } return await self.db_pool.runInteraction( "get_account_data_for_room", get_account_data_for_room_txn) @cached(num_args=3, max_entries=5000) async def get_account_data_for_room_and_type( self, user_id: str, room_id: str, account_data_type: str) -> Optional[JsonDict]: """Get the client account_data of given type for a user for a room. Args: user_id: The user to get the account_data for. room_id: The room to get the account_data for. account_data_type: The account data type to get. Returns: The room account_data for that type, or None if there isn't any set. """ def get_account_data_for_room_and_type_txn(txn): content_json = self.db_pool.simple_select_one_onecol_txn( txn, table="room_account_data", keyvalues={ "user_id": user_id, "room_id": room_id, "account_data_type": account_data_type, }, retcol="content", allow_none=True, ) return db_to_json(content_json) if content_json else None return await self.db_pool.runInteraction( "get_account_data_for_room_and_type", get_account_data_for_room_and_type_txn) async def get_updated_global_account_data( self, last_id: int, current_id: int, limit: int) -> List[Tuple[int, str, str]]: """Get the global account_data that has changed, for the account_data stream Args: last_id: the last stream_id from the previous batch. current_id: the maximum stream_id to return up to limit: the maximum number of rows to return Returns: A list of tuples of stream_id int, user_id string, and type string. """ if last_id == current_id: return [] def get_updated_global_account_data_txn(txn): sql = ("SELECT stream_id, user_id, account_data_type" " FROM account_data WHERE ? < stream_id AND stream_id <= ?" " ORDER BY stream_id ASC LIMIT ?") txn.execute(sql, (last_id, current_id, limit)) return txn.fetchall() return await self.db_pool.runInteraction( "get_updated_global_account_data", get_updated_global_account_data_txn) async def get_updated_room_account_data( self, last_id: int, current_id: int, limit: int) -> List[Tuple[int, str, str, str]]: """Get the global account_data that has changed, for the account_data stream Args: last_id: the last stream_id from the previous batch. current_id: the maximum stream_id to return up to limit: the maximum number of rows to return Returns: A list of tuples of stream_id int, user_id string, room_id string and type string. """ if last_id == current_id: return [] def get_updated_room_account_data_txn(txn): sql = ( "SELECT stream_id, user_id, room_id, account_data_type" " FROM room_account_data WHERE ? < stream_id AND stream_id <= ?" " ORDER BY stream_id ASC LIMIT ?") txn.execute(sql, (last_id, current_id, limit)) return txn.fetchall() return await self.db_pool.runInteraction( "get_updated_room_account_data", get_updated_room_account_data_txn) async def get_updated_account_data_for_user( self, user_id: str, stream_id: int ) -> Tuple[Dict[str, JsonDict], Dict[str, Dict[str, JsonDict]]]: """Get all the client account_data for a that's changed for a user Args: user_id: The user to get the account_data for. stream_id: The point in the stream since which to get updates Returns: A deferred pair of a dict of global account_data and a dict mapping from room_id string to per room account_data dicts. """ def get_updated_account_data_for_user_txn(txn): sql = ("SELECT account_data_type, content FROM account_data" " WHERE user_id = ? AND stream_id > ?") txn.execute(sql, (user_id, stream_id)) global_account_data = {row[0]: db_to_json(row[1]) for row in txn} sql = ( "SELECT room_id, account_data_type, content FROM room_account_data" " WHERE user_id = ? AND stream_id > ?") txn.execute(sql, (user_id, stream_id)) account_data_by_room = {} for row in txn: room_account_data = account_data_by_room.setdefault(row[0], {}) room_account_data[row[1]] = db_to_json(row[2]) return global_account_data, account_data_by_room changed = self._account_data_stream_cache.has_entity_changed( user_id, int(stream_id)) if not changed: return ({}, {}) return await self.db_pool.runInteraction( "get_updated_account_data_for_user", get_updated_account_data_for_user_txn) @cached(max_entries=5000, iterable=True) async def ignored_by(self, user_id: str) -> Set[str]: """ Get users which ignore the given user. Params: user_id: The user ID which might be ignored. Return: The user IDs which ignore the given user. """ return set(await self.db_pool.simple_select_onecol( table="ignored_users", keyvalues={"ignored_user_id": user_id}, retcol="ignorer_user_id", desc="ignored_by", )) def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == TagAccountDataStream.NAME: self._account_data_id_gen.advance(instance_name, token) for row in rows: self.get_tags_for_user.invalidate((row.user_id, )) self._account_data_stream_cache.entity_has_changed( row.user_id, token) elif stream_name == AccountDataStream.NAME: self._account_data_id_gen.advance(instance_name, token) for row in rows: if not row.room_id: self.get_global_account_data_by_type_for_user.invalidate( (row.data_type, row.user_id)) self.get_account_data_for_user.invalidate((row.user_id, )) self.get_account_data_for_room.invalidate( (row.user_id, row.room_id)) self.get_account_data_for_room_and_type.invalidate( (row.user_id, row.room_id, row.data_type)) self._account_data_stream_cache.entity_has_changed( row.user_id, token) return super().process_replication_rows(stream_name, instance_name, token, rows) async def add_account_data_to_room(self, user_id: str, room_id: str, account_data_type: str, content: JsonDict) -> int: """Add some account_data to a room for a user. Args: user_id: The user to add a tag for. 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. Returns: The maximum stream ID. """ assert self._can_write_to_account_data content_json = json_encoder.encode(content) async with self._account_data_id_gen.get_next() as next_id: # no need to lock here as room_account_data has a unique constraint # on (user_id, room_id, account_data_type) so simple_upsert will # retry if there is a conflict. await self.db_pool.simple_upsert( desc="add_room_account_data", 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 }, lock=False, ) self._account_data_stream_cache.entity_has_changed( user_id, next_id) self.get_account_data_for_user.invalidate((user_id, )) self.get_account_data_for_room.invalidate((user_id, room_id)) self.get_account_data_for_room_and_type.prefill( (user_id, room_id, account_data_type), content) 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) -> int: """Add some account_data to a room 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. Returns: The maximum stream ID. """ assert self._can_write_to_account_data async with self._account_data_id_gen.get_next() as next_id: await self.db_pool.runInteraction( "add_user_account_data", self._add_account_data_for_user, next_id, user_id, account_data_type, content, ) self._account_data_stream_cache.entity_has_changed( user_id, next_id) self.get_account_data_for_user.invalidate((user_id, )) self.get_global_account_data_by_type_for_user.invalidate( (account_data_type, user_id)) return self._account_data_id_gen.get_current_token() def _add_account_data_for_user( self, txn, next_id: int, user_id: str, account_data_type: str, content: JsonDict, ) -> None: content_json = json_encoder.encode(content) # no need to lock here as account_data has a unique constraint on # (user_id, account_data_type) so simple_upsert will retry if # there is a conflict. self.db_pool.simple_upsert_txn( txn, table="account_data", keyvalues={ "user_id": user_id, "account_data_type": account_data_type }, values={ "stream_id": next_id, "content": content_json }, lock=False, ) # Ignored users get denormalized into a separate table as an optimisation. if account_data_type != AccountDataTypes.IGNORED_USER_LIST: return # Insert / delete to sync the list of ignored users. previously_ignored_users = set( self.db_pool.simple_select_onecol_txn( txn, table="ignored_users", keyvalues={"ignorer_user_id": user_id}, retcol="ignored_user_id", )) # If the data is invalid, no one is ignored. ignored_users_content = content.get("ignored_users", {}) if isinstance(ignored_users_content, dict): currently_ignored_users = set(ignored_users_content) else: currently_ignored_users = set() # Delete entries which are no longer ignored. self.db_pool.simple_delete_many_txn( txn, table="ignored_users", column="ignored_user_id", iterable=previously_ignored_users - currently_ignored_users, keyvalues={"ignorer_user_id": user_id}, ) # Add entries which are newly ignored. self.db_pool.simple_insert_many_txn( txn, table="ignored_users", values=[{ "ignorer_user_id": user_id, "ignored_user_id": u } for u in currently_ignored_users - previously_ignored_users], ) # Invalidate the cache for any ignored users which were added or removed. for ignored_user_id in previously_ignored_users ^ currently_ignored_users: self._invalidate_cache_and_stream(txn, self.ignored_by, (ignored_user_id, ))
def test_get_entities_changed(self): """ StreamChangeCache.get_entities_changed will return the entities in the given list that have changed since the provided stream ID. If the stream position is earlier than the earliest known position, it will return all of the entities queried for. """ cache = StreamChangeCache("#test", 1) cache.entity_has_changed("*****@*****.**", 2) cache.entity_has_changed("*****@*****.**", 3) cache.entity_has_changed("*****@*****.**", 4) # Query all the entries, but mid-way through the stream. We should only # get the ones after that point. self.assertEqual( cache.get_entities_changed( ["*****@*****.**", "*****@*****.**", "*****@*****.**"], stream_pos=2 ), set(["*****@*****.**", "*****@*****.**"]), ) # Query all the entries mid-way through the stream, but include one # that doesn't exist in it. We shouldn't get back the one that doesn't # exist. self.assertEqual( cache.get_entities_changed( [ "*****@*****.**", "*****@*****.**", "*****@*****.**", "*****@*****.**", ], stream_pos=2, ), set(["*****@*****.**", "*****@*****.**"]), ) # Query all the entries, but before the first known point. We will get # all the entries we queried for, including ones that don't exist. self.assertEqual( cache.get_entities_changed( [ "*****@*****.**", "*****@*****.**", "*****@*****.**", "*****@*****.**", ], stream_pos=0, ), set( [ "*****@*****.**", "*****@*****.**", "*****@*****.**", "*****@*****.**", ] ), ) # Query a subset of the entries mid-way through the stream. We should # only get back the subset. self.assertEqual( cache.get_entities_changed(["*****@*****.**"], stream_pos=2), set(["*****@*****.**"]), )
class TypingWriterHandler(FollowerTypingHandler): def __init__(self, hs: "HomeServer"): super().__init__(hs) assert hs.config.worker.writers.typing == hs.get_instance_name() self.auth = hs.get_auth() self.notifier = hs.get_notifier() self.hs = hs hs.get_federation_registry().register_edu_handler( "m.typing", self._recv_edu) hs.get_distributor().observe("user_left_room", self.user_left_room) # clock time we expect to stop self._member_typing_until = {} # type: Dict[RoomMember, int] # caches which room_ids changed at which serials self._typing_stream_change_cache = StreamChangeCache( "TypingStreamChangeCache", self._latest_room_serial) def _handle_timeout_for_member(self, now: int, member: RoomMember) -> None: super()._handle_timeout_for_member(now, member) if not self.is_typing(member): # Nothing to do if they're no longer typing return until = self._member_typing_until.get(member, None) if not until or until <= now: logger.info("Timing out typing for: %s", member.user_id) self._stopped_typing(member) return async def started_typing(self, target_user: UserID, requester: Requester, room_id: str, timeout: int) -> None: target_user_id = target_user.to_string() auth_user_id = requester.user.to_string() if not self.is_mine_id(target_user_id): raise SynapseError(400, "User is not hosted on this homeserver") if target_user_id != auth_user_id: raise AuthError(400, "Cannot set another user's typing state") if requester.shadow_banned: # We randomly sleep a bit just to annoy the requester. await self.clock.sleep(random.randint(1, 10)) raise ShadowBanError() await self.auth.check_user_in_room(room_id, target_user_id) logger.debug("%s has started typing in %s", target_user_id, room_id) member = RoomMember(room_id=room_id, user_id=target_user_id) was_present = member.user_id in self._room_typing.get(room_id, set()) now = self.clock.time_msec() self._member_typing_until[member] = now + timeout self.wheel_timer.insert(now=now, obj=member, then=now + timeout) if was_present: # No point sending another notification return self._push_update(member=member, typing=True) async def stopped_typing(self, target_user: UserID, requester: Requester, room_id: str) -> None: target_user_id = target_user.to_string() auth_user_id = requester.user.to_string() if not self.is_mine_id(target_user_id): raise SynapseError(400, "User is not hosted on this homeserver") if target_user_id != auth_user_id: raise AuthError(400, "Cannot set another user's typing state") if requester.shadow_banned: # We randomly sleep a bit just to annoy the requester. await self.clock.sleep(random.randint(1, 10)) raise ShadowBanError() await self.auth.check_user_in_room(room_id, target_user_id) logger.debug("%s has stopped typing in %s", target_user_id, room_id) member = RoomMember(room_id=room_id, user_id=target_user_id) self._stopped_typing(member) def user_left_room(self, user: UserID, room_id: str) -> None: user_id = user.to_string() if self.is_mine_id(user_id): member = RoomMember(room_id=room_id, user_id=user_id) self._stopped_typing(member) def _stopped_typing(self, member: RoomMember) -> None: if member.user_id not in self._room_typing.get(member.room_id, set()): # No point return self._member_typing_until.pop(member, None) self._member_last_federation_poke.pop(member, None) self._push_update(member=member, typing=False) def _push_update(self, member: RoomMember, typing: bool) -> None: if self.hs.is_mine_id(member.user_id): # Only send updates for changes to our own users. run_as_background_process("typing._push_remote", self._push_remote, member, typing) self._push_update_local(member=member, typing=typing) async def _recv_edu(self, origin: str, content: JsonDict) -> None: room_id = content["room_id"] user_id = content["user_id"] member = RoomMember(user_id=user_id, room_id=room_id) # Check that the string is a valid user id user = UserID.from_string(user_id) if user.domain != origin: logger.info("Got typing update from %r with bad 'user_id': %r", origin, user_id) return users = await self.store.get_users_in_room(room_id) domains = {get_domain_from_id(u) for u in users} if self.server_name in domains: logger.info("Got typing update from %s: %r", user_id, content) now = self.clock.time_msec() self._member_typing_until[member] = now + FEDERATION_TIMEOUT self.wheel_timer.insert(now=now, obj=member, then=now + FEDERATION_TIMEOUT) self._push_update_local(member=member, typing=content["typing"]) def _push_update_local(self, member: RoomMember, typing: bool) -> None: room_set = self._room_typing.setdefault(member.room_id, set()) if typing: room_set.add(member.user_id) else: room_set.discard(member.user_id) self._latest_room_serial += 1 self._room_serials[member.room_id] = self._latest_room_serial self._typing_stream_change_cache.entity_has_changed( member.room_id, self._latest_room_serial) self.notifier.on_new_event("typing_key", self._latest_room_serial, rooms=[member.room_id]) async def get_all_typing_updates( self, instance_name: str, last_id: int, current_id: int, limit: int) -> Tuple[List[Tuple[int, list]], int, bool]: """Get updates for typing replication stream. Args: instance_name: The writer we want to fetch updates from. Unused here since there is only ever one writer. last_id: The token to fetch updates from. Exclusive. current_id: The token to fetch updates up to. Inclusive. limit: The requested limit for the number of rows to return. The function may return more or fewer rows. Returns: A tuple consisting of: the updates, a token to use to fetch subsequent updates, and whether we returned fewer rows than exists between the requested tokens due to the limit. The token returned can be used in a subsequent call to this function to get further updates. The updates are a list of 2-tuples of stream ID and the row data """ if last_id == current_id: return [], current_id, False changed_rooms = self._typing_stream_change_cache.get_all_entities_changed( last_id) # type: Optional[Iterable[str]] if changed_rooms is None: changed_rooms = self._room_serials rows = [] for room_id in changed_rooms: serial = self._room_serials[room_id] if last_id < serial <= current_id: typing = self._room_typing[room_id] rows.append((serial, [room_id, list(typing)])) rows.sort() limited = False # We, unusually, use a strict limit here as we have all the rows in # memory rather than pulling them out of the database with a `LIMIT ?` # clause. if len(rows) > limit: rows = rows[:limit] current_id = rows[-1][0] limited = True return rows, current_id, limited def process_replication_rows( self, token: int, rows: List[TypingStream.TypingStreamRow]) -> None: # The writing process should never get updates from replication. raise Exception( "Typing writer instance got typing info over replication")
class SlavedEventStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedEventStore, self).__init__(db_conn, hs) self._stream_id_gen = SlavedIdTracker( db_conn, "events", "stream_ordering", ) self._backfill_id_gen = SlavedIdTracker(db_conn, "events", "stream_ordering", step=-1) events_max = self._stream_id_gen.get_current_token() event_cache_prefill, min_event_val = self._get_cache_dict( db_conn, "events", entity_column="room_id", stream_column="stream_ordering", max_value=events_max, ) self._events_stream_cache = StreamChangeCache( "EventsRoomStreamChangeCache", min_event_val, prefilled_cache=event_cache_prefill, ) self._membership_stream_cache = StreamChangeCache( "MembershipStreamChangeCache", events_max, ) self.stream_ordering_month_ago = 0 self._stream_order_on_start = self.get_room_max_stream_ordering() # Cached functions can't be accessed through a class instance so we need # to reach inside the __dict__ to extract them. get_rooms_for_user = RoomMemberStore.__dict__["get_rooms_for_user"] get_users_in_room = RoomMemberStore.__dict__["get_users_in_room"] get_latest_event_ids_in_room = EventFederationStore.__dict__[ "get_latest_event_ids_in_room"] _get_current_state_for_key = StateStore.__dict__[ "_get_current_state_for_key"] get_invited_rooms_for_user = RoomMemberStore.__dict__[ "get_invited_rooms_for_user"] get_unread_event_push_actions_by_room_for_user = ( EventPushActionsStore. __dict__["get_unread_event_push_actions_by_room_for_user"]) _get_state_group_for_events = ( StateStore.__dict__["_get_state_group_for_events"]) _get_state_group_for_event = ( StateStore.__dict__["_get_state_group_for_event"]) _get_state_groups_from_groups = ( StateStore.__dict__["_get_state_groups_from_groups"]) _get_state_groups_from_groups_txn = ( DataStore._get_state_groups_from_groups_txn.__func__) _get_state_group_from_group = ( StateStore.__dict__["_get_state_group_from_group"]) get_recent_event_ids_for_room = ( StreamStore.__dict__["get_recent_event_ids_for_room"]) get_unread_push_actions_for_user_in_range_for_http = ( DataStore.get_unread_push_actions_for_user_in_range_for_http.__func__) get_unread_push_actions_for_user_in_range_for_email = ( DataStore.get_unread_push_actions_for_user_in_range_for_email.__func__) get_push_action_users_in_range = ( DataStore.get_push_action_users_in_range.__func__) get_event = DataStore.get_event.__func__ get_events = DataStore.get_events.__func__ get_current_state = DataStore.get_current_state.__func__ get_current_state_for_key = DataStore.get_current_state_for_key.__func__ get_rooms_for_user_where_membership_is = ( DataStore.get_rooms_for_user_where_membership_is.__func__) get_membership_changes_for_user = ( DataStore.get_membership_changes_for_user.__func__) get_room_events_max_id = DataStore.get_room_events_max_id.__func__ get_room_events_stream_for_room = ( DataStore.get_room_events_stream_for_room.__func__) get_events_around = DataStore.get_events_around.__func__ get_state_for_event = DataStore.get_state_for_event.__func__ get_state_for_events = DataStore.get_state_for_events.__func__ get_state_groups = DataStore.get_state_groups.__func__ get_state_groups_ids = DataStore.get_state_groups_ids.__func__ get_state_ids_for_event = DataStore.get_state_ids_for_event.__func__ get_state_ids_for_events = DataStore.get_state_ids_for_events.__func__ get_joined_users_from_state = DataStore.get_joined_users_from_state.__func__ get_joined_users_from_context = DataStore.get_joined_users_from_context.__func__ _get_joined_users_from_context = ( RoomMemberStore.__dict__["_get_joined_users_from_context"]) get_recent_events_for_room = DataStore.get_recent_events_for_room.__func__ get_room_events_stream_for_rooms = ( DataStore.get_room_events_stream_for_rooms.__func__) is_host_joined = DataStore.is_host_joined.__func__ _is_host_joined = RoomMemberStore.__dict__["_is_host_joined"] get_stream_token_for_event = DataStore.get_stream_token_for_event.__func__ _set_before_and_after = staticmethod(DataStore._set_before_and_after) _get_events = DataStore._get_events.__func__ _get_events_from_cache = DataStore._get_events_from_cache.__func__ _invalidate_get_event_cache = DataStore._invalidate_get_event_cache.__func__ _enqueue_events = DataStore._enqueue_events.__func__ _do_fetch = DataStore._do_fetch.__func__ _fetch_event_rows = DataStore._fetch_event_rows.__func__ _get_event_from_row = DataStore._get_event_from_row.__func__ _get_rooms_for_user_where_membership_is_txn = ( DataStore._get_rooms_for_user_where_membership_is_txn.__func__) _get_members_rows_txn = DataStore._get_members_rows_txn.__func__ _get_state_for_groups = DataStore._get_state_for_groups.__func__ _get_all_state_from_cache = DataStore._get_all_state_from_cache.__func__ _get_events_around_txn = DataStore._get_events_around_txn.__func__ _get_some_state_from_cache = DataStore._get_some_state_from_cache.__func__ get_backfill_events = DataStore.get_backfill_events.__func__ _get_backfill_events = DataStore._get_backfill_events.__func__ get_missing_events = DataStore.get_missing_events.__func__ _get_missing_events = DataStore._get_missing_events.__func__ get_auth_chain = DataStore.get_auth_chain.__func__ get_auth_chain_ids = DataStore.get_auth_chain_ids.__func__ _get_auth_chain_ids_txn = DataStore._get_auth_chain_ids_txn.__func__ get_room_max_stream_ordering = DataStore.get_room_max_stream_ordering.__func__ get_forward_extremeties_for_room = ( DataStore.get_forward_extremeties_for_room.__func__) _get_forward_extremeties_for_room = ( EventFederationStore.__dict__["_get_forward_extremeties_for_room"]) def stream_positions(self): result = super(SlavedEventStore, self).stream_positions() result["events"] = self._stream_id_gen.get_current_token() result["backfill"] = -self._backfill_id_gen.get_current_token() return result def process_replication(self, result): state_resets = set( r[0] for r in result.get("state_resets", {"rows": []})["rows"]) stream = result.get("events") if stream: self._stream_id_gen.advance(int(stream["position"])) for row in stream["rows"]: self._process_replication_row(row, backfilled=False, state_resets=state_resets) stream = result.get("backfill") if stream: self._backfill_id_gen.advance(-int(stream["position"])) for row in stream["rows"]: self._process_replication_row(row, backfilled=True, state_resets=state_resets) stream = result.get("forward_ex_outliers") if stream: self._stream_id_gen.advance(int(stream["position"])) for row in stream["rows"]: event_id = row[1] self._invalidate_get_event_cache(event_id) stream = result.get("backward_ex_outliers") if stream: self._backfill_id_gen.advance(-int(stream["position"])) for row in stream["rows"]: event_id = row[1] self._invalidate_get_event_cache(event_id) return super(SlavedEventStore, self).process_replication(result) def _process_replication_row(self, row, backfilled, state_resets): position = row[0] internal = json.loads(row[1]) event_json = json.loads(row[2]) event = FrozenEvent(event_json, internal_metadata_dict=internal) self.invalidate_caches_for_event(event, backfilled, reset_state=position in state_resets) def invalidate_caches_for_event(self, event, backfilled, reset_state): if reset_state: self._get_current_state_for_key.invalidate_all() self.get_rooms_for_user.invalidate_all() self.get_users_in_room.invalidate((event.room_id, )) self._invalidate_get_event_cache(event.event_id) self.get_latest_event_ids_in_room.invalidate((event.room_id, )) self.get_unread_event_push_actions_by_room_for_user.invalidate_many( (event.room_id, )) if not backfilled: self._events_stream_cache.entity_has_changed( event.room_id, event.internal_metadata.stream_ordering) # self.get_unread_event_push_actions_by_room_for_user.invalidate_many( # (event.room_id,) # ) if event.type == EventTypes.Redaction: self._invalidate_get_event_cache(event.redacts) if event.type == EventTypes.Member: self.get_rooms_for_user.invalidate((event.state_key, )) self.get_users_in_room.invalidate((event.room_id, )) self._membership_stream_cache.entity_has_changed( event.state_key, event.internal_metadata.stream_ordering) self.get_invited_rooms_for_user.invalidate((event.state_key, )) if not event.is_state(): return if backfilled: return if (not event.internal_metadata.is_invite_from_remote() and event.internal_metadata.is_outlier()): return self._get_current_state_for_key.invalidate( (event.room_id, event.type, event.state_key))
class DeviceInboxWorkerStore(SQLBaseStore): def __init__(self, database: DatabasePool, db_conn, hs): super().__init__(database, db_conn, hs) self._instance_name = hs.get_instance_name() # Map of (user_id, device_id) to the last stream_id that has been # deleted up to. This is so that we can no op deletions. self._last_device_delete_cache = ExpiringCache( cache_name="last_device_delete_cache", clock=self._clock, max_len=10000, expiry_ms=30 * 60 * 1000, ) if isinstance(database.engine, PostgresEngine): self._can_write_to_device = (self._instance_name in hs.config.worker.writers.to_device) self._device_inbox_id_gen = MultiWriterIdGenerator( db_conn=db_conn, db=database, stream_name="to_device", instance_name=self._instance_name, tables=[("device_inbox", "instance_name", "stream_id")], sequence_name="device_inbox_sequence", writers=hs.config.worker.writers.to_device, ) else: self._can_write_to_device = True self._device_inbox_id_gen = StreamIdGenerator( db_conn, "device_inbox", "stream_id") max_device_inbox_id = self._device_inbox_id_gen.get_current_token() device_inbox_prefill, min_device_inbox_id = self.db_pool.get_cache_dict( db_conn, "device_inbox", entity_column="user_id", stream_column="stream_id", max_value=max_device_inbox_id, limit=1000, ) self._device_inbox_stream_cache = StreamChangeCache( "DeviceInboxStreamChangeCache", min_device_inbox_id, prefilled_cache=device_inbox_prefill, ) # The federation outbox and the local device inbox uses the same # stream_id generator. device_outbox_prefill, min_device_outbox_id = self.db_pool.get_cache_dict( db_conn, "device_federation_outbox", entity_column="destination", stream_column="stream_id", max_value=max_device_inbox_id, limit=1000, ) self._device_federation_outbox_stream_cache = StreamChangeCache( "DeviceFederationOutboxStreamChangeCache", min_device_outbox_id, prefilled_cache=device_outbox_prefill, ) def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == ToDeviceStream.NAME: self._device_inbox_id_gen.advance(instance_name, token) for row in rows: if row.entity.startswith("@"): self._device_inbox_stream_cache.entity_has_changed( row.entity, token) else: self._device_federation_outbox_stream_cache.entity_has_changed( row.entity, token) return super().process_replication_rows(stream_name, instance_name, token, rows) def get_to_device_stream_token(self): return self._device_inbox_id_gen.get_current_token() async def get_new_messages_for_device( self, user_id: str, device_id: str, last_stream_id: int, current_stream_id: int, limit: int = 100, ) -> Tuple[List[dict], int]: """ Args: user_id: The recipient user_id. device_id: The recipient device_id. last_stream_id: The last stream ID checked. current_stream_id: The current position of the to device message stream. limit: The maximum number of messages to retrieve. Returns: A list of messages for the device and where in the stream the messages got to. """ has_changed = self._device_inbox_stream_cache.has_entity_changed( user_id, last_stream_id) if not has_changed: return ([], current_stream_id) def get_new_messages_for_device_txn(txn): sql = ("SELECT stream_id, message_json FROM device_inbox" " WHERE user_id = ? AND device_id = ?" " AND ? < stream_id AND stream_id <= ?" " ORDER BY stream_id ASC" " LIMIT ?") txn.execute( sql, (user_id, device_id, last_stream_id, current_stream_id, limit)) messages = [] for row in txn: stream_pos = row[0] messages.append(db_to_json(row[1])) if len(messages) < limit: stream_pos = current_stream_id return messages, stream_pos return await self.db_pool.runInteraction( "get_new_messages_for_device", get_new_messages_for_device_txn) @trace async def delete_messages_for_device(self, user_id: str, device_id: str, up_to_stream_id: int) -> int: """ Args: user_id: The recipient user_id. device_id: The recipient device_id. up_to_stream_id: Where to delete messages up to. Returns: The number of messages deleted. """ # If we have cached the last stream id we've deleted up to, we can # check if there is likely to be anything that needs deleting last_deleted_stream_id = self._last_device_delete_cache.get( (user_id, device_id), None) set_tag("last_deleted_stream_id", last_deleted_stream_id) if last_deleted_stream_id: has_changed = self._device_inbox_stream_cache.has_entity_changed( user_id, last_deleted_stream_id) if not has_changed: log_kv({"message": "No changes in cache since last check"}) return 0 def delete_messages_for_device_txn(txn): sql = ("DELETE FROM device_inbox" " WHERE user_id = ? AND device_id = ?" " AND stream_id <= ?") txn.execute(sql, (user_id, device_id, up_to_stream_id)) return txn.rowcount count = await self.db_pool.runInteraction( "delete_messages_for_device", delete_messages_for_device_txn) log_kv({ "message": "deleted {} messages for device".format(count), "count": count }) # Update the cache, ensuring that we only ever increase the value last_deleted_stream_id = self._last_device_delete_cache.get( (user_id, device_id), 0) self._last_device_delete_cache[(user_id, device_id)] = max( last_deleted_stream_id, up_to_stream_id) return count @trace async def get_new_device_msgs_for_remote(self, destination, last_stream_id, current_stream_id, limit) -> Tuple[List[dict], int]: """ Args: destination(str): The name of the remote server. last_stream_id(int|long): The last position of the device message stream that the server sent up to. current_stream_id(int|long): The current position of the device message stream. Returns: A list of messages for the device and where in the stream the messages got to. """ set_tag("destination", destination) set_tag("last_stream_id", last_stream_id) set_tag("current_stream_id", current_stream_id) set_tag("limit", limit) has_changed = self._device_federation_outbox_stream_cache.has_entity_changed( destination, last_stream_id) if not has_changed or last_stream_id == current_stream_id: log_kv({"message": "No new messages in stream"}) return ([], current_stream_id) if limit <= 0: # This can happen if we run out of room for EDUs in the transaction. return ([], last_stream_id) @trace def get_new_messages_for_remote_destination_txn(txn): sql = ( "SELECT stream_id, messages_json FROM device_federation_outbox" " WHERE destination = ?" " AND ? < stream_id AND stream_id <= ?" " ORDER BY stream_id ASC" " LIMIT ?") txn.execute( sql, (destination, last_stream_id, current_stream_id, limit)) messages = [] for row in txn: stream_pos = row[0] messages.append(db_to_json(row[1])) if len(messages) < limit: log_kv({"message": "Set stream position to current position"}) stream_pos = current_stream_id return messages, stream_pos return await self.db_pool.runInteraction( "get_new_device_msgs_for_remote", get_new_messages_for_remote_destination_txn, ) @trace async def delete_device_msgs_for_remote(self, destination: str, up_to_stream_id: int) -> None: """Used to delete messages when the remote destination acknowledges their receipt. Args: destination: The destination server_name up_to_stream_id: Where to delete messages up to. """ def delete_messages_for_remote_destination_txn(txn): sql = ("DELETE FROM device_federation_outbox" " WHERE destination = ?" " AND stream_id <= ?") txn.execute(sql, (destination, up_to_stream_id)) await self.db_pool.runInteraction( "delete_device_msgs_for_remote", delete_messages_for_remote_destination_txn) async def get_all_new_device_messages( self, instance_name: str, last_id: int, current_id: int, limit: int) -> Tuple[List[Tuple[int, tuple]], int, bool]: """Get updates for to device replication stream. Args: instance_name: The writer we want to fetch updates from. Unused here since there is only ever one writer. last_id: The token to fetch updates from. Exclusive. current_id: The token to fetch updates up to. Inclusive. limit: The requested limit for the number of rows to return. The function may return more or fewer rows. Returns: A tuple consisting of: the updates, a token to use to fetch subsequent updates, and whether we returned fewer rows than exists between the requested tokens due to the limit. The token returned can be used in a subsequent call to this function to get further updatees. The updates are a list of 2-tuples of stream ID and the row data """ if last_id == current_id: return [], current_id, False def get_all_new_device_messages_txn(txn): # We limit like this as we might have multiple rows per stream_id, and # we want to make sure we always get all entries for any stream_id # we return. upper_pos = min(current_id, last_id + limit) sql = ("SELECT max(stream_id), user_id" " FROM device_inbox" " WHERE ? < stream_id AND stream_id <= ?" " GROUP BY user_id") txn.execute(sql, (last_id, upper_pos)) updates = [(row[0], row[1:]) for row in txn] sql = ("SELECT max(stream_id), destination" " FROM device_federation_outbox" " WHERE ? < stream_id AND stream_id <= ?" " GROUP BY destination") txn.execute(sql, (last_id, upper_pos)) updates.extend((row[0], row[1:]) for row in txn) # Order by ascending stream ordering updates.sort() limited = False upto_token = current_id if len(updates) >= limit: upto_token = updates[-1][0] limited = True return updates, upto_token, limited return await self.db_pool.runInteraction( "get_all_new_device_messages", get_all_new_device_messages_txn) @trace async def add_messages_to_device_inbox( self, local_messages_by_user_then_device: dict, remote_messages_by_destination: dict, ) -> int: """Used to send messages from this server. Args: local_messages_by_user_and_device: Dictionary of user_id to device_id to message. remote_messages_by_destination: Dictionary of destination server_name to the EDU JSON to send. Returns: The new stream_id. """ assert self._can_write_to_device def add_messages_txn(txn, now_ms, stream_id): # Add the local messages directly to the local inbox. self._add_messages_to_local_device_inbox_txn( txn, stream_id, local_messages_by_user_then_device) # Add the remote messages to the federation outbox. # We'll send them to a remote server when we next send a # federation transaction to that destination. self.db_pool.simple_insert_many_txn( txn, table="device_federation_outbox", values=[{ "destination": destination, "stream_id": stream_id, "queued_ts": now_ms, "messages_json": json_encoder.encode(edu), "instance_name": self._instance_name, } for destination, edu in remote_messages_by_destination.items()], ) async with self._device_inbox_id_gen.get_next() as stream_id: now_ms = self.clock.time_msec() await self.db_pool.runInteraction("add_messages_to_device_inbox", add_messages_txn, now_ms, stream_id) for user_id in local_messages_by_user_then_device.keys(): self._device_inbox_stream_cache.entity_has_changed( user_id, stream_id) for destination in remote_messages_by_destination.keys(): self._device_federation_outbox_stream_cache.entity_has_changed( destination, stream_id) return self._device_inbox_id_gen.get_current_token() async def add_messages_from_remote_to_device_inbox( self, origin: str, message_id: str, local_messages_by_user_then_device: dict) -> int: assert self._can_write_to_device def add_messages_txn(txn, now_ms, stream_id): # Check if we've already inserted a matching message_id for that # origin. This can happen if the origin doesn't receive our # acknowledgement from the first time we received the message. already_inserted = self.db_pool.simple_select_one_txn( txn, table="device_federation_inbox", keyvalues={ "origin": origin, "message_id": message_id }, retcols=("message_id", ), allow_none=True, ) if already_inserted is not None: return # Add an entry for this message_id so that we know we've processed # it. self.db_pool.simple_insert_txn( txn, table="device_federation_inbox", values={ "origin": origin, "message_id": message_id, "received_ts": now_ms, }, ) # Add the messages to the approriate local device inboxes so that # they'll be sent to the devices when they next sync. self._add_messages_to_local_device_inbox_txn( txn, stream_id, local_messages_by_user_then_device) async with self._device_inbox_id_gen.get_next() as stream_id: now_ms = self.clock.time_msec() await self.db_pool.runInteraction( "add_messages_from_remote_to_device_inbox", add_messages_txn, now_ms, stream_id, ) for user_id in local_messages_by_user_then_device.keys(): self._device_inbox_stream_cache.entity_has_changed( user_id, stream_id) return stream_id def _add_messages_to_local_device_inbox_txn(self, txn, stream_id, messages_by_user_then_device): assert self._can_write_to_device local_by_user_then_device = {} for user_id, messages_by_device in messages_by_user_then_device.items( ): messages_json_for_user = {} devices = list(messages_by_device.keys()) if len(devices) == 1 and devices[0] == "*": # Handle wildcard device_ids. devices = self.db_pool.simple_select_onecol_txn( txn, table="devices", keyvalues={"user_id": user_id}, retcol="device_id", ) message_json = json_encoder.encode(messages_by_device["*"]) for device_id in devices: # Add the message for all devices for this user on this # server. messages_json_for_user[device_id] = message_json else: if not devices: continue rows = self.db_pool.simple_select_many_txn( txn, table="devices", keyvalues={"user_id": user_id}, column="device_id", iterable=devices, retcols=("device_id", ), ) for row in rows: # Only insert into the local inbox if the device exists on # this server device_id = row["device_id"] message_json = json_encoder.encode( messages_by_device[device_id]) messages_json_for_user[device_id] = message_json if messages_json_for_user: local_by_user_then_device[user_id] = messages_json_for_user if not local_by_user_then_device: return self.db_pool.simple_insert_many_txn( txn, table="device_inbox", values=[{ "user_id": user_id, "device_id": device_id, "stream_id": stream_id, "message_json": message_json, "instance_name": self._instance_name, } for user_id, messages_by_device in local_by_user_then_device.items() for device_id, message_json in messages_by_device.items()], )
class SlavedEventStore( EventFederationWorkerStore, RoomMemberWorkerStore, EventPushActionsWorkerStore, StreamWorkerStore, StateGroupWorkerStore, EventsWorkerStore, SignatureWorkerStore, UserErasureWorkerStore, RelationsWorkerStore, BaseSlavedStore, ): def __init__(self, database: Database, db_conn, hs): self._stream_id_gen = SlavedIdTracker(db_conn, "events", "stream_ordering") self._backfill_id_gen = SlavedIdTracker(db_conn, "events", "stream_ordering", step=-1) super(SlavedEventStore, self).__init__(database, db_conn, hs) events_max = self._stream_id_gen.get_current_token() curr_state_delta_prefill, min_curr_state_delta_id = self.db.get_cache_dict( db_conn, "current_state_delta_stream", entity_column="room_id", stream_column="stream_id", max_value=events_max, # As we share the stream id with events token limit=1000, ) self._curr_state_delta_stream_cache = StreamChangeCache( "_curr_state_delta_stream_cache", min_curr_state_delta_id, prefilled_cache=curr_state_delta_prefill, ) # Cached functions can't be accessed through a class instance so we need # to reach inside the __dict__ to extract them. def get_room_max_stream_ordering(self): return self._stream_id_gen.get_current_token() def get_room_min_stream_ordering(self): return self._backfill_id_gen.get_current_token() def process_replication_rows(self, stream_name, token, rows): if stream_name == "events": self._stream_id_gen.advance(token) for row in rows: self._process_event_stream_row(token, row) elif stream_name == "backfill": self._backfill_id_gen.advance(-token) for row in rows: self.invalidate_caches_for_event( -token, row.event_id, row.room_id, row.type, row.state_key, row.redacts, row.relates_to, backfilled=True, ) return super(SlavedEventStore, self).process_replication_rows(stream_name, token, rows) def _process_event_stream_row(self, token, row): data = row.data if row.type == EventsStreamEventRow.TypeId: self.invalidate_caches_for_event( token, data.event_id, data.room_id, data.type, data.state_key, data.redacts, data.relates_to, backfilled=False, ) elif row.type == EventsStreamCurrentStateRow.TypeId: self._curr_state_delta_stream_cache.entity_has_changed( row.data.room_id, token) if data.type == EventTypes.Member: self.get_rooms_for_user_with_stream_ordering.invalidate( (data.state_key, )) else: raise Exception("Unknown events stream row type %s" % (row.type, )) def invalidate_caches_for_event( self, stream_ordering, event_id, room_id, etype, state_key, redacts, relates_to, backfilled, ): self._invalidate_get_event_cache(event_id) self.get_latest_event_ids_in_room.invalidate((room_id, )) self.get_unread_event_push_actions_by_room_for_user.invalidate_many( (room_id, )) if not backfilled: self._events_stream_cache.entity_has_changed( room_id, stream_ordering) if redacts: self._invalidate_get_event_cache(redacts) if etype == EventTypes.Member: self._membership_stream_cache.entity_has_changed( state_key, stream_ordering) self.get_invited_rooms_for_local_user.invalidate((state_key, )) if relates_to: self.get_relations_for_event.invalidate_many((relates_to, )) self.get_aggregation_groups_for_event.invalidate_many( (relates_to, )) self.get_applicable_edit.invalidate((relates_to, ))
class TypingHandler(object): def __init__(self, hs): self.store = hs.get_datastore() self.server_name = hs.config.server_name self.auth = hs.get_auth() self.is_mine_id = hs.is_mine_id self.notifier = hs.get_notifier() self.state = hs.get_state_handler() self.hs = hs self.clock = hs.get_clock() self.wheel_timer = WheelTimer(bucket_size=5000) self.federation = hs.get_federation_sender() hs.get_federation_registry().register_edu_handler( "m.typing", self._recv_edu) hs.get_distributor().observe("user_left_room", self.user_left_room) self._member_typing_until = {} # clock time we expect to stop self._member_last_federation_poke = {} self._latest_room_serial = 0 self._reset() # caches which room_ids changed at which serials self._typing_stream_change_cache = StreamChangeCache( "TypingStreamChangeCache", self._latest_room_serial, ) self.clock.looping_call( self._handle_timeouts, 5000, ) def _reset(self): """ Reset the typing handler's data caches. """ # map room IDs to serial numbers self._room_serials = {} # map room IDs to sets of users currently typing self._room_typing = {} def _handle_timeouts(self): logger.info("Checking for typing timeouts") now = self.clock.time_msec() members = set(self.wheel_timer.fetch(now)) for member in members: if not self.is_typing(member): # Nothing to do if they're no longer typing continue until = self._member_typing_until.get(member, None) if not until or until <= now: logger.info("Timing out typing for: %s", member.user_id) self._stopped_typing(member) continue # Check if we need to resend a keep alive over federation for this # user. if self.hs.is_mine_id(member.user_id): last_fed_poke = self._member_last_federation_poke.get( member, None) if not last_fed_poke or last_fed_poke + FEDERATION_PING_INTERVAL <= now: run_in_background(self._push_remote, member=member, typing=True) # Add a paranoia timer to ensure that we always have a timer for # each person typing. self.wheel_timer.insert( now=now, obj=member, then=now + 60 * 1000, ) def is_typing(self, member): return member.user_id in self._room_typing.get(member.room_id, []) @defer.inlineCallbacks def started_typing(self, target_user, auth_user, room_id, timeout): target_user_id = target_user.to_string() auth_user_id = auth_user.to_string() if not self.is_mine_id(target_user_id): raise SynapseError(400, "User is not hosted on this Home Server") if target_user_id != auth_user_id: raise AuthError(400, "Cannot set another user's typing state") yield self.auth.check_joined_room(room_id, target_user_id) logger.debug("%s has started typing in %s", target_user_id, room_id) member = RoomMember(room_id=room_id, user_id=target_user_id) was_present = member.user_id in self._room_typing.get(room_id, set()) now = self.clock.time_msec() self._member_typing_until[member] = now + timeout self.wheel_timer.insert( now=now, obj=member, then=now + timeout, ) if was_present: # No point sending another notification defer.returnValue(None) self._push_update( member=member, typing=True, ) @defer.inlineCallbacks def stopped_typing(self, target_user, auth_user, room_id): target_user_id = target_user.to_string() auth_user_id = auth_user.to_string() if not self.is_mine_id(target_user_id): raise SynapseError(400, "User is not hosted on this Home Server") if target_user_id != auth_user_id: raise AuthError(400, "Cannot set another user's typing state") yield self.auth.check_joined_room(room_id, target_user_id) logger.debug("%s has stopped typing in %s", target_user_id, room_id) member = RoomMember(room_id=room_id, user_id=target_user_id) self._stopped_typing(member) @defer.inlineCallbacks def user_left_room(self, user, room_id): user_id = user.to_string() if self.is_mine_id(user_id): member = RoomMember(room_id=room_id, user_id=user_id) yield self._stopped_typing(member) def _stopped_typing(self, member): if member.user_id not in self._room_typing.get(member.room_id, set()): # No point defer.returnValue(None) self._member_typing_until.pop(member, None) self._member_last_federation_poke.pop(member, None) self._push_update( member=member, typing=False, ) def _push_update(self, member, typing): if self.hs.is_mine_id(member.user_id): # Only send updates for changes to our own users. run_in_background(self._push_remote, member, typing) self._push_update_local(member=member, typing=typing) @defer.inlineCallbacks def _push_remote(self, member, typing): try: users = yield self.state.get_current_user_in_room(member.room_id) self._member_last_federation_poke[member] = self.clock.time_msec() now = self.clock.time_msec() self.wheel_timer.insert( now=now, obj=member, then=now + FEDERATION_PING_INTERVAL, ) for domain in set(get_domain_from_id(u) for u in users): if domain != self.server_name: logger.debug("sending typing update to %s", domain) self.federation.build_and_send_edu( destination=domain, edu_type="m.typing", content={ "room_id": member.room_id, "user_id": member.user_id, "typing": typing, }, key=member, ) except Exception: logger.exception("Error pushing typing notif to remotes") @defer.inlineCallbacks def _recv_edu(self, origin, content): room_id = content["room_id"] user_id = content["user_id"] member = RoomMember(user_id=user_id, room_id=room_id) # Check that the string is a valid user id user = UserID.from_string(user_id) if user.domain != origin: logger.info( "Got typing update from %r with bad 'user_id': %r", origin, user_id, ) return users = yield self.state.get_current_user_in_room(room_id) domains = set(get_domain_from_id(u) for u in users) if self.server_name in domains: logger.info("Got typing update from %s: %r", user_id, content) now = self.clock.time_msec() self._member_typing_until[member] = now + FEDERATION_TIMEOUT self.wheel_timer.insert( now=now, obj=member, then=now + FEDERATION_TIMEOUT, ) self._push_update_local(member=member, typing=content["typing"]) def _push_update_local(self, member, typing): room_set = self._room_typing.setdefault(member.room_id, set()) if typing: room_set.add(member.user_id) else: room_set.discard(member.user_id) self._latest_room_serial += 1 self._room_serials[member.room_id] = self._latest_room_serial self._typing_stream_change_cache.entity_has_changed( member.room_id, self._latest_room_serial, ) self.notifier.on_new_event("typing_key", self._latest_room_serial, rooms=[member.room_id]) def get_all_typing_updates(self, last_id, current_id): if last_id == current_id: return [] changed_rooms = self._typing_stream_change_cache.get_all_entities_changed( last_id, ) if changed_rooms is None: changed_rooms = self._room_serials rows = [] for room_id in changed_rooms: serial = self._room_serials[room_id] if last_id < serial <= current_id: typing = self._room_typing[room_id] rows.append((serial, room_id, list(typing))) rows.sort() return rows def get_current_token(self): return self._latest_room_serial
class SlavedEventStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedEventStore, self).__init__(db_conn, hs) self._stream_id_gen = SlavedIdTracker( db_conn, "events", "stream_ordering", ) self._backfill_id_gen = SlavedIdTracker( db_conn, "events", "stream_ordering", step=-1 ) events_max = self._stream_id_gen.get_current_token() event_cache_prefill, min_event_val = self._get_cache_dict( db_conn, "events", entity_column="room_id", stream_column="stream_ordering", max_value=events_max, ) self._events_stream_cache = StreamChangeCache( "EventsRoomStreamChangeCache", min_event_val, prefilled_cache=event_cache_prefill, ) self._membership_stream_cache = StreamChangeCache( "MembershipStreamChangeCache", events_max, ) self.stream_ordering_month_ago = 0 self._stream_order_on_start = self.get_room_max_stream_ordering() # Cached functions can't be accessed through a class instance so we need # to reach inside the __dict__ to extract them. get_rooms_for_user = RoomMemberStore.__dict__["get_rooms_for_user"] get_users_in_room = RoomMemberStore.__dict__["get_users_in_room"] get_latest_event_ids_in_room = EventFederationStore.__dict__[ "get_latest_event_ids_in_room" ] _get_current_state_for_key = StateStore.__dict__[ "_get_current_state_for_key" ] get_invited_rooms_for_user = RoomMemberStore.__dict__[ "get_invited_rooms_for_user" ] get_unread_event_push_actions_by_room_for_user = ( EventPushActionsStore.__dict__["get_unread_event_push_actions_by_room_for_user"] ) _get_state_group_for_events = ( StateStore.__dict__["_get_state_group_for_events"] ) _get_state_group_for_event = ( StateStore.__dict__["_get_state_group_for_event"] ) _get_state_groups_from_groups = ( StateStore.__dict__["_get_state_groups_from_groups"] ) _get_state_groups_from_groups_txn = ( DataStore._get_state_groups_from_groups_txn.__func__ ) _get_state_group_from_group = ( StateStore.__dict__["_get_state_group_from_group"] ) get_recent_event_ids_for_room = ( StreamStore.__dict__["get_recent_event_ids_for_room"] ) get_unread_push_actions_for_user_in_range_for_http = ( DataStore.get_unread_push_actions_for_user_in_range_for_http.__func__ ) get_unread_push_actions_for_user_in_range_for_email = ( DataStore.get_unread_push_actions_for_user_in_range_for_email.__func__ ) get_push_action_users_in_range = ( DataStore.get_push_action_users_in_range.__func__ ) get_event = DataStore.get_event.__func__ get_events = DataStore.get_events.__func__ get_current_state = DataStore.get_current_state.__func__ get_current_state_for_key = DataStore.get_current_state_for_key.__func__ get_rooms_for_user_where_membership_is = ( DataStore.get_rooms_for_user_where_membership_is.__func__ ) get_membership_changes_for_user = ( DataStore.get_membership_changes_for_user.__func__ ) get_room_events_max_id = DataStore.get_room_events_max_id.__func__ get_room_events_stream_for_room = ( DataStore.get_room_events_stream_for_room.__func__ ) get_events_around = DataStore.get_events_around.__func__ get_state_for_event = DataStore.get_state_for_event.__func__ get_state_for_events = DataStore.get_state_for_events.__func__ get_state_groups = DataStore.get_state_groups.__func__ get_state_groups_ids = DataStore.get_state_groups_ids.__func__ get_state_ids_for_event = DataStore.get_state_ids_for_event.__func__ get_state_ids_for_events = DataStore.get_state_ids_for_events.__func__ get_joined_users_from_state = DataStore.get_joined_users_from_state.__func__ get_joined_users_from_context = DataStore.get_joined_users_from_context.__func__ _get_joined_users_from_context = ( RoomMemberStore.__dict__["_get_joined_users_from_context"] ) get_recent_events_for_room = DataStore.get_recent_events_for_room.__func__ get_room_events_stream_for_rooms = ( DataStore.get_room_events_stream_for_rooms.__func__ ) is_host_joined = DataStore.is_host_joined.__func__ _is_host_joined = RoomMemberStore.__dict__["_is_host_joined"] get_stream_token_for_event = DataStore.get_stream_token_for_event.__func__ _set_before_and_after = staticmethod(DataStore._set_before_and_after) _get_events = DataStore._get_events.__func__ _get_events_from_cache = DataStore._get_events_from_cache.__func__ _invalidate_get_event_cache = DataStore._invalidate_get_event_cache.__func__ _enqueue_events = DataStore._enqueue_events.__func__ _do_fetch = DataStore._do_fetch.__func__ _fetch_event_rows = DataStore._fetch_event_rows.__func__ _get_event_from_row = DataStore._get_event_from_row.__func__ _get_rooms_for_user_where_membership_is_txn = ( DataStore._get_rooms_for_user_where_membership_is_txn.__func__ ) _get_members_rows_txn = DataStore._get_members_rows_txn.__func__ _get_state_for_groups = DataStore._get_state_for_groups.__func__ _get_all_state_from_cache = DataStore._get_all_state_from_cache.__func__ _get_events_around_txn = DataStore._get_events_around_txn.__func__ _get_some_state_from_cache = DataStore._get_some_state_from_cache.__func__ get_backfill_events = DataStore.get_backfill_events.__func__ _get_backfill_events = DataStore._get_backfill_events.__func__ get_missing_events = DataStore.get_missing_events.__func__ _get_missing_events = DataStore._get_missing_events.__func__ get_auth_chain = DataStore.get_auth_chain.__func__ get_auth_chain_ids = DataStore.get_auth_chain_ids.__func__ _get_auth_chain_ids_txn = DataStore._get_auth_chain_ids_txn.__func__ get_room_max_stream_ordering = DataStore.get_room_max_stream_ordering.__func__ get_forward_extremeties_for_room = ( DataStore.get_forward_extremeties_for_room.__func__ ) _get_forward_extremeties_for_room = ( EventFederationStore.__dict__["_get_forward_extremeties_for_room"] ) def stream_positions(self): result = super(SlavedEventStore, self).stream_positions() result["events"] = self._stream_id_gen.get_current_token() result["backfill"] = -self._backfill_id_gen.get_current_token() return result def process_replication(self, result): state_resets = set( r[0] for r in result.get("state_resets", {"rows": []})["rows"] ) stream = result.get("events") if stream: self._stream_id_gen.advance(int(stream["position"])) for row in stream["rows"]: self._process_replication_row( row, backfilled=False, state_resets=state_resets ) stream = result.get("backfill") if stream: self._backfill_id_gen.advance(-int(stream["position"])) for row in stream["rows"]: self._process_replication_row( row, backfilled=True, state_resets=state_resets ) stream = result.get("forward_ex_outliers") if stream: self._stream_id_gen.advance(int(stream["position"])) for row in stream["rows"]: event_id = row[1] self._invalidate_get_event_cache(event_id) stream = result.get("backward_ex_outliers") if stream: self._backfill_id_gen.advance(-int(stream["position"])) for row in stream["rows"]: event_id = row[1] self._invalidate_get_event_cache(event_id) return super(SlavedEventStore, self).process_replication(result) def _process_replication_row(self, row, backfilled, state_resets): position = row[0] internal = json.loads(row[1]) event_json = json.loads(row[2]) event = FrozenEvent(event_json, internal_metadata_dict=internal) self.invalidate_caches_for_event( event, backfilled, reset_state=position in state_resets ) def invalidate_caches_for_event(self, event, backfilled, reset_state): if reset_state: self._get_current_state_for_key.invalidate_all() self.get_rooms_for_user.invalidate_all() self.get_users_in_room.invalidate((event.room_id,)) self._invalidate_get_event_cache(event.event_id) self.get_latest_event_ids_in_room.invalidate((event.room_id,)) self.get_unread_event_push_actions_by_room_for_user.invalidate_many( (event.room_id,) ) if not backfilled: self._events_stream_cache.entity_has_changed( event.room_id, event.internal_metadata.stream_ordering ) # self.get_unread_event_push_actions_by_room_for_user.invalidate_many( # (event.room_id,) # ) if event.type == EventTypes.Redaction: self._invalidate_get_event_cache(event.redacts) if event.type == EventTypes.Member: self.get_rooms_for_user.invalidate((event.state_key,)) self.get_users_in_room.invalidate((event.room_id,)) self._membership_stream_cache.entity_has_changed( event.state_key, event.internal_metadata.stream_ordering ) self.get_invited_rooms_for_user.invalidate((event.state_key,)) if not event.is_state(): return if backfilled: return if (not event.internal_metadata.is_invite_from_remote() and event.internal_metadata.is_outlier()): return self._get_current_state_for_key.invalidate(( event.room_id, event.type, event.state_key ))
class SlavedAccountDataStore(BaseSlavedStore): def __init__(self, db_conn, hs): super(SlavedAccountDataStore, self).__init__(db_conn, hs) self._account_data_id_gen = SlavedIdTracker( db_conn, "account_data_max_stream_id", "stream_id", ) self._account_data_stream_cache = StreamChangeCache( "AccountDataAndTagsChangeCache", self._account_data_id_gen.get_current_token(), ) get_account_data_for_user = ( AccountDataStore.__dict__["get_account_data_for_user"] ) get_global_account_data_by_type_for_users = ( AccountDataStore.__dict__["get_global_account_data_by_type_for_users"] ) get_global_account_data_by_type_for_user = ( AccountDataStore.__dict__["get_global_account_data_by_type_for_user"] ) get_tags_for_user = TagsStore.__dict__["get_tags_for_user"] get_updated_tags = DataStore.get_updated_tags.__func__ get_updated_account_data_for_user = ( DataStore.get_updated_account_data_for_user.__func__ ) def get_max_account_data_stream_id(self): return self._account_data_id_gen.get_current_token() def stream_positions(self): result = super(SlavedAccountDataStore, self).stream_positions() position = self._account_data_id_gen.get_current_token() result["user_account_data"] = position result["room_account_data"] = position result["tag_account_data"] = position return result def process_replication(self, result): stream = result.get("user_account_data") if stream: self._account_data_id_gen.advance(int(stream["position"])) for row in stream["rows"]: position, user_id, data_type = row[:3] self.get_global_account_data_by_type_for_user.invalidate( (data_type, user_id,) ) self.get_account_data_for_user.invalidate((user_id,)) self._account_data_stream_cache.entity_has_changed( user_id, position ) stream = result.get("room_account_data") if stream: self._account_data_id_gen.advance(int(stream["position"])) for row in stream["rows"]: position, user_id = row[:2] self.get_account_data_for_user.invalidate((user_id,)) self._account_data_stream_cache.entity_has_changed( user_id, position ) stream = result.get("tag_account_data") if stream: self._account_data_id_gen.advance(int(stream["position"])) for row in stream["rows"]: position, user_id = row[:2] self.get_tags_for_user.invalidate((user_id,)) self._account_data_stream_cache.entity_has_changed( user_id, position ) return super(SlavedAccountDataStore, self).process_replication(result)
class PresenceStore(PresenceBackgroundUpdateStore): def __init__( self, database: DatabasePool, db_conn: LoggingDatabaseConnection, hs: "HomeServer", ): super().__init__(database, db_conn, hs) self._can_persist_presence = (hs.get_instance_name() in hs.config.worker.writers.presence) if isinstance(database.engine, PostgresEngine): self._presence_id_gen = MultiWriterIdGenerator( db_conn=db_conn, db=database, stream_name="presence_stream", instance_name=self._instance_name, tables=[("presence_stream", "instance_name", "stream_id")], sequence_name="presence_stream_sequence", writers=hs.config.worker.writers.presence, ) else: self._presence_id_gen = StreamIdGenerator(db_conn, "presence_stream", "stream_id") self.hs = hs self._presence_on_startup = self._get_active_presence(db_conn) presence_cache_prefill, min_presence_val = self.db_pool.get_cache_dict( db_conn, "presence_stream", entity_column="user_id", stream_column="stream_id", max_value=self._presence_id_gen.get_current_token(), ) self.presence_stream_cache = StreamChangeCache( "PresenceStreamChangeCache", min_presence_val, prefilled_cache=presence_cache_prefill, ) async def update_presence(self, presence_states) -> Tuple[int, int]: assert self._can_persist_presence stream_ordering_manager = self._presence_id_gen.get_next_mult( len(presence_states)) async with stream_ordering_manager as stream_orderings: await self.db_pool.runInteraction( "update_presence", self._update_presence_txn, stream_orderings, presence_states, ) return stream_orderings[-1], self._presence_id_gen.get_current_token() def _update_presence_txn(self, txn, stream_orderings, presence_states): for stream_id, state in zip(stream_orderings, presence_states): txn.call_after(self.presence_stream_cache.entity_has_changed, state.user_id, stream_id) txn.call_after(self._get_presence_for_user.invalidate, (state.user_id, )) # Delete old rows to stop database from getting really big sql = "DELETE FROM presence_stream WHERE stream_id < ? AND " for states in batch_iter(presence_states, 50): clause, args = make_in_list_sql_clause(self.database_engine, "user_id", [s.user_id for s in states]) txn.execute(sql + clause, [stream_id] + list(args)) # Actually insert new rows self.db_pool.simple_insert_many_txn( txn, table="presence_stream", keys=( "stream_id", "user_id", "state", "last_active_ts", "last_federation_update_ts", "last_user_sync_ts", "status_msg", "currently_active", "instance_name", ), values=[( stream_id, state.user_id, state.state, state.last_active_ts, state.last_federation_update_ts, state.last_user_sync_ts, state.status_msg, state.currently_active, self._instance_name, ) for stream_id, state in zip(stream_orderings, presence_states)], ) async def get_all_presence_updates( self, instance_name: str, last_id: int, current_id: int, limit: int) -> Tuple[List[Tuple[int, list]], int, bool]: """Get updates for presence replication stream. Args: instance_name: The writer we want to fetch updates from. Unused here since there is only ever one writer. last_id: The token to fetch updates from. Exclusive. current_id: The token to fetch updates up to. Inclusive. limit: The requested limit for the number of rows to return. The function may return more or fewer rows. Returns: A tuple consisting of: the updates, a token to use to fetch subsequent updates, and whether we returned fewer rows than exists between the requested tokens due to the limit. The token returned can be used in a subsequent call to this function to get further updatees. The updates are a list of 2-tuples of stream ID and the row data """ if last_id == current_id: return [], current_id, False def get_all_presence_updates_txn(txn): sql = """ SELECT stream_id, user_id, state, last_active_ts, last_federation_update_ts, last_user_sync_ts, status_msg, currently_active FROM presence_stream WHERE ? < stream_id AND stream_id <= ? ORDER BY stream_id ASC LIMIT ? """ txn.execute(sql, (last_id, current_id, limit)) updates = [(row[0], row[1:]) for row in txn] upper_bound = current_id limited = False if len(updates) >= limit: upper_bound = updates[-1][0] limited = True return updates, upper_bound, limited return await self.db_pool.runInteraction("get_all_presence_updates", get_all_presence_updates_txn) @cached() def _get_presence_for_user(self, user_id): raise NotImplementedError() @cachedList( cached_method_name="_get_presence_for_user", list_name="user_ids", num_args=1, ) async def get_presence_for_users(self, user_ids): rows = await self.db_pool.simple_select_many_batch( table="presence_stream", column="user_id", iterable=user_ids, keyvalues={}, retcols=( "user_id", "state", "last_active_ts", "last_federation_update_ts", "last_user_sync_ts", "status_msg", "currently_active", ), desc="get_presence_for_users", ) for row in rows: row["currently_active"] = bool(row["currently_active"]) return {row["user_id"]: UserPresenceState(**row) for row in rows} async def should_user_receive_full_presence_with_token( self, user_id: str, from_token: int, ) -> bool: """Check whether the given user should receive full presence using the stream token they're updating from. Args: user_id: The ID of the user to check. from_token: The stream token included in their /sync token. Returns: True if the user should have full presence sent to them, False otherwise. """ def _should_user_receive_full_presence_with_token_txn(txn): sql = """ SELECT 1 FROM users_to_send_full_presence_to WHERE user_id = ? AND presence_stream_id >= ? """ txn.execute(sql, (user_id, from_token)) return bool(txn.fetchone()) return await self.db_pool.runInteraction( "should_user_receive_full_presence_with_token", _should_user_receive_full_presence_with_token_txn, ) async def add_users_to_send_full_presence_to(self, user_ids: Iterable[str]): """Adds to the list of users who should receive a full snapshot of presence upon their next sync. Args: user_ids: An iterable of user IDs. """ # Add user entries to the table, updating the presence_stream_id column if the user already # exists in the table. presence_stream_id = self._presence_id_gen.get_current_token() await self.db_pool.simple_upsert_many( table="users_to_send_full_presence_to", key_names=("user_id", ), key_values=[(user_id, ) for user_id in user_ids], value_names=("presence_stream_id", ), # We save the current presence stream ID token along with the user ID entry so # that when a user /sync's, even if they syncing multiple times across separate # devices at different times, each device will receive full presence once - when # the presence stream ID in their sync token is less than the one in the table # for their user ID. value_values=[(presence_stream_id, ) for _ in user_ids], desc="add_users_to_send_full_presence_to", ) async def get_presence_for_all_users( self, include_offline: bool = True, ) -> Dict[str, UserPresenceState]: """Retrieve the current presence state for all users. Note that the presence_stream table is culled frequently, so it should only contain the latest presence state for each user. Args: include_offline: Whether to include offline presence states Returns: A dict of user IDs to their current UserPresenceState. """ users_to_state = {} exclude_keyvalues = None if not include_offline: # Exclude offline presence state exclude_keyvalues = {"state": "offline"} # This may be a very heavy database query. # We paginate in order to not block a database connection. limit = 100 offset = 0 while True: rows = await self.db_pool.runInteraction( "get_presence_for_all_users", self.db_pool.simple_select_list_paginate_txn, "presence_stream", orderby="stream_id", start=offset, limit=limit, exclude_keyvalues=exclude_keyvalues, retcols=( "user_id", "state", "last_active_ts", "last_federation_update_ts", "last_user_sync_ts", "status_msg", "currently_active", ), order_direction="ASC", ) for row in rows: users_to_state[row["user_id"]] = UserPresenceState(**row) # We've run out of updates to query if len(rows) < limit: break offset += limit return users_to_state def get_current_presence_token(self): return self._presence_id_gen.get_current_token() def _get_active_presence(self, db_conn: Connection): """Fetch non-offline presence from the database so that we can register the appropriate time outs. """ # The `presence_stream_state_not_offline_idx` index should be used for this # query. sql = ( "SELECT user_id, state, last_active_ts, last_federation_update_ts," " last_user_sync_ts, status_msg, currently_active FROM presence_stream" " WHERE state != ?") txn = db_conn.cursor() txn.execute(sql, (PresenceState.OFFLINE, )) rows = self.db_pool.cursor_to_dict(txn) txn.close() for row in rows: row["currently_active"] = bool(row["currently_active"]) return [UserPresenceState(**row) for row in rows] def take_presence_startup_info(self): active_on_startup = self._presence_on_startup self._presence_on_startup = None return active_on_startup def process_replication_rows(self, stream_name, instance_name, token, rows): if stream_name == PresenceStream.NAME: self._presence_id_gen.advance(instance_name, token) for row in rows: self.presence_stream_cache.entity_has_changed( row.user_id, token) self._get_presence_for_user.invalidate((row.user_id, )) return super().process_replication_rows(stream_name, instance_name, token, rows)
class TypingHandler(object): def __init__(self, hs): self.store = hs.get_datastore() self.server_name = hs.config.server_name self.auth = hs.get_auth() self.is_mine_id = hs.is_mine_id self.notifier = hs.get_notifier() self.state = hs.get_state_handler() self.hs = hs self.clock = hs.get_clock() self.wheel_timer = WheelTimer(bucket_size=5000) self.federation = hs.get_federation_sender() hs.get_federation_registry().register_edu_handler("m.typing", self._recv_edu) hs.get_distributor().observe("user_left_room", self.user_left_room) self._member_typing_until = {} # clock time we expect to stop self._member_last_federation_poke = {} self._latest_room_serial = 0 self._reset() # caches which room_ids changed at which serials self._typing_stream_change_cache = StreamChangeCache( "TypingStreamChangeCache", self._latest_room_serial, ) self.clock.looping_call( self._handle_timeouts, 5000, ) def _reset(self): """ Reset the typing handler's data caches. """ # map room IDs to serial numbers self._room_serials = {} # map room IDs to sets of users currently typing self._room_typing = {} def _handle_timeouts(self): logger.info("Checking for typing timeouts") now = self.clock.time_msec() members = set(self.wheel_timer.fetch(now)) for member in members: if not self.is_typing(member): # Nothing to do if they're no longer typing continue until = self._member_typing_until.get(member, None) if not until or until <= now: logger.info("Timing out typing for: %s", member.user_id) self._stopped_typing(member) continue # Check if we need to resend a keep alive over federation for this # user. if self.hs.is_mine_id(member.user_id): last_fed_poke = self._member_last_federation_poke.get(member, None) if not last_fed_poke or last_fed_poke + FEDERATION_PING_INTERVAL <= now: run_in_background( self._push_remote, member=member, typing=True ) # Add a paranoia timer to ensure that we always have a timer for # each person typing. self.wheel_timer.insert( now=now, obj=member, then=now + 60 * 1000, ) def is_typing(self, member): return member.user_id in self._room_typing.get(member.room_id, []) @defer.inlineCallbacks def started_typing(self, target_user, auth_user, room_id, timeout): target_user_id = target_user.to_string() auth_user_id = auth_user.to_string() if not self.is_mine_id(target_user_id): raise SynapseError(400, "User is not hosted on this Home Server") if target_user_id != auth_user_id: raise AuthError(400, "Cannot set another user's typing state") yield self.auth.check_joined_room(room_id, target_user_id) logger.debug( "%s has started typing in %s", target_user_id, room_id ) member = RoomMember(room_id=room_id, user_id=target_user_id) was_present = member.user_id in self._room_typing.get(room_id, set()) now = self.clock.time_msec() self._member_typing_until[member] = now + timeout self.wheel_timer.insert( now=now, obj=member, then=now + timeout, ) if was_present: # No point sending another notification defer.returnValue(None) self._push_update( member=member, typing=True, ) @defer.inlineCallbacks def stopped_typing(self, target_user, auth_user, room_id): target_user_id = target_user.to_string() auth_user_id = auth_user.to_string() if not self.is_mine_id(target_user_id): raise SynapseError(400, "User is not hosted on this Home Server") if target_user_id != auth_user_id: raise AuthError(400, "Cannot set another user's typing state") yield self.auth.check_joined_room(room_id, target_user_id) logger.debug( "%s has stopped typing in %s", target_user_id, room_id ) member = RoomMember(room_id=room_id, user_id=target_user_id) self._stopped_typing(member) @defer.inlineCallbacks def user_left_room(self, user, room_id): user_id = user.to_string() if self.is_mine_id(user_id): member = RoomMember(room_id=room_id, user_id=user_id) yield self._stopped_typing(member) def _stopped_typing(self, member): if member.user_id not in self._room_typing.get(member.room_id, set()): # No point defer.returnValue(None) self._member_typing_until.pop(member, None) self._member_last_federation_poke.pop(member, None) self._push_update( member=member, typing=False, ) def _push_update(self, member, typing): if self.hs.is_mine_id(member.user_id): # Only send updates for changes to our own users. run_in_background(self._push_remote, member, typing) self._push_update_local( member=member, typing=typing ) @defer.inlineCallbacks def _push_remote(self, member, typing): try: users = yield self.state.get_current_user_in_room(member.room_id) self._member_last_federation_poke[member] = self.clock.time_msec() now = self.clock.time_msec() self.wheel_timer.insert( now=now, obj=member, then=now + FEDERATION_PING_INTERVAL, ) for domain in set(get_domain_from_id(u) for u in users): if domain != self.server_name: logger.debug("sending typing update to %s", domain) self.federation.send_edu( destination=domain, edu_type="m.typing", content={ "room_id": member.room_id, "user_id": member.user_id, "typing": typing, }, key=member, ) except Exception: logger.exception("Error pushing typing notif to remotes") @defer.inlineCallbacks def _recv_edu(self, origin, content): room_id = content["room_id"] user_id = content["user_id"] member = RoomMember(user_id=user_id, room_id=room_id) # Check that the string is a valid user id user = UserID.from_string(user_id) if user.domain != origin: logger.info( "Got typing update from %r with bad 'user_id': %r", origin, user_id, ) return users = yield self.state.get_current_user_in_room(room_id) domains = set(get_domain_from_id(u) for u in users) if self.server_name in domains: logger.info("Got typing update from %s: %r", user_id, content) now = self.clock.time_msec() self._member_typing_until[member] = now + FEDERATION_TIMEOUT self.wheel_timer.insert( now=now, obj=member, then=now + FEDERATION_TIMEOUT, ) self._push_update_local( member=member, typing=content["typing"] ) def _push_update_local(self, member, typing): room_set = self._room_typing.setdefault(member.room_id, set()) if typing: room_set.add(member.user_id) else: room_set.discard(member.user_id) self._latest_room_serial += 1 self._room_serials[member.room_id] = self._latest_room_serial self._typing_stream_change_cache.entity_has_changed( member.room_id, self._latest_room_serial, ) self.notifier.on_new_event( "typing_key", self._latest_room_serial, rooms=[member.room_id] ) def get_all_typing_updates(self, last_id, current_id): if last_id == current_id: return [] changed_rooms = self._typing_stream_change_cache.get_all_entities_changed( last_id, ) if changed_rooms is None: changed_rooms = self._room_serials rows = [] for room_id in changed_rooms: serial = self._room_serials[room_id] if last_id < serial <= current_id: typing = self._room_typing[room_id] rows.append((serial, room_id, list(typing))) rows.sort() return rows def get_current_token(self): return self._latest_room_serial