refactor to iterator inputs for auth_chain/short batch functions

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

src/admin/debug/commands.rs

@@ -1,6 +1,7 @@
 use std::{
 	collections::HashMap,
 	fmt::Write,
+	iter::once,
 	sync::Arc,
 	time::{Instant, SystemTime},
 };
@@ -43,7 +44,7 @@ pub(super) async fn get_auth_chain(&self, event_id: Box<EventId>) -> Result<Room
 		.services
 		.rooms
 		.auth_chain
-		.event_ids_iter(room_id, &[&event_id])
+		.event_ids_iter(room_id, once(event_id.as_ref()))
 		.await?
 		.count()
 		.await;

src/api/server/event_auth.rs

@@ -1,4 +1,4 @@
-use std::borrow::Borrow;
+use std::{borrow::Borrow, iter::once};
 
 use axum::extract::State;
 use conduit::{Error, Result};
@@ -46,7 +46,7 @@ pub(crate) async fn get_event_authorization_route(
 	let auth_chain = services
 		.rooms
 		.auth_chain
-		.event_ids_iter(room_id, &[body.event_id.borrow()])
+		.event_ids_iter(room_id, once(body.event_id.borrow()))
 		.await?
 		.filter_map(|id| async move { services.rooms.timeline.get_pdu_json(&id).await.ok() })
 		.then(|pdu| services.sending.convert_to_outgoing_federation_event(pdu))

src/api/server/send_join.rs

@@ -11,7 +11,7 @@ use ruma::{
 		room::member::{MembershipState, RoomMemberEventContent},
 		StateEventType,
 	},
-	CanonicalJsonValue, EventId, OwnedServerName, OwnedUserId, RoomId, ServerName,
+	CanonicalJsonValue, OwnedServerName, OwnedUserId, RoomId, ServerName,
 };
 use serde_json::value::{to_raw_value, RawValue as RawJsonValue};
 use service::Services;
@@ -184,11 +184,11 @@ async fn create_join_event(
 		.try_collect()
 		.await?;
 
-	let starting_events: Vec<&EventId> = state_ids.values().map(Borrow::borrow).collect();
+	let starting_events = state_ids.values().map(Borrow::borrow);
 	let auth_chain = services
 		.rooms
 		.auth_chain
-		.event_ids_iter(room_id, &starting_events)
+		.event_ids_iter(room_id, starting_events)
 		.await?
 		.map(Ok)
 		.and_then(|event_id| async move { services.rooms.timeline.get_pdu_json(&event_id).await })

src/api/server/state.rs

@@ -1,4 +1,4 @@
-use std::borrow::Borrow;
+use std::{borrow::Borrow, iter::once};
 
 use axum::extract::State;
 use conduit::{err, result::LogErr, utils::IterStream, Result};
@@ -52,7 +52,7 @@ pub(crate) async fn get_room_state_route(
 	let auth_chain = services
 		.rooms
 		.auth_chain
-		.event_ids_iter(&body.room_id, &[body.event_id.borrow()])
+		.event_ids_iter(&body.room_id, once(body.event_id.borrow()))
 		.await?
 		.map(Ok)
 		.and_then(|id| async move { services.rooms.timeline.get_pdu_json(&id).await })

src/api/server/state_ids.rs

@@ -1,4 +1,4 @@
-use std::borrow::Borrow;
+use std::{borrow::Borrow, iter::once};
 
 use axum::extract::State;
 use conduit::{err, Result};
@@ -44,7 +44,7 @@ pub(crate) async fn get_room_state_ids_route(
 	let auth_chain_ids = services
 		.rooms
 		.auth_chain
-		.event_ids_iter(&body.room_id, &[body.event_id.borrow()])
+		.event_ids_iter(&body.room_id, once(body.event_id.borrow()))
 		.await?
 		.map(|id| (*id).to_owned())
 		.collect()

src/database/map/get.rs

@@ -80,8 +80,8 @@ where
 #[tracing::instrument(skip(self, keys), fields(%self), level = "trace")]
 pub fn get_batch<'a, I, K>(&self, keys: I) -> impl Stream<Item = Result<Handle<'_>>>
 where
-	I: Iterator<Item = &'a K> + ExactSizeIterator + Send + Debug,
+	I: Iterator<Item = &'a K> + ExactSizeIterator + Debug + Send,
-	K: AsRef<[u8]> + Send + Sync + Sized + Debug + 'a,
+	K: AsRef<[u8]> + Debug + Send + ?Sized + Sync + 'a,
 {
 	self.get_batch_blocking(keys).stream()
 }
@@ -89,8 +89,8 @@ where
 #[implement(super::Map)]
 pub fn get_batch_blocking<'a, I, K>(&self, keys: I) -> impl Iterator<Item = Result<Handle<'_>>>
 where
-	I: Iterator<Item = &'a K> + ExactSizeIterator + Send,
+	I: Iterator<Item = &'a K> + ExactSizeIterator + Debug + Send,
-	K: AsRef<[u8]> + Sized + 'a,
+	K: AsRef<[u8]> + Debug + Send + ?Sized + Sync + 'a,
 {
 	// Optimization can be `true` if key vector is pre-sorted **by the column
 	// comparator**.

src/service/rooms/auth_chain/mod.rs

@@ -2,6 +2,7 @@ mod data;
 
 use std::{
 	collections::{BTreeSet, HashSet},
+	fmt::Debug,
 	sync::Arc,
 };
 
@@ -37,9 +38,12 @@ impl crate::Service for Service {
 }
 
 impl Service {
-	pub async fn event_ids_iter(
+	pub async fn event_ids_iter<'a, I>(
-		&self, room_id: &RoomId, starting_events: &[&EventId],
+		&'a self, room_id: &RoomId, starting_events: I,
-	) -> Result<impl Stream<Item = Arc<EventId>> + Send + '_> {
+	) -> Result<impl Stream<Item = Arc<EventId>> + Send + '_>
+	where
+		I: Iterator<Item = &'a EventId> + Clone + Debug + ExactSizeIterator + Send + 'a,
+	{
 		let stream = self
 			.get_event_ids(room_id, starting_events)
 			.await?
@@ -49,12 +53,15 @@ impl Service {
 		Ok(stream)
 	}
 
-	pub async fn get_event_ids(&self, room_id: &RoomId, starting_events: &[&EventId]) -> Result<Vec<Arc<EventId>>> {
+	pub async fn get_event_ids<'a, I>(&'a self, room_id: &RoomId, starting_events: I) -> Result<Vec<Arc<EventId>>>
+	where
+		I: Iterator<Item = &'a EventId> + Clone + Debug + ExactSizeIterator + Send + 'a,
+	{
 		let chain = self.get_auth_chain(room_id, starting_events).await?;
 		let event_ids = self
 			.services
 			.short
-			.multi_get_eventid_from_short(&chain)
+			.multi_get_eventid_from_short(chain.into_iter())
 			.await
 			.into_iter()
 			.filter_map(Result::ok)
@@ -64,7 +71,10 @@ impl Service {
 	}
 
 	#[tracing::instrument(skip_all, name = "auth_chain")]
-	pub async fn get_auth_chain(&self, room_id: &RoomId, starting_events: &[&EventId]) -> Result<Vec<ShortEventId>> {
+	pub async fn get_auth_chain<'a, I>(&'a self, room_id: &RoomId, starting_events: I) -> Result<Vec<ShortEventId>>
+	where
+		I: Iterator<Item = &'a EventId> + Clone + Debug + ExactSizeIterator + Send + 'a,
+	{
 		const NUM_BUCKETS: usize = 50; //TODO: change possible w/o disrupting db?
 		const BUCKET: BTreeSet<(u64, &EventId)> = BTreeSet::new();
 
@@ -72,19 +82,19 @@ impl Service {
 		let mut starting_ids = self
 			.services
 			.short
-			.multi_get_or_create_shorteventid(starting_events)
+			.multi_get_or_create_shorteventid(starting_events.clone())
-			.enumerate()
+			.zip(starting_events.clone().stream())
 			.boxed();
 
 		let mut buckets = [BUCKET; NUM_BUCKETS];
-		while let Some((i, short)) = starting_ids.next().await {
+		while let Some((short, starting_event)) = starting_ids.next().await {
 			let bucket: usize = short.try_into()?;
 			let bucket: usize = validated!(bucket % NUM_BUCKETS);
-			buckets[bucket].insert((short, starting_events[i]));
+			buckets[bucket].insert((short, starting_event));
 		}
 
 		debug!(
-			starting_events = ?starting_events.len(),
+			starting_events = ?starting_events.count(),
 			elapsed = ?started.elapsed(),
 			"start",
 		);

src/service/rooms/event_handler/resolve_state.rs

@@ -35,12 +35,12 @@ pub async fn resolve_state(
 	let fork_states = [current_state_ids, incoming_state];
 	let mut auth_chain_sets = Vec::with_capacity(fork_states.len());
 	for state in &fork_states {
-		let starting_events: Vec<&EventId> = state.values().map(Borrow::borrow).collect();
+		let starting_events = state.values().map(Borrow::borrow);
 
 		let auth_chain: HashSet<Arc<EventId>> = self
 			.services
 			.auth_chain
-			.get_event_ids(room_id, &starting_events)
+			.get_event_ids(room_id, starting_events)
 			.await?
 			.into_iter()
 			.collect();

src/service/rooms/event_handler/state_at_incoming.rs

@@ -139,7 +139,7 @@ pub(super) async fn state_at_incoming_resolved(
 		let auth_chain: HashSet<Arc<EventId>> = self
 			.services
 			.auth_chain
-			.get_event_ids(room_id, &starting_events)
+			.get_event_ids(room_id, starting_events.into_iter())
 			.await?
 			.into_iter()
 			.collect();

src/service/rooms/short/mod.rs

@@ -1,7 +1,7 @@
-use std::{mem::size_of_val, sync::Arc};
+use std::{fmt::Debug, mem::size_of_val, sync::Arc};
 
 pub use conduit::pdu::{ShortEventId, ShortId, ShortRoomId};
-use conduit::{err, implement, utils, Result};
+use conduit::{err, implement, utils, utils::stream::ReadyExt, Result};
 use database::{Deserialized, Map};
 use futures::{Stream, StreamExt};
 use ruma::{events::StateEventType, EventId, RoomId};
@@ -51,52 +51,46 @@ impl crate::Service for Service {
 
 #[implement(Service)]
 pub async fn get_or_create_shorteventid(&self, event_id: &EventId) -> ShortEventId {
-	const BUFSIZE: usize = size_of::<ShortEventId>();
-
 	if let Ok(shorteventid) = self.get_shorteventid(event_id).await {
 		return shorteventid;
 	}
 
-	let shorteventid = self.services.globals.next_count().unwrap();
+	self.create_shorteventid(event_id)
-	debug_assert!(size_of_val(&shorteventid) == BUFSIZE, "buffer requirement changed");
-
-	self.db
-		.eventid_shorteventid
-		.raw_aput::<BUFSIZE, _, _>(event_id, shorteventid);
-
-	self.db
-		.shorteventid_eventid
-		.aput_raw::<BUFSIZE, _, _>(shorteventid, event_id);
-
-	shorteventid
 }
 
 #[implement(Service)]
-pub fn multi_get_or_create_shorteventid<'a>(
+pub fn multi_get_or_create_shorteventid<'a, I>(&'a self, event_ids: I) -> impl Stream<Item = ShortEventId> + Send + '_
-	&'a self, event_ids: &'a [&EventId],
+where
-) -> impl Stream<Item = ShortEventId> + Send + 'a {
+	I: Iterator<Item = &'a EventId> + Clone + Debug + ExactSizeIterator + Send + 'a,
+	<I as Iterator>::Item: AsRef<[u8]> + Send + Sync + 'a,
+{
 	self.db
 		.eventid_shorteventid
-		.get_batch(event_ids.iter())
+		.get_batch(event_ids.clone())
-		.enumerate()
+		.ready_scan(event_ids, |event_ids, result| {
-		.map(|(i, result)| match result {
+			event_ids.next().map(|event_id| (event_id, result))
-			Ok(ref short) => utils::u64_from_u8(short),
-			Err(_) => {
-				const BUFSIZE: usize = size_of::<ShortEventId>();
-
-				let short = self.services.globals.next_count().unwrap();
-				debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed");
-
-				self.db
-					.eventid_shorteventid
-					.raw_aput::<BUFSIZE, _, _>(event_ids[i], short);
-				self.db
-					.shorteventid_eventid
-					.aput_raw::<BUFSIZE, _, _>(short, event_ids[i]);
-
-				short
-			},
 		})
+		.map(|(event_id, result)| match result {
+			Ok(ref short) => utils::u64_from_u8(short),
+			Err(_) => self.create_shorteventid(event_id),
+		})
+}
+
+#[implement(Service)]
+fn create_shorteventid(&self, event_id: &EventId) -> ShortEventId {
+	const BUFSIZE: usize = size_of::<ShortEventId>();
+
+	let short = self.services.globals.next_count().unwrap();
+	debug_assert!(size_of_val(&short) == BUFSIZE, "buffer requirement changed");
+
+	self.db
+		.eventid_shorteventid
+		.raw_aput::<BUFSIZE, _, _>(event_id, short);
+	self.db
+		.shorteventid_eventid
+		.aput_raw::<BUFSIZE, _, _>(short, event_id);
+
+	short
 }
 
 #[implement(Service)]
@@ -154,13 +148,13 @@ pub async fn get_eventid_from_short(&self, shorteventid: ShortEventId) -> Result
 }
 
 #[implement(Service)]
-pub async fn multi_get_eventid_from_short(&self, shorteventid: &[ShortEventId]) -> Vec<Result<Arc<EventId>>> {
+pub async fn multi_get_eventid_from_short<I>(&self, shorteventid: I) -> Vec<Result<Arc<EventId>>>
+where
+	I: Iterator<Item = ShortEventId> + Send,
+{
 	const BUFSIZE: usize = size_of::<ShortEventId>();
 
-	let keys: Vec<[u8; BUFSIZE]> = shorteventid
+	let keys: Vec<[u8; BUFSIZE]> = shorteventid.map(u64::to_be_bytes).collect();
-		.iter()
-		.map(|short| short.to_be_bytes())
-		.collect();
 
 	self.db
 		.shorteventid_eventid