tuwunel_service/rooms/lazy_loading/
mod.rs

1//! Lazy Loading
2
3use std::{collections::HashSet, sync::Arc};
4
5use futures::{Stream, StreamExt, pin_mut};
6use ruma::{DeviceId, OwnedUserId, RoomId, UserId, api::client::filter::LazyLoadOptions};
7use tuwunel_core::{
8	Result, implement,
9	utils::{IterStream, ReadyExt, stream::TryIgnore},
10};
11use tuwunel_database::{Database, Deserialized, Handle, Interfix, Map, Qry};
12
13pub struct Service {
14	db: Data,
15}
16
17struct Data {
18	lazyloadedids: Arc<Map>,
19	db: Arc<Database>,
20}
21
22pub trait Options: Send + Sync {
23	fn is_enabled(&self) -> bool;
24	fn include_redundant_members(&self) -> bool;
25}
26
27#[derive(Clone, Debug)]
28pub struct Context<'a> {
29	pub user_id: &'a UserId,
30	pub device_id: Option<&'a DeviceId>,
31	pub room_id: &'a RoomId,
32	pub token: Option<u64>,
33	pub options: Option<&'a LazyLoadOptions>,
34	pub mode: Mode,
35}
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
38pub enum Mode {
39	Read,
40	Update,
41	Prefetch,
42}
43
44#[derive(Clone, Copy, Debug, Eq, PartialEq)]
45pub enum Status {
46	Unseen,
47	Seen(u64),
48}
49
50pub type Witness = HashSet<OwnedUserId>;
51type Key<'a> = (&'a UserId, Option<&'a DeviceId>, &'a RoomId, &'a UserId);
52
53impl crate::Service for Service {
54	fn build(args: &crate::Args<'_>) -> Result<Arc<Self>> {
55		Ok(Arc::new(Self {
56			db: Data {
57				lazyloadedids: args.db["lazyloadedids"].clone(),
58				db: args.db.clone(),
59			},
60		}))
61	}
62
63	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
64}
65
66#[implement(Service)]
67#[tracing::instrument(skip(self), level = "debug")]
68pub async fn reset(&self, ctx: &Context<'_>) {
69	let prefix = (ctx.user_id, ctx.device_id, ctx.room_id, Interfix);
70	self.db
71		.lazyloadedids
72		.keys_prefix_raw(&prefix)
73		.ignore_err()
74		.ready_for_each(|key| self.db.lazyloadedids.remove(key))
75		.await;
76}
77
78#[implement(Service)]
79#[tracing::instrument(name = "retain", level = "debug", skip_all)]
80pub async fn witness_retain(&self, senders: Witness, ctx: &Context<'_>) -> Witness {
81	debug_assert!(
82		ctx.options.is_none_or(Options::is_enabled),
83		"lazy loading should be enabled by your options"
84	);
85
86	let include_redundant = cfg!(feature = "element_hacks")
87		|| ctx
88			.options
89			.is_some_and(Options::include_redundant_members);
90
91	let witness = self
92		.witness(ctx, senders.iter().map(AsRef::as_ref))
93		.zip(senders.iter().stream());
94
95	pin_mut!(witness);
96	let _cork = self.db.db.cork();
97	let mut senders = Witness::with_capacity(senders.len());
98	while let Some((status, sender)) = witness.next().await {
99		if ctx.mode == Mode::Prefetch {
100			continue;
101		}
102
103		if include_redundant || status == Status::Unseen {
104			senders.insert(sender.into());
105			continue;
106		}
107
108		if let Status::Seen(seen) = status
109			&& (seen == 0 || ctx.token == Some(seen))
110		{
111			senders.insert(sender.into());
112		}
113	}
114
115	senders
116}
117
118#[implement(Service)]
119fn witness<'a, I>(
120	&'a self,
121	ctx: &'a Context<'a>,
122	senders: I,
123) -> impl Stream<Item = Status> + Send + 'a
124where
125	I: Iterator<Item = &'a UserId> + Send + Clone + 'a,
126{
127	senders
128		.clone()
129		.stream()
130		.map(|sender| make_key(ctx, sender))
131		.qry(&self.db.lazyloadedids)
132		.map(into_status)
133		.zip(senders.stream())
134		.map(move |(status, sender)| {
135			if matches!(ctx.mode, Mode::Update) {
136				self.update(ctx, &status, sender);
137			}
138
139			status
140		})
141}
142
143#[implement(Service)]
144fn update(&self, ctx: &Context<'_>, status: &Status, sender: &UserId) {
145	if matches!(status, Status::Unseen) {
146		self.db
147			.lazyloadedids
148			.put_aput::<8, _, _>(make_key(ctx, sender), 0_u64);
149	} else if matches!(status, Status::Seen(0)) {
150		self.db
151			.lazyloadedids
152			.put_aput::<8, _, _>(make_key(ctx, sender), ctx.token.unwrap_or(0_u64));
153	}
154}
155
156fn into_status(result: Result<Handle<'_>>) -> Status {
157	match result.and_then(|handle| handle.deserialized()) {
158		| Ok(seen) => Status::Seen(seen),
159		| Err(_) => Status::Unseen,
160	}
161}
162
163fn make_key<'a>(ctx: &'a Context<'a>, sender: &'a UserId) -> Key<'a> {
164	(ctx.user_id, ctx.device_id, ctx.room_id, sender)
165}
166
167impl Options for LazyLoadOptions {
168	fn include_redundant_members(&self) -> bool {
169		if let Self::Enabled { include_redundant_members } = self {
170			*include_redundant_members
171		} else {
172			false
173		}
174	}
175
176	fn is_enabled(&self) -> bool { !self.is_disabled() }
177}
tuwunel_service/rooms/lazy_loading/mod.rs

tuwunel_service/rooms/lazy_loading/
mod.rs
