Use different stream IDs for device_list_outbound_pokes

synapse/storage/databases/main/devices.py

@@ -1553,28 +1553,24 @@ async def add_device_change_to_streams(
  if not device_ids:
  return None
 
- num_stream_ids = len(device_ids)
- if hosts:
- # The `device_lists_outbound_pokes` wants a different stream ID for
- # each row, which is a row per host per device update.
- num_stream_ids = len(hosts) * len(device_ids)
-
  context = get_active_span_text_map()
 
- def add_device_changes_txn(txn, stream_ids):
+ def add_device_changes_txn(
+ txn, stream_ids_for_device_change, stream_ids_for_outbound_pokes
+ ):
  self._add_device_change_to_stream_txn(
  txn,
  user_id,
  device_ids,
- stream_ids,
+ stream_ids_for_device_change,
  )
 
  self._add_device_outbound_room_poke_txn(
  txn,
  user_id,
  device_ids,
  room_ids,
- stream_ids,
+ stream_ids_for_device_change,
  context,
  hosts_have_been_calculated=hosts is not None,
  )
@@ -1590,15 +1586,27 @@ def add_device_changes_txn(txn, stream_ids):
  user_id,
  device_ids,
  hosts,
- stream_ids,
+ stream_ids_for_outbound_pokes,
  context,
  )
 
+ # The `device_lists_stream` wants a stream ID per device update.
+ num_stream_ids = len(device_ids)
+
+ if hosts:
+ # The `device_lists_outbound_pokes` wants a different stream ID for
+ # each row, which is a row per host per device update.
+ num_stream_ids += len(hosts) * len(device_ids)
+
  async with self._device_list_id_gen.get_next_mult(num_stream_ids) as stream_ids:
+ stream_ids_for_device_change = stream_ids[: len(device_ids)]
+ stream_ids_for_outbound_pokes = stream_ids[len(device_ids) :]
+
  await self.db_pool.runInteraction(
  "add_device_change_to_stream",
  add_device_changes_txn,
- stream_ids,
+ stream_ids_for_device_change,
+ stream_ids_for_outbound_pokes,
  )
 
  return stream_ids[-1]