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__(
hs.get_instance_name(),
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._device_inbox_id_gen: AbstractStreamIdGenerator = (
self._to_device_msg_id_gen: AbstractStreamIdGenerator = (
MultiWriterIdGenerator(
db_conn=db_conn,
db=database,
notifier=hs.get_replication_notifier(),
stream_name="to_device",
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",
writers=hs.config.worker.writers.to_device,
)
)
else:
self._can_write_to_device = True
self._device_inbox_id_gen = StreamIdGenerator(
db_conn, hs.get_replication_notifier(), "device_inbox", "stream_id"
self._to_device_msg_id_gen = StreamIdGenerator(
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(
db_conn,
"device_inbox",
@ -145,8 +152,8 @@ class DeviceInboxWorkerStore(SQLBaseStore):
) -> None:
if stream_name == ToDeviceStream.NAME:
# If replication is happening than postgres must be being used.
assert isinstance(self._device_inbox_id_gen, MultiWriterIdGenerator)
self._device_inbox_id_gen.advance(instance_name, token)
assert isinstance(self._to_device_msg_id_gen, MultiWriterIdGenerator)
self._to_device_msg_id_gen.advance(instance_name, token)
for row in rows:
if row.entity.startswith("@"):
self._device_inbox_stream_cache.entity_has_changed(
@ -162,11 +169,11 @@ class DeviceInboxWorkerStore(SQLBaseStore):
self, stream_name: str, instance_name: str, token: int
) -> None:
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)
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(
self,
@ -801,7 +808,7 @@ class DeviceInboxWorkerStore(SQLBaseStore):
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()
await self.db_pool.runInteraction(
"add_messages_to_device_inbox", add_messages_txn, now_ms, stream_id
@ -813,7 +820,7 @@ class DeviceInboxWorkerStore(SQLBaseStore):
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(
self,
@ -857,7 +864,7 @@ class DeviceInboxWorkerStore(SQLBaseStore):
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()
await self.db_pool.runInteraction(
"add_messages_from_remote_to_device_inbox",