feat: lazy loading

This commit is contained in:
Timo Kösters 2022-01-04 14:30:13 +01:00
parent 5bcc1324ed
commit 68e910bb77
No known key found for this signature in database
GPG key ID: 356E705610F626D5
5 changed files with 321 additions and 42 deletions

View file

@ -1,5 +1,9 @@
use crate::{database::DatabaseGuard, ConduitResult, Error, Ruma};
use ruma::api::client::{error::ErrorKind, r0::context::get_context};
use ruma::{
api::client::{error::ErrorKind, r0::context::get_context},
events::EventType,
};
use std::collections::HashSet;
use std::convert::TryFrom;
#[cfg(feature = "conduit_bin")]
@ -21,6 +25,7 @@ pub async fn get_context_route(
body: Ruma<get_context::Request<'_>>,
) -> ConduitResult<get_context::Response> {
let sender_user = body.sender_user.as_ref().expect("user is authenticated");
let sender_device = body.sender_device.as_ref().expect("user is authenticated");
if !db.rooms.is_joined(sender_user, &body.room_id)? {
return Err(Error::BadRequest(
@ -29,6 +34,8 @@ pub async fn get_context_route(
));
}
let mut lazy_loaded = HashSet::new();
let base_pdu_id = db
.rooms
.get_pdu_id(&body.event_id)?
@ -45,8 +52,18 @@ pub async fn get_context_route(
.ok_or(Error::BadRequest(
ErrorKind::NotFound,
"Base event not found.",
))?
.to_room_event();
))?;
if !db.rooms.lazy_load_was_sent_before(
&sender_user,
&sender_device,
&body.room_id,
&base_event.sender,
)? {
lazy_loaded.insert(base_event.sender.clone());
}
let base_event = base_event.to_room_event();
let events_before: Vec<_> = db
.rooms
@ -60,6 +77,17 @@ pub async fn get_context_route(
.filter_map(|r| r.ok()) // Remove buggy events
.collect();
for (_, event) in &events_before {
if !db.rooms.lazy_load_was_sent_before(
&sender_user,
&sender_device,
&body.room_id,
&event.sender,
)? {
lazy_loaded.insert(event.sender.clone());
}
}
let start_token = events_before
.last()
.and_then(|(pdu_id, _)| db.rooms.pdu_count(pdu_id).ok())
@ -82,6 +110,17 @@ pub async fn get_context_route(
.filter_map(|r| r.ok()) // Remove buggy events
.collect();
for (_, event) in &events_after {
if !db.rooms.lazy_load_was_sent_before(
&sender_user,
&sender_device,
&body.room_id,
&event.sender,
)? {
lazy_loaded.insert(event.sender.clone());
}
}
let end_token = events_after
.last()
.and_then(|(pdu_id, _)| db.rooms.pdu_count(pdu_id).ok())
@ -92,18 +131,23 @@ pub async fn get_context_route(
.map(|(_, pdu)| pdu.to_room_event())
.collect();
let mut state = Vec::new();
for ll_id in &lazy_loaded {
if let Some(member_event) =
db.rooms
.room_state_get(&body.room_id, &EventType::RoomMember, ll_id.as_str())?
{
state.push(member_event.to_state_event());
}
}
let resp = get_context::Response {
start: start_token,
end: end_token,
events_before,
event: Some(base_event),
events_after,
state: db // TODO: State at event
.rooms
.room_state_full(&body.room_id)?
.values()
.map(|pdu| pdu.to_state_event())
.collect(),
state,
};
Ok(resp.into())

View file

@ -6,7 +6,11 @@ use ruma::{
},
events::EventType,
};
use std::{collections::BTreeMap, convert::TryInto, sync::Arc};
use std::{
collections::{BTreeMap, HashSet},
convert::TryInto,
sync::Arc,
};
#[cfg(feature = "conduit_bin")]
use rocket::{get, put};
@ -117,6 +121,7 @@ pub async fn get_message_events_route(
body: Ruma<get_message_events::Request<'_>>,
) -> ConduitResult<get_message_events::Response> {
let sender_user = body.sender_user.as_ref().expect("user is authenticated");
let sender_device = body.sender_device.as_ref().expect("user is authenticated");
if !db.rooms.is_joined(sender_user, &body.room_id)? {
return Err(Error::BadRequest(
@ -136,6 +141,12 @@ pub async fn get_message_events_route(
// Use limit or else 10
let limit = body.limit.try_into().map_or(10_usize, |l: u32| l as usize);
let next_token;
let mut resp = get_message_events::Response::new();
let mut lazy_loaded = HashSet::new();
match body.dir {
get_message_events::Direction::Forward => {
let events_after: Vec<_> = db
@ -152,21 +163,27 @@ pub async fn get_message_events_route(
.take_while(|&(k, _)| Some(Ok(k)) != to) // Stop at `to`
.collect();
let end_token = events_after.last().map(|(count, _)| count.to_string());
for (_, event) in &events_after {
if !db.rooms.lazy_load_was_sent_before(
&sender_user,
&sender_device,
&body.room_id,
&event.sender,
)? {
lazy_loaded.insert(event.sender.clone());
}
}
next_token = events_after.last().map(|(count, _)| count).copied();
let events_after: Vec<_> = events_after
.into_iter()
.map(|(_, pdu)| pdu.to_room_event())
.collect();
let resp = get_message_events::Response {
start: body.from.to_owned(),
end: end_token,
chunk: events_after,
state: Vec::new(),
};
Ok(resp.into())
resp.start = body.from.to_owned();
resp.end = next_token.map(|count| count.to_string());
resp.chunk = events_after;
}
get_message_events::Direction::Backward => {
let events_before: Vec<_> = db
@ -183,21 +200,51 @@ pub async fn get_message_events_route(
.take_while(|&(k, _)| Some(Ok(k)) != to) // Stop at `to`
.collect();
let start_token = events_before.last().map(|(count, _)| count.to_string());
for (_, event) in &events_before {
if !db.rooms.lazy_load_was_sent_before(
&sender_user,
&sender_device,
&body.room_id,
&event.sender,
)? {
lazy_loaded.insert(event.sender.clone());
}
}
next_token = events_before.last().map(|(count, _)| count).copied();
let events_before: Vec<_> = events_before
.into_iter()
.map(|(_, pdu)| pdu.to_room_event())
.collect();
let resp = get_message_events::Response {
start: body.from.to_owned(),
end: start_token,
chunk: events_before,
state: Vec::new(),
};
Ok(resp.into())
resp.start = body.from.to_owned();
resp.end = next_token.map(|count| count.to_string());
resp.chunk = events_before;
}
}
db.rooms
.lazy_load_confirm_delivery(&sender_user, &sender_device, &body.room_id, from)?;
resp.state = Vec::new();
for ll_id in &lazy_loaded {
if let Some(member_event) =
db.rooms
.room_state_get(&body.room_id, &EventType::RoomMember, ll_id.as_str())?
{
resp.state.push(member_event.to_state_event());
}
}
if let Some(next_token) = next_token {
db.rooms.lazy_load_mark_sent(
&sender_user,
&sender_device,
&body.room_id,
lazy_loaded.into_iter().collect(),
next_token,
);
}
Ok(resp.into())
}

View file

@ -264,6 +264,14 @@ async fn sync_helper(
// limited unless there are events in non_timeline_pdus
let limited = non_timeline_pdus.next().is_some();
let mut timeline_users = HashSet::new();
for (_, event) in &timeline_pdus {
timeline_users.insert(event.sender.as_str().to_owned());
}
db.rooms
.lazy_load_confirm_delivery(&sender_user, &sender_device, &room_id, since)?;
// Database queries:
let current_shortstatehash = db
@ -344,14 +352,51 @@ async fn sync_helper(
state_events,
) = if since_shortstatehash.is_none() {
// Probably since = 0, we will do an initial sync
let (joined_member_count, invited_member_count, heroes) = calculate_counts()?;
let current_state_ids = db.rooms.state_full_ids(current_shortstatehash)?;
let state_events: Vec<_> = current_state_ids
.iter()
.map(|(_, id)| db.rooms.get_pdu(id))
.filter_map(|r| r.ok().flatten())
.collect();
let mut state_events = Vec::new();
let mut lazy_loaded = Vec::new();
for (_, id) in current_state_ids {
let pdu = match db.rooms.get_pdu(&id)? {
Some(pdu) => pdu,
None => {
error!("Pdu in state not found: {}", id);
continue;
}
};
let state_key = pdu
.state_key
.as_ref()
.expect("state events have state keys");
if pdu.kind != EventType::RoomMember {
state_events.push(pdu);
} else if full_state || timeline_users.contains(state_key) {
// TODO: check filter: is ll enabled?
lazy_loaded.push(
UserId::parse(state_key.as_ref())
.expect("they are in timeline_users, so they should be correct"),
);
state_events.push(pdu);
}
}
// Reset lazy loading because this is an initial sync
db.rooms
.lazy_load_reset(&sender_user, &sender_device, &room_id)?;
// The state_events above should contain all timeline_users, let's mark them as lazy
// loaded.
db.rooms.lazy_load_mark_sent(
&sender_user,
&sender_device,
&room_id,
lazy_loaded,
next_batch,
);
(
heroes,
@ -387,20 +432,66 @@ async fn sync_helper(
let since_state_ids = db.rooms.state_full_ids(since_shortstatehash)?;
let state_events = if joined_since_last_sync {
/*
let state_events = if joined_since_last_sync || full_state {
current_state_ids
.iter()
.map(|(_, id)| db.rooms.get_pdu(id))
.filter_map(|r| r.ok().flatten())
.collect::<Vec<_>>()
} else {
current_state_ids
.iter()
.filter(|(key, id)| since_state_ids.get(key) != Some(id))
.map(|(_, id)| db.rooms.get_pdu(id))
.filter_map(|r| r.ok().flatten())
.collect()
};
*/
let mut state_events = Vec::new();
let mut lazy_loaded = Vec::new();
for (key, id) in current_state_ids {
let pdu = match db.rooms.get_pdu(&id)? {
Some(pdu) => pdu,
None => {
error!("Pdu in state not found: {}", id);
continue;
}
};
let state_key = pdu
.state_key
.as_ref()
.expect("state events have state keys");
if pdu.kind != EventType::RoomMember {
if full_state || since_state_ids.get(&key) != Some(&id) {
state_events.push(pdu);
}
continue;
}
// Pdu has to be a member event
let state_key_userid = UserId::parse(state_key.as_ref())
.expect("they are in timeline_users, so they should be correct");
if full_state || since_state_ids.get(&key) != Some(&id) {
lazy_loaded.push(state_key_userid);
state_events.push(pdu);
} else if timeline_users.contains(state_key)
&& !db.rooms.lazy_load_was_sent_before(
&sender_user,
&sender_device,
&room_id,
&state_key_userid,
)?
{
lazy_loaded.push(state_key_userid);
state_events.push(pdu);
}
}
db.rooms.lazy_load_mark_sent(
&sender_user,
&sender_device,
&room_id,
lazy_loaded,
next_batch,
);
let encrypted_room = db
.rooms