diff --git a/src/service/rooms/lazy_loading/mod.rs b/src/service/rooms/lazy_loading/mod.rs index 71fe82960..60abdbdfe 100644 --- a/src/service/rooms/lazy_loading/mod.rs +++ b/src/service/rooms/lazy_loading/mod.rs @@ -3,7 +3,7 @@ use std::{collections::HashSet, sync::Arc}; use conduwuit::{ - Result, implement, + Result, utils::{IterStream, ReadyExt, stream::TryIgnore}, }; use database::{Database, Deserialized, Handle, Interfix, Map, Qry}; @@ -55,87 +55,87 @@ fn build(args: crate::Args<'_>) -> Result> { fn name(&self) -> &str { crate::service::make_name(std::module_path!()) } } -#[implement(Service)] -#[tracing::instrument(skip(self), level = "debug")] -pub async fn reset(&self, ctx: &Context<'_>) { - let prefix = (ctx.user_id, ctx.device_id, ctx.room_id, Interfix); - self.db - .lazyloadedids - .keys_prefix_raw(&prefix) - .ignore_err() - .ready_for_each(|key| self.db.lazyloadedids.remove(key)) - .await; -} +impl Service { + /// Resets the lazy loading context. + pub async fn reset(&self, ctx: &Context<'_>) { + let prefix = (ctx.user_id, ctx.device_id, ctx.room_id, Interfix); + self.db + .lazyloadedids + .keys_prefix_raw(&prefix) + .ignore_err() + .ready_for_each(|key| self.db.lazyloadedids.remove(key)) + .await; + } -/// Returns only the subset of `senders` which should be sent to the client -/// according to the provided lazy loading context. -#[implement(Service)] -#[tracing::instrument(name = "retain", level = "debug", skip_all)] -pub async fn retain_lazy_members(&self, senders: MemberSet, ctx: &Context<'_>) -> MemberSet { - debug_assert!( - ctx.options.is_none_or(Options::is_enabled), - "lazy loading should be enabled by your options" - ); + /// Returns only the subset of `senders` which should be sent to the client + /// according to the provided lazy loading context. + pub async fn retain_lazy_members(&self, senders: MemberSet, ctx: &Context<'_>) -> MemberSet { + debug_assert!( + ctx.options.is_none_or(Options::is_enabled), + "lazy loading should be enabled by your options" + ); - let include_redundant = cfg!(feature = "element_hacks") - || ctx.options.is_some_and(Options::include_redundant_members); + let include_redundant = cfg!(feature = "element_hacks") + || ctx.options.is_some_and(Options::include_redundant_members); - let witness = self - .witness(ctx, senders.iter().map(AsRef::as_ref)) - .zip(senders.iter().stream()); + let witness = self + .witness(ctx, senders.iter().map(AsRef::as_ref)) + .zip(senders.iter().stream()); - pin_mut!(witness); - let _cork = self.db.db.cork(); - let mut senders = MemberSet::with_capacity(senders.len()); - while let Some((status, sender)) = witness.next().await { - if include_redundant || status == Status::Unseen { - senders.insert(sender.clone()); - continue; - } - - if let Status::Seen(seen) = status { - if seen == 0 || ctx.token == Some(seen) { + pin_mut!(witness); + let _cork = self.db.db.cork(); + let mut senders = MemberSet::with_capacity(senders.len()); + while let Some((status, sender)) = witness.next().await { + if include_redundant || status == Status::Unseen { senders.insert(sender.clone()); continue; } + + if let Status::Seen(seen) = status { + if seen == 0 || ctx.token == Some(seen) { + senders.insert(sender.clone()); + continue; + } + } } + + senders } - senders -} + /// Witnesses a set of users via the lazy loading context, returning a + /// stream of their statuses. + fn witness<'a, I>( + &'a self, + ctx: &'a Context<'a>, + senders: I, + ) -> impl Stream + Send + 'a + where + I: Iterator + Send + Clone + 'a, + { + let make_key = + |sender: &'a UserId| -> Key<'a> { (ctx.user_id, ctx.device_id, ctx.room_id, sender) }; -#[implement(Service)] -fn witness<'a, I>( - &'a self, - ctx: &'a Context<'a>, - senders: I, -) -> impl Stream + Send + 'a -where - I: Iterator + Send + Clone + 'a, -{ - let make_key = - |sender: &'a UserId| -> Key<'a> { (ctx.user_id, ctx.device_id, ctx.room_id, sender) }; + senders + .clone() + .stream() + .map(make_key) + .qry(&self.db.lazyloadedids) + .map(into_status) + .zip(senders.stream()) + .map(move |(status, sender)| { + if matches!(status, Status::Unseen) { + self.db + .lazyloadedids + .put_aput::<8, _, _>(make_key(sender), 0_u64); + } else if matches!(status, Status::Seen(0)) { + self.db + .lazyloadedids + .put_aput::<8, _, _>(make_key(sender), ctx.token.unwrap_or(0_u64)); + } - senders - .clone() - .stream() - .map(make_key) - .qry(&self.db.lazyloadedids) - .map(into_status) - .zip(senders.stream()) - .map(move |(status, sender)| { - if matches!(status, Status::Unseen) { - self.db - .lazyloadedids - .put_aput::<8, _, _>(make_key(sender), 0_u64); - } else if matches!(status, Status::Seen(0)) { - self.db - .lazyloadedids - .put_aput::<8, _, _>(make_key(sender), ctx.token.unwrap_or(0_u64)); - } - - status - }) + status + }) + } } fn into_status(result: Result>) -> Status { @@ -146,6 +146,8 @@ fn into_status(result: Result>) -> Status { } impl Options for LazyLoadOptions { + fn is_enabled(&self) -> bool { !self.is_disabled() } + fn include_redundant_members(&self) -> bool { if let Self::Enabled { include_redundant_members } = self { *include_redundant_members @@ -153,6 +155,4 @@ fn include_redundant_members(&self) -> bool { false } } - - fn is_enabled(&self) -> bool { !self.is_disabled() } } diff --git a/src/service/rooms/metadata/mod.rs b/src/service/rooms/metadata/mod.rs index f0de153f7..b224ac232 100644 --- a/src/service/rooms/metadata/mod.rs +++ b/src/service/rooms/metadata/mod.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use conduwuit::{Result, implement, utils::stream::TryIgnore}; +use conduwuit::{Result, utils::stream::TryIgnore}; use database::Map; use futures::{Stream, StreamExt}; use ruma::{OwnedRoomId, RoomId}; diff --git a/src/service/rooms/outlier/mod.rs b/src/service/rooms/outlier/mod.rs index 55a3ce7b5..afd671999 100644 --- a/src/service/rooms/outlier/mod.rs +++ b/src/service/rooms/outlier/mod.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use conduwuit::{Result, implement, matrix::PduEvent}; +use conduwuit::{Result, matrix::PduEvent}; use database::{Deserialized, Json, Map}; use ruma::{CanonicalJsonObject, EventId};