implement lazy-loading for incremental sync

Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
Jason Volk 2025-01-31 08:34:32 +00:00
parent a4ef04cd14
commit 6983798487

View file

@ -55,7 +55,10 @@ use ruma::{
}; };
use super::{load_timeline, share_encrypted_room}; use super::{load_timeline, share_encrypted_room};
use crate::{client::ignored_filter, Ruma, RumaResponse}; use crate::{
client::{ignored_filter, lazy_loading_witness},
Ruma, RumaResponse,
};
#[derive(Default)] #[derive(Default)]
struct StateChanges { struct StateChanges {
@ -633,10 +636,6 @@ async fn load_joined_room(
}) })
.into(); .into();
let no_state_changes = timeline_pdus.is_empty()
&& (since_shortstatehash.is_none()
|| since_shortstatehash.is_some_and(is_equal_to!(current_shortstatehash)));
let since_sender_member: OptionFuture<_> = since_shortstatehash let since_sender_member: OptionFuture<_> = since_shortstatehash
.map(|short| { .map(|short| {
services services
@ -658,11 +657,7 @@ async fn load_joined_room(
let lazy_loading_enabled = filter.room.state.lazy_load_options.is_enabled() let lazy_loading_enabled = filter.room.state.lazy_load_options.is_enabled()
|| filter.room.timeline.lazy_load_options.is_enabled(); || filter.room.timeline.lazy_load_options.is_enabled();
let generate_witness = let lazy_reset = since_shortstatehash.is_none();
lazy_loading_enabled && (since_shortstatehash.is_none() || joined_since_last_sync);
let lazy_reset = lazy_loading_enabled && since_shortstatehash.is_none();
let lazy_loading_context = &lazy_loading::Context { let lazy_loading_context = &lazy_loading::Context {
user_id: sender_user, user_id: sender_user,
device_id: sender_device, device_id: sender_device,
@ -677,24 +672,10 @@ async fn load_joined_room(
.into(); .into();
lazy_load_reset.await; lazy_load_reset.await;
let witness: Option<Witness> = generate_witness.then(|| { let witness: OptionFuture<_> = lazy_loading_enabled
timeline_pdus .then(|| lazy_loading_witness(services, lazy_loading_context, timeline_pdus.iter()))
.iter()
.map(|(_, pdu)| pdu.sender.clone())
.chain(receipt_events.keys().cloned())
.collect()
});
let witness: OptionFuture<_> = witness
.map(|witness| {
services
.rooms
.lazy_loading
.witness_retain(witness, lazy_loading_context)
})
.into(); .into();
let witness = witness.await;
let StateChanges { let StateChanges {
heroes, heroes,
joined_member_count, joined_member_count,
@ -703,23 +684,19 @@ async fn load_joined_room(
state_events, state_events,
mut device_list_updates, mut device_list_updates,
left_encrypted_users, left_encrypted_users,
} = if no_state_changes { } = calculate_state_changes(
StateChanges::default() services,
} else { sender_user,
calculate_state_changes( room_id,
services, full_state,
sender_user, filter,
room_id, since_shortstatehash,
full_state, current_shortstatehash,
filter, joined_since_last_sync,
since_shortstatehash, witness.await.as_ref(),
current_shortstatehash, )
joined_since_last_sync, .boxed()
witness.as_ref(), .await?;
)
.boxed()
.await?
};
let account_data_events = services let account_data_events = services
.account_data .account_data
@ -908,6 +885,7 @@ async fn calculate_state_changes(
since_shortstatehash, since_shortstatehash,
current_shortstatehash, current_shortstatehash,
joined_since_last_sync, joined_since_last_sync,
witness,
) )
.await .await
} }
@ -920,7 +898,7 @@ async fn calculate_state_initial(
sender_user: &UserId, sender_user: &UserId,
room_id: &RoomId, room_id: &RoomId,
full_state: bool, full_state: bool,
filter: &FilterDefinition, _filter: &FilterDefinition,
current_shortstatehash: ShortStateHash, current_shortstatehash: ShortStateHash,
witness: Option<&Witness>, witness: Option<&Witness>,
) -> Result<StateChanges> { ) -> Result<StateChanges> {
@ -938,20 +916,14 @@ async fn calculate_state_initial(
.zip(event_ids.into_iter().stream()) .zip(event_ids.into_iter().stream())
.ready_filter_map(|item| Some((item.0.ok()?, item.1))) .ready_filter_map(|item| Some((item.0.ok()?, item.1)))
.ready_filter_map(|((event_type, state_key), event_id)| { .ready_filter_map(|((event_type, state_key), event_id)| {
let lazy_load_enabled = filter.room.state.lazy_load_options.is_enabled() let lazy = !full_state
|| filter.room.timeline.lazy_load_options.is_enabled();
if lazy_load_enabled
&& event_type == StateEventType::RoomMember && event_type == StateEventType::RoomMember
&& !full_state
&& state_key.as_str().try_into().is_ok_and(|user_id: &UserId| { && state_key.as_str().try_into().is_ok_and(|user_id: &UserId| {
sender_user != user_id sender_user != user_id
&& witness.is_some_and(|witness| !witness.contains(user_id)) && witness.is_some_and(|witness| !witness.contains(user_id))
}) { });
return None;
}
Some(event_id) lazy.or_some(event_id)
}) })
.broad_filter_map(|event_id: OwnedEventId| async move { .broad_filter_map(|event_id: OwnedEventId| async move {
services.rooms.timeline.get_pdu(&event_id).await.ok() services.rooms.timeline.get_pdu(&event_id).await.ok()
@ -978,7 +950,7 @@ async fn calculate_state_initial(
#[tracing::instrument(name = "incremental", level = "trace", skip_all)] #[tracing::instrument(name = "incremental", level = "trace", skip_all)]
#[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)]
async fn calculate_state_incremental( async fn calculate_state_incremental<'a>(
services: &Services, services: &Services,
sender_user: &UserId, sender_user: &UserId,
room_id: &RoomId, room_id: &RoomId,
@ -987,39 +959,80 @@ async fn calculate_state_incremental(
since_shortstatehash: Option<ShortStateHash>, since_shortstatehash: Option<ShortStateHash>,
current_shortstatehash: ShortStateHash, current_shortstatehash: ShortStateHash,
joined_since_last_sync: bool, joined_since_last_sync: bool,
witness: Option<&'a Witness>,
) -> Result<StateChanges> { ) -> Result<StateChanges> {
// Incremental /sync let since_shortstatehash = since_shortstatehash.unwrap_or(current_shortstatehash);
let since_shortstatehash =
since_shortstatehash.expect("missing since_shortstatehash on incremental sync");
let mut delta_state_events = Vec::new(); let state_changed = since_shortstatehash != current_shortstatehash;
if since_shortstatehash != current_shortstatehash { let state_get_id = |user_id: &'a UserId| {
let current_state_ids = services services
.rooms .rooms
.state_accessor .state_accessor
.state_full_ids(current_shortstatehash) .state_get_id(current_shortstatehash, &StateEventType::RoomMember, user_id.as_str())
.collect(); .ok()
};
let since_state_ids = services let lazy_state_ids: OptionFuture<_> = witness
.rooms .map(|witness| {
.state_accessor witness
.state_full_ids(since_shortstatehash) .iter()
.collect(); .stream()
.broad_filter_map(|user_id| state_get_id(user_id))
.collect::<Vec<OwnedEventId>>()
})
.into();
let (current_state_ids, since_state_ids): ( let current_state_ids: OptionFuture<_> = state_changed
HashMap<_, OwnedEventId>, .then(|| {
HashMap<_, OwnedEventId>, services
) = join(current_state_ids, since_state_ids).await; .rooms
.state_accessor
.state_full_ids(current_shortstatehash)
.collect::<Vec<(_, OwnedEventId)>>()
})
.into();
current_state_ids let since_state_ids: OptionFuture<_> = (state_changed && !full_state)
.iter() .then(|| {
.stream() services
.ready_filter(|(key, id)| full_state || since_state_ids.get(key) != Some(id)) .rooms
.wide_filter_map(|(_, id)| services.rooms.timeline.get_pdu(id).ok()) .state_accessor
.ready_for_each(|pdu| delta_state_events.push(pdu)) .state_full_ids(since_shortstatehash)
.await; .collect::<HashMap<_, OwnedEventId>>()
} })
.into();
let lazy_state_ids = lazy_state_ids
.map(Option::into_iter)
.map(|iter| iter.flat_map(Vec::into_iter))
.map(IterStream::stream)
.flatten_stream();
let ref since_state_ids = since_state_ids.shared();
let delta_state_events = current_state_ids
.map(Option::into_iter)
.map(|iter| iter.flat_map(Vec::into_iter))
.map(IterStream::stream)
.flatten_stream()
.filter_map(|(shortstatekey, event_id): (u64, OwnedEventId)| async move {
since_state_ids
.clone()
.await
.is_none_or(|since_state| since_state.get(&shortstatekey) != Some(&event_id))
.then_some(event_id)
})
.chain(lazy_state_ids)
.broad_filter_map(|event_id: OwnedEventId| async move {
services
.rooms
.timeline
.get_pdu(&event_id)
.await
.map(move |pdu| (event_id, pdu))
.ok()
})
.collect::<HashMap<_, _>>();
let since_encryption = services let since_encryption = services
.rooms .rooms
@ -1031,11 +1044,12 @@ async fn calculate_state_incremental(
.rooms .rooms
.state_accessor .state_accessor
.state_get(current_shortstatehash, &StateEventType::RoomEncryption, "") .state_get(current_shortstatehash, &StateEventType::RoomEncryption, "")
.is_ok() .is_ok();
.await;
let (delta_state_events, encrypted_room) = join(delta_state_events, encrypted_room).await;
let (mut device_list_updates, left_encrypted_users) = delta_state_events let (mut device_list_updates, left_encrypted_users) = delta_state_events
.iter() .values()
.stream() .stream()
.ready_filter(|_| encrypted_room) .ready_filter(|_| encrypted_room)
.ready_filter(|state_event| state_event.kind == RoomMember) .ready_filter(|state_event| state_event.kind == RoomMember)
@ -1084,7 +1098,7 @@ async fn calculate_state_incremental(
} }
let send_member_count = delta_state_events let send_member_count = delta_state_events
.iter() .values()
.any(|event| event.kind == RoomMember); .any(|event| event.kind == RoomMember);
let (joined_member_count, invited_member_count, heroes) = if send_member_count { let (joined_member_count, invited_member_count, heroes) = if send_member_count {
@ -1098,9 +1112,9 @@ async fn calculate_state_incremental(
joined_member_count, joined_member_count,
invited_member_count, invited_member_count,
joined_since_last_sync, joined_since_last_sync,
state_events: delta_state_events,
device_list_updates, device_list_updates,
left_encrypted_users, left_encrypted_users,
state_events: delta_state_events.into_values().collect(),
}) })
} }