split event_handler service

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

src/service/rooms/event_handler/state_at_incoming.rs

@@ -0,0 +1,178 @@
+use std::{
+	borrow::Borrow,
+	collections::{HashMap, HashSet},
+	sync::Arc,
+};
+
+use conduit::{debug, err, implement, result::LogErr, utils::IterStream, PduEvent, Result};
+use futures::{FutureExt, StreamExt};
+use ruma::{
+	state_res::{self, StateMap},
+	EventId, RoomId, RoomVersionId,
+};
+
+// TODO: if we know the prev_events of the incoming event we can avoid the
+#[implement(super::Service)]
+// request and build the state from a known point and resolve if > 1 prev_event
+#[tracing::instrument(skip_all, name = "state")]
+pub(super) async fn state_at_incoming_degree_one(
+	&self, incoming_pdu: &Arc<PduEvent>,
+) -> Result<Option<HashMap<u64, Arc<EventId>>>> {
+	let prev_event = &*incoming_pdu.prev_events[0];
+	let Ok(prev_event_sstatehash) = self
+		.services
+		.state_accessor
+		.pdu_shortstatehash(prev_event)
+		.await
+	else {
+		return Ok(None);
+	};
+
+	let Ok(mut state) = self
+		.services
+		.state_accessor
+		.state_full_ids(prev_event_sstatehash)
+		.await
+		.log_err()
+	else {
+		return Ok(None);
+	};
+
+	debug!("Using cached state");
+	let prev_pdu = self
+		.services
+		.timeline
+		.get_pdu(prev_event)
+		.await
+		.map_err(|e| err!(Database("Could not find prev event, but we know the state: {e:?}")))?;
+
+	if let Some(state_key) = &prev_pdu.state_key {
+		let shortstatekey = self
+			.services
+			.short
+			.get_or_create_shortstatekey(&prev_pdu.kind.to_string().into(), state_key)
+			.await;
+
+		state.insert(shortstatekey, Arc::from(prev_event));
+		// Now it's the state after the pdu
+	}
+
+	debug_assert!(!state.is_empty(), "should be returning None for empty HashMap result");
+
+	Ok(Some(state))
+}
+
+#[implement(super::Service)]
+#[tracing::instrument(skip_all, name = "state")]
+pub(super) async fn state_at_incoming_resolved(
+	&self, incoming_pdu: &Arc<PduEvent>, room_id: &RoomId, room_version_id: &RoomVersionId,
+) -> Result<Option<HashMap<u64, Arc<EventId>>>> {
+	debug!("Calculating state at event using state res");
+	let mut extremity_sstatehashes = HashMap::with_capacity(incoming_pdu.prev_events.len());
+
+	let mut okay = true;
+	for prev_eventid in &incoming_pdu.prev_events {
+		let Ok(prev_event) = self.services.timeline.get_pdu(prev_eventid).await else {
+			okay = false;
+			break;
+		};
+
+		let Ok(sstatehash) = self
+			.services
+			.state_accessor
+			.pdu_shortstatehash(prev_eventid)
+			.await
+		else {
+			okay = false;
+			break;
+		};
+
+		extremity_sstatehashes.insert(sstatehash, prev_event);
+	}
+
+	if !okay {
+		return Ok(None);
+	}
+
+	let mut fork_states = Vec::with_capacity(extremity_sstatehashes.len());
+	let mut auth_chain_sets = Vec::with_capacity(extremity_sstatehashes.len());
+	for (sstatehash, prev_event) in extremity_sstatehashes {
+		let Ok(mut leaf_state) = self
+			.services
+			.state_accessor
+			.state_full_ids(sstatehash)
+			.await
+		else {
+			continue;
+		};
+
+		if let Some(state_key) = &prev_event.state_key {
+			let shortstatekey = self
+				.services
+				.short
+				.get_or_create_shortstatekey(&prev_event.kind.to_string().into(), state_key)
+				.await;
+
+			let event_id = &prev_event.event_id;
+			leaf_state.insert(shortstatekey, event_id.clone());
+			// Now it's the state after the pdu
+		}
+
+		let mut state = StateMap::with_capacity(leaf_state.len());
+		let mut starting_events = Vec::with_capacity(leaf_state.len());
+		for (k, id) in &leaf_state {
+			if let Ok((ty, st_key)) = self
+				.services
+				.short
+				.get_statekey_from_short(*k)
+				.await
+				.log_err()
+			{
+				// FIXME: Undo .to_string().into() when StateMap
+				//        is updated to use StateEventType
+				state.insert((ty.to_string().into(), st_key), id.clone());
+			}
+
+			starting_events.push(id.borrow());
+		}
+
+		let auth_chain: HashSet<Arc<EventId>> = self
+			.services
+			.auth_chain
+			.get_event_ids(room_id, &starting_events)
+			.await?
+			.into_iter()
+			.collect();
+
+		auth_chain_sets.push(auth_chain);
+		fork_states.push(state);
+	}
+
+	let lock = self.services.globals.stateres_mutex.lock();
+
+	let event_fetch = |event_id| self.event_fetch(event_id);
+	let event_exists = |event_id| self.event_exists(event_id);
+	let result = state_res::resolve(room_version_id, &fork_states, &auth_chain_sets, &event_fetch, &event_exists)
+		.await
+		.map_err(|e| err!(Database(warn!(?e, "State resolution on prev events failed."))));
+
+	drop(lock);
+
+	let Ok(new_state) = result else {
+		return Ok(None);
+	};
+
+	new_state
+		.iter()
+		.stream()
+		.then(|((event_type, state_key), event_id)| {
+			self.services
+				.short
+				.get_or_create_shortstatekey(event_type, state_key)
+				.map(move |shortstatekey| (shortstatekey, event_id.clone()))
+		})
+		.collect()
+		.map(Some)
+		.map(Ok)
+		.await
+}