Split state_cache service.

Signed-off-by: Jason Volk <jason@zemos.net>

src/service/rooms/state_cache/update.rs

@@ -0,0 +1,369 @@
+use std::collections::HashSet;
+
+use conduwuit::{Result, implement, is_not_empty, utils::ReadyExt, warn};
+use database::{Json, serialize_key};
+use futures::StreamExt;
+use ruma::{
+	OwnedServerName, RoomId, UserId,
+	events::{
+		AnyStrippedStateEvent, AnySyncStateEvent, GlobalAccountDataEventType,
+		RoomAccountDataEventType, StateEventType,
+		direct::DirectEvent,
+		room::{
+			create::RoomCreateEventContent,
+			member::{MembershipState, RoomMemberEventContent},
+		},
+	},
+	serde::Raw,
+};
+
+/// Update current membership data.
+#[implement(super::Service)]
+#[tracing::instrument(
+		level = "debug",
+		skip_all,
+		fields(
+			%room_id,
+			%user_id,
+			%sender,
+			?membership_event,
+		),
+	)]
+#[allow(clippy::too_many_arguments)]
+pub async fn update_membership(
+	&self,
+	room_id: &RoomId,
+	user_id: &UserId,
+	membership_event: RoomMemberEventContent,
+	sender: &UserId,
+	last_state: Option<Vec<Raw<AnyStrippedStateEvent>>>,
+	invite_via: Option<Vec<OwnedServerName>>,
+	update_joined_count: bool,
+) -> Result {
+	let membership = membership_event.membership;
+
+	// Keep track what remote users exist by adding them as "deactivated" users
+	//
+	// TODO: use futures to update remote profiles without blocking the membership
+	// update
+	#[allow(clippy::collapsible_if)]
+	if !self.services.globals.user_is_local(user_id) {
+		if !self.services.users.exists(user_id).await {
+			self.services.users.create(user_id, None)?;
+		}
+	}
+
+	match &membership {
+		| MembershipState::Join => {
+			// Check if the user never joined this room
+			if !self.once_joined(user_id, room_id).await {
+				// Add the user ID to the join list then
+				self.mark_as_once_joined(user_id, room_id);
+
+				// Check if the room has a predecessor
+				if let Ok(Some(predecessor)) = self
+					.services
+					.state_accessor
+					.room_state_get_content(room_id, &StateEventType::RoomCreate, "")
+					.await
+					.map(|content: RoomCreateEventContent| content.predecessor)
+				{
+					// Copy old tags to new room
+					if let Ok(tag_event) = self
+						.services
+						.account_data
+						.get_room(&predecessor.room_id, user_id, RoomAccountDataEventType::Tag)
+						.await
+					{
+						self.services
+							.account_data
+							.update(
+								Some(room_id),
+								user_id,
+								RoomAccountDataEventType::Tag,
+								&tag_event,
+							)
+							.await
+							.ok();
+					}
+
+					// Copy direct chat flag
+					if let Ok(mut direct_event) = self
+						.services
+						.account_data
+						.get_global::<DirectEvent>(user_id, GlobalAccountDataEventType::Direct)
+						.await
+					{
+						let mut room_ids_updated = false;
+						for room_ids in direct_event.content.0.values_mut() {
+							if room_ids.iter().any(|r| r == &predecessor.room_id) {
+								room_ids.push(room_id.to_owned());
+								room_ids_updated = true;
+							}
+						}
+
+						if room_ids_updated {
+							self.services
+								.account_data
+								.update(
+									None,
+									user_id,
+									GlobalAccountDataEventType::Direct.to_string().into(),
+									&serde_json::to_value(&direct_event)
+										.expect("to json always works"),
+								)
+								.await?;
+						}
+					}
+				}
+			}
+
+			self.mark_as_joined(user_id, room_id);
+		},
+		| MembershipState::Invite => {
+			// We want to know if the sender is ignored by the receiver
+			if self.services.users.user_is_ignored(sender, user_id).await {
+				return Ok(());
+			}
+
+			self.mark_as_invited(user_id, room_id, last_state, invite_via)
+				.await;
+		},
+		| MembershipState::Leave | MembershipState::Ban => {
+			self.mark_as_left(user_id, room_id);
+
+			if self.services.globals.user_is_local(user_id)
+				&& (self.services.config.forget_forced_upon_leave
+					|| self.services.metadata.is_banned(room_id).await
+					|| self.services.metadata.is_disabled(room_id).await)
+			{
+				self.forget(room_id, user_id);
+			}
+		},
+		| _ => {},
+	}
+
+	if update_joined_count {
+		self.update_joined_count(room_id).await;
+	}
+
+	Ok(())
+}
+
+#[implement(super::Service)]
+#[tracing::instrument(level = "debug", skip(self))]
+pub async fn update_joined_count(&self, room_id: &RoomId) {
+	let mut joinedcount = 0_u64;
+	let mut invitedcount = 0_u64;
+	let mut knockedcount = 0_u64;
+	let mut joined_servers = HashSet::new();
+
+	self.room_members(room_id)
+		.ready_for_each(|joined| {
+			joined_servers.insert(joined.server_name().to_owned());
+			joinedcount = joinedcount.saturating_add(1);
+		})
+		.await;
+
+	invitedcount = invitedcount.saturating_add(
+		self.room_members_invited(room_id)
+			.count()
+			.await
+			.try_into()
+			.unwrap_or(0),
+	);
+
+	knockedcount = knockedcount.saturating_add(
+		self.room_members_knocked(room_id)
+			.count()
+			.await
+			.try_into()
+			.unwrap_or(0),
+	);
+
+	self.db.roomid_joinedcount.raw_put(room_id, joinedcount);
+	self.db.roomid_invitedcount.raw_put(room_id, invitedcount);
+	self.db
+		.roomuserid_knockedcount
+		.raw_put(room_id, knockedcount);
+
+	self.room_servers(room_id)
+		.ready_for_each(|old_joined_server| {
+			if joined_servers.remove(old_joined_server) {
+				return;
+			}
+
+			// Server not in room anymore
+			let roomserver_id = (room_id, old_joined_server);
+			let serverroom_id = (old_joined_server, room_id);
+
+			self.db.roomserverids.del(roomserver_id);
+			self.db.serverroomids.del(serverroom_id);
+		})
+		.await;
+
+	// Now only new servers are in joined_servers anymore
+	for server in &joined_servers {
+		let roomserver_id = (room_id, server);
+		let serverroom_id = (server, room_id);
+
+		self.db.roomserverids.put_raw(roomserver_id, []);
+		self.db.serverroomids.put_raw(serverroom_id, []);
+	}
+
+	self.appservice_in_room_cache
+		.write()
+		.expect("locked")
+		.remove(room_id);
+}
+
+/// Direct DB function to directly mark a user as joined. It is not
+/// recommended to use this directly. You most likely should use
+/// `update_membership` instead
+#[implement(super::Service)]
+#[tracing::instrument(skip(self), level = "debug")]
+pub fn mark_as_joined(&self, user_id: &UserId, room_id: &RoomId) {
+	let userroom_id = (user_id, room_id);
+	let userroom_id = serialize_key(userroom_id).expect("failed to serialize userroom_id");
+
+	let roomuser_id = (room_id, user_id);
+	let roomuser_id = serialize_key(roomuser_id).expect("failed to serialize roomuser_id");
+
+	self.db.userroomid_joined.insert(&userroom_id, []);
+	self.db.roomuserid_joined.insert(&roomuser_id, []);
+
+	self.db.userroomid_invitestate.remove(&userroom_id);
+	self.db.roomuserid_invitecount.remove(&roomuser_id);
+
+	self.db.userroomid_leftstate.remove(&userroom_id);
+	self.db.roomuserid_leftcount.remove(&roomuser_id);
+
+	self.db.userroomid_knockedstate.remove(&userroom_id);
+	self.db.roomuserid_knockedcount.remove(&roomuser_id);
+
+	self.db.roomid_inviteviaservers.remove(room_id);
+}
+
+/// Direct DB function to directly mark a user as left. It is not
+/// recommended to use this directly. You most likely should use
+/// `update_membership` instead
+#[implement(super::Service)]
+#[tracing::instrument(skip(self), level = "debug")]
+pub fn mark_as_left(&self, user_id: &UserId, room_id: &RoomId) {
+	let userroom_id = (user_id, room_id);
+	let userroom_id = serialize_key(userroom_id).expect("failed to serialize userroom_id");
+
+	let roomuser_id = (room_id, user_id);
+	let roomuser_id = serialize_key(roomuser_id).expect("failed to serialize roomuser_id");
+
+	// (timo) TODO
+	let leftstate = Vec::<Raw<AnySyncStateEvent>>::new();
+
+	self.db
+		.userroomid_leftstate
+		.raw_put(&userroom_id, Json(leftstate));
+	self.db
+		.roomuserid_leftcount
+		.raw_aput::<8, _, _>(&roomuser_id, self.services.globals.next_count().unwrap());
+
+	self.db.userroomid_joined.remove(&userroom_id);
+	self.db.roomuserid_joined.remove(&roomuser_id);
+
+	self.db.userroomid_invitestate.remove(&userroom_id);
+	self.db.roomuserid_invitecount.remove(&roomuser_id);
+
+	self.db.userroomid_knockedstate.remove(&userroom_id);
+	self.db.roomuserid_knockedcount.remove(&roomuser_id);
+
+	self.db.roomid_inviteviaservers.remove(room_id);
+}
+
+/// Direct DB function to directly mark a user as knocked. It is not
+/// recommended to use this directly. You most likely should use
+/// `update_membership` instead
+#[implement(super::Service)]
+#[tracing::instrument(skip(self), level = "debug")]
+pub fn mark_as_knocked(
+	&self,
+	user_id: &UserId,
+	room_id: &RoomId,
+	knocked_state: Option<Vec<Raw<AnyStrippedStateEvent>>>,
+) {
+	let userroom_id = (user_id, room_id);
+	let userroom_id = serialize_key(userroom_id).expect("failed to serialize userroom_id");
+
+	let roomuser_id = (room_id, user_id);
+	let roomuser_id = serialize_key(roomuser_id).expect("failed to serialize roomuser_id");
+
+	self.db
+		.userroomid_knockedstate
+		.raw_put(&userroom_id, Json(knocked_state.unwrap_or_default()));
+	self.db
+		.roomuserid_knockedcount
+		.raw_aput::<8, _, _>(&roomuser_id, self.services.globals.next_count().unwrap());
+
+	self.db.userroomid_joined.remove(&userroom_id);
+	self.db.roomuserid_joined.remove(&roomuser_id);
+
+	self.db.userroomid_invitestate.remove(&userroom_id);
+	self.db.roomuserid_invitecount.remove(&roomuser_id);
+
+	self.db.userroomid_leftstate.remove(&userroom_id);
+	self.db.roomuserid_leftcount.remove(&roomuser_id);
+
+	self.db.roomid_inviteviaservers.remove(room_id);
+}
+
+/// Makes a user forget a room.
+#[implement(super::Service)]
+#[tracing::instrument(skip(self), level = "debug")]
+pub fn forget(&self, room_id: &RoomId, user_id: &UserId) {
+	let userroom_id = (user_id, room_id);
+	let roomuser_id = (room_id, user_id);
+
+	self.db.userroomid_leftstate.del(userroom_id);
+	self.db.roomuserid_leftcount.del(roomuser_id);
+}
+
+#[implement(super::Service)]
+#[tracing::instrument(level = "debug", skip(self))]
+fn mark_as_once_joined(&self, user_id: &UserId, room_id: &RoomId) {
+	let key = (user_id, room_id);
+	self.db.roomuseroncejoinedids.put_raw(key, []);
+}
+
+#[implement(super::Service)]
+#[tracing::instrument(level = "debug", skip(self, last_state, invite_via))]
+pub async fn mark_as_invited(
+	&self,
+	user_id: &UserId,
+	room_id: &RoomId,
+	last_state: Option<Vec<Raw<AnyStrippedStateEvent>>>,
+	invite_via: Option<Vec<OwnedServerName>>,
+) {
+	let roomuser_id = (room_id, user_id);
+	let roomuser_id = serialize_key(roomuser_id).expect("failed to serialize roomuser_id");
+
+	let userroom_id = (user_id, room_id);
+	let userroom_id = serialize_key(userroom_id).expect("failed to serialize userroom_id");
+
+	self.db
+		.userroomid_invitestate
+		.raw_put(&userroom_id, Json(last_state.unwrap_or_default()));
+	self.db
+		.roomuserid_invitecount
+		.raw_aput::<8, _, _>(&roomuser_id, self.services.globals.next_count().unwrap());
+
+	self.db.userroomid_joined.remove(&userroom_id);
+	self.db.roomuserid_joined.remove(&roomuser_id);
+
+	self.db.userroomid_leftstate.remove(&userroom_id);
+	self.db.roomuserid_leftcount.remove(&roomuser_id);
+
+	self.db.userroomid_knockedstate.remove(&userroom_id);
+	self.db.roomuserid_knockedcount.remove(&roomuser_id);
+
+	if let Some(servers) = invite_via.filter(is_not_empty!()) {
+		self.add_servers_invite_via(room_id, servers).await;
+	}
+}

src/service/rooms/state_cache/via.rs

@@ -0,0 +1,92 @@
+use conduwuit::{
+	Result, implement,
+	utils::{StreamTools, stream::TryIgnore},
+	warn,
+};
+use database::Ignore;
+use futures::{Stream, StreamExt, stream::iter};
+use itertools::Itertools;
+use ruma::{
+	OwnedServerName, RoomId, ServerName,
+	events::{StateEventType, room::power_levels::RoomPowerLevelsEventContent},
+	int,
+};
+
+#[implement(super::Service)]
+#[tracing::instrument(level = "debug", skip(self, servers))]
+pub async fn add_servers_invite_via(&self, room_id: &RoomId, servers: Vec<OwnedServerName>) {
+	let mut servers: Vec<_> = self
+		.servers_invite_via(room_id)
+		.map(ToOwned::to_owned)
+		.chain(iter(servers.into_iter()))
+		.collect()
+		.await;
+
+	servers.sort_unstable();
+	servers.dedup();
+
+	let servers = servers
+		.iter()
+		.map(|server| server.as_bytes())
+		.collect_vec()
+		.join(&[0xFF][..]);
+
+	self.db
+		.roomid_inviteviaservers
+		.insert(room_id.as_bytes(), &servers);
+}
+
+/// Gets up to five servers that are likely to be in the room in the
+/// distant future.
+///
+/// See <https://spec.matrix.org/latest/appendices/#routing>
+#[implement(super::Service)]
+#[tracing::instrument(skip(self), level = "trace")]
+pub async fn servers_route_via(&self, room_id: &RoomId) -> Result<Vec<OwnedServerName>> {
+	let most_powerful_user_server = self
+		.services
+		.state_accessor
+		.room_state_get_content(room_id, &StateEventType::RoomPowerLevels, "")
+		.await
+		.map(|content: RoomPowerLevelsEventContent| {
+			content
+				.users
+				.iter()
+				.max_by_key(|(_, power)| *power)
+				.and_then(|x| (x.1 >= &int!(50)).then_some(x))
+				.map(|(user, _power)| user.server_name().to_owned())
+		});
+
+	let mut servers: Vec<OwnedServerName> = self
+		.room_members(room_id)
+		.counts_by(|user| user.server_name().to_owned())
+		.await
+		.into_iter()
+		.sorted_by_key(|(_, users)| *users)
+		.map(|(server, _)| server)
+		.rev()
+		.take(5)
+		.collect();
+
+	if let Ok(Some(server)) = most_powerful_user_server {
+		servers.insert(0, server);
+		servers.truncate(5);
+	}
+
+	Ok(servers)
+}
+
+#[implement(super::Service)]
+#[tracing::instrument(skip(self), level = "debug")]
+pub fn servers_invite_via<'a>(
+	&'a self,
+	room_id: &'a RoomId,
+) -> impl Stream<Item = &ServerName> + Send + 'a {
+	type KeyVal<'a> = (Ignore, Vec<&'a ServerName>);
+
+	self.db
+		.roomid_inviteviaservers
+		.stream_raw_prefix(room_id)
+		.ignore_err()
+		.map(|(_, servers): KeyVal<'_>| *servers.last().expect("at least one server"))
+}