Correctly read to-device stream pos on SQLite (#16682)

This commit is contained in:
David Robertson 2023-11-24 13:42:38 +00:00 committed by GitHub
parent 32a59a6495
commit c3627d0f99
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 21 additions and 13 deletions

1
changelog.d/16682.misc Normal file
View File

@ -0,0 +1 @@
Correctly read the to-device stream ID on startup using SQLite.

View File

@ -621,7 +621,7 @@ class ToDeviceStream(_StreamFromIdGen):
super().__init__( super().__init__(
hs.get_instance_name(), hs.get_instance_name(),
store.get_all_new_device_messages, store.get_all_new_device_messages,
store._device_inbox_id_gen, store._to_device_msg_id_gen,
) )

View File

@ -87,25 +87,32 @@ class DeviceInboxWorkerStore(SQLBaseStore):
self._instance_name in hs.config.worker.writers.to_device self._instance_name in hs.config.worker.writers.to_device
) )
self._device_inbox_id_gen: AbstractStreamIdGenerator = ( self._to_device_msg_id_gen: AbstractStreamIdGenerator = (
MultiWriterIdGenerator( MultiWriterIdGenerator(
db_conn=db_conn, db_conn=db_conn,
db=database, db=database,
notifier=hs.get_replication_notifier(), notifier=hs.get_replication_notifier(),
stream_name="to_device", stream_name="to_device",
instance_name=self._instance_name, instance_name=self._instance_name,
tables=[("device_inbox", "instance_name", "stream_id")], tables=[
("device_inbox", "instance_name", "stream_id"),
("device_federation_outbox", "instance_name", "stream_id"),
],
sequence_name="device_inbox_sequence", sequence_name="device_inbox_sequence",
writers=hs.config.worker.writers.to_device, writers=hs.config.worker.writers.to_device,
) )
) )
else: else:
self._can_write_to_device = True self._can_write_to_device = True
self._device_inbox_id_gen = StreamIdGenerator( self._to_device_msg_id_gen = StreamIdGenerator(
db_conn, hs.get_replication_notifier(), "device_inbox", "stream_id" db_conn,
hs.get_replication_notifier(),
"device_inbox",
"stream_id",
extra_tables=[("device_federation_outbox", "stream_id")],
) )
max_device_inbox_id = self._device_inbox_id_gen.get_current_token() max_device_inbox_id = self._to_device_msg_id_gen.get_current_token()
device_inbox_prefill, min_device_inbox_id = self.db_pool.get_cache_dict( device_inbox_prefill, min_device_inbox_id = self.db_pool.get_cache_dict(
db_conn, db_conn,
"device_inbox", "device_inbox",
@ -145,8 +152,8 @@ class DeviceInboxWorkerStore(SQLBaseStore):
) -> None: ) -> None:
if stream_name == ToDeviceStream.NAME: if stream_name == ToDeviceStream.NAME:
# If replication is happening than postgres must be being used. # If replication is happening than postgres must be being used.
assert isinstance(self._device_inbox_id_gen, MultiWriterIdGenerator) assert isinstance(self._to_device_msg_id_gen, MultiWriterIdGenerator)
self._device_inbox_id_gen.advance(instance_name, token) self._to_device_msg_id_gen.advance(instance_name, token)
for row in rows: for row in rows:
if row.entity.startswith("@"): if row.entity.startswith("@"):
self._device_inbox_stream_cache.entity_has_changed( self._device_inbox_stream_cache.entity_has_changed(
@ -162,11 +169,11 @@ class DeviceInboxWorkerStore(SQLBaseStore):
self, stream_name: str, instance_name: str, token: int self, stream_name: str, instance_name: str, token: int
) -> None: ) -> None:
if stream_name == ToDeviceStream.NAME: if stream_name == ToDeviceStream.NAME:
self._device_inbox_id_gen.advance(instance_name, token) self._to_device_msg_id_gen.advance(instance_name, token)
super().process_replication_position(stream_name, instance_name, token) super().process_replication_position(stream_name, instance_name, token)
def get_to_device_stream_token(self) -> int: def get_to_device_stream_token(self) -> int:
return self._device_inbox_id_gen.get_current_token() return self._to_device_msg_id_gen.get_current_token()
async def get_messages_for_user_devices( async def get_messages_for_user_devices(
self, self,
@ -801,7 +808,7 @@ class DeviceInboxWorkerStore(SQLBaseStore):
msg.get(EventContentFields.TO_DEVICE_MSGID), msg.get(EventContentFields.TO_DEVICE_MSGID),
) )
async with self._device_inbox_id_gen.get_next() as stream_id: async with self._to_device_msg_id_gen.get_next() as stream_id:
now_ms = self._clock.time_msec() now_ms = self._clock.time_msec()
await self.db_pool.runInteraction( await self.db_pool.runInteraction(
"add_messages_to_device_inbox", add_messages_txn, now_ms, stream_id "add_messages_to_device_inbox", add_messages_txn, now_ms, stream_id
@ -813,7 +820,7 @@ class DeviceInboxWorkerStore(SQLBaseStore):
destination, stream_id destination, stream_id
) )
return self._device_inbox_id_gen.get_current_token() return self._to_device_msg_id_gen.get_current_token()
async def add_messages_from_remote_to_device_inbox( async def add_messages_from_remote_to_device_inbox(
self, self,
@ -857,7 +864,7 @@ class DeviceInboxWorkerStore(SQLBaseStore):
txn, stream_id, local_messages_by_user_then_device txn, stream_id, local_messages_by_user_then_device
) )
async with self._device_inbox_id_gen.get_next() as stream_id: async with self._to_device_msg_id_gen.get_next() as stream_id:
now_ms = self._clock.time_msec() now_ms = self._clock.time_msec()
await self.db_pool.runInteraction( await self.db_pool.runInteraction(
"add_messages_from_remote_to_device_inbox", "add_messages_from_remote_to_device_inbox",