flatten timeline pdus iterations; increase concurrency

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

src/api/client/context.rs

@@ -1,14 +1,12 @@
-use std::iter::once;
-
 use axum::extract::State;
 use conduwuit::{
 	at, err, ref_at,
 	utils::{
 		future::TryExtExt,
-		stream::{BroadbandExt, ReadyExt, WidebandExt},
+		stream::{BroadbandExt, ReadyExt, TryIgnore, WidebandExt},
 		IterStream,
 	},
-	Err, Result,
+	Err, PduEvent, Result,
 };
 use futures::{join, try_join, FutureExt, StreamExt, TryFutureExt};
 use ruma::{
@@ -59,13 +57,13 @@ pub(crate) async fn get_context_route(
 		false
 	};
 
-	let base_token = services
+	let base_id = services
 		.rooms
 		.timeline
-		.get_pdu_count(&body.event_id)
+		.get_pdu_id(&body.event_id)
 		.map_err(|_| err!(Request(NotFound("Event not found."))));
 
-	let base_event = services
+	let base_pdu = services
 		.rooms
 		.timeline
 		.get_pdu(&body.event_id)
@@ -77,48 +75,44 @@ pub(crate) async fn get_context_route(
 		.user_can_see_event(sender_user, &body.room_id, &body.event_id)
 		.map(Ok);
 
-	let (base_token, base_event, visible) = try_join!(base_token, base_event, visible)?;
+	let (base_id, base_pdu, visible) = try_join!(base_id, base_pdu, visible)?;
 
-	if base_event.room_id != body.room_id || base_event.event_id != body.event_id {
+	if base_pdu.room_id != body.room_id || base_pdu.event_id != body.event_id {
 		return Err!(Request(NotFound("Base event not found.")));
 	}
 
-	if !visible
-		|| ignored_filter(&services, (base_token, base_event.clone()), sender_user)
-			.await
-			.is_none()
-	{
+	if !visible {
 		return Err!(Request(Forbidden("You don't have permission to view this event.")));
 	}
 
-	let events_before =
-		services
-			.rooms
-			.timeline
-			.pdus_rev(Some(sender_user), room_id, Some(base_token));
+	let base_count = base_id.pdu_count();
+
+	let base_event = ignored_filter(&services, (base_count, base_pdu), sender_user);
+
+	let events_before = services
+		.rooms
+		.timeline
+		.pdus_rev(Some(sender_user), room_id, Some(base_count))
+		.ignore_err()
+		.ready_filter_map(|item| event_filter(item, filter))
+		.wide_filter_map(|item| ignored_filter(&services, item, sender_user))
+		.wide_filter_map(|item| visibility_filter(&services, item, sender_user))
+		.take(limit / 2)
+		.collect();
 
 	let events_after = services
 		.rooms
 		.timeline
-		.pdus(Some(sender_user), room_id, Some(base_token));
-
-	let (events_before, events_after) = try_join!(events_before, events_after)?;
-
-	let events_before = events_before
+		.pdus(Some(sender_user), room_id, Some(base_count))
+		.ignore_err()
 		.ready_filter_map(|item| event_filter(item, filter))
 		.wide_filter_map(|item| ignored_filter(&services, item, sender_user))
 		.wide_filter_map(|item| visibility_filter(&services, item, sender_user))
 		.take(limit / 2)
 		.collect();
 
-	let events_after = events_after
-		.ready_filter_map(|item| event_filter(item, filter))
-		.wide_filter_map(|item| ignored_filter(&services, item, sender_user))
-		.wide_filter_map(|item| visibility_filter(&services, item, sender_user))
-		.take(limit / 2)
-		.collect();
-
-	let (events_before, events_after): (Vec<_>, Vec<_>) = join!(events_before, events_after);
+	let (base_event, events_before, events_after): (_, Vec<_>, Vec<_>) =
+		join!(base_event, events_before, events_after);
 
 	let state_at = events_after
 		.last()
@@ -134,7 +128,8 @@ pub(crate) async fn get_context_route(
 		.map_err(|e| err!(Database("State not found: {e}")))
 		.await?;
 
-	let lazy = once(&(base_token, base_event.clone()))
+	let lazy = base_event
+		.iter()
 		.chain(events_before.iter())
 		.chain(events_after.iter())
 		.stream()
@@ -175,19 +170,19 @@ pub(crate) async fn get_context_route(
 		.await;
 
 	Ok(get_context::v3::Response {
-		event: Some(base_event.to_room_event()),
+		event: base_event.map(at!(1)).as_ref().map(PduEvent::to_room_event),
 
 		start: events_before
 			.last()
 			.map(at!(0))
-			.or(Some(base_token))
+			.or(Some(base_count))
 			.as_ref()
 			.map(ToString::to_string),
 
 		end: events_after
 			.last()
 			.map(at!(0))
-			.or(Some(base_token))
+			.or(Some(base_count))
 			.as_ref()
 			.map(ToString::to_string),
 

src/api/client/membership.rs

@@ -1314,6 +1314,7 @@ async fn join_room_by_id_helper_local(
 			.rooms
 			.event_handler
 			.handle_incoming_pdu(&remote_server, room_id, &signed_event_id, signed_value, true)
+			.boxed()
 			.await?;
 	} else {
 		return Err(error);
@@ -1491,6 +1492,7 @@ pub(crate) async fn invite_helper(
 			.rooms
 			.event_handler
 			.handle_incoming_pdu(&origin, room_id, &event_id, value, true)
+			.boxed()
 			.await?
 			.ok_or_else(|| {
 				err!(Request(InvalidParam("Could not accept incoming PDU as timeline event.")))

src/api/client/message.rs

@@ -5,7 +5,7 @@ use conduwuit::{
 	at, is_equal_to,
 	utils::{
 		result::{FlatOk, LogErr},
-		stream::{BroadbandExt, WidebandExt},
+		stream::{BroadbandExt, TryIgnore, WidebandExt},
 		IterStream, ReadyExt,
 	},
 	Event, PduCount, Result,
@@ -107,14 +107,14 @@ pub(crate) async fn get_message_events_route(
 			.rooms
 			.timeline
 			.pdus(Some(sender_user), room_id, Some(from))
-			.await?
+			.ignore_err()
 			.boxed(),
 
 		| Direction::Backward => services
 			.rooms
 			.timeline
 			.pdus_rev(Some(sender_user), room_id, Some(from))
-			.await?
+			.ignore_err()
 			.boxed(),
 	};
 

src/api/client/room/initial_sync.rs

@@ -1,6 +1,10 @@
 use axum::extract::State;
-use conduwuit::{at, utils::BoolExt, Err, Result};
-use futures::StreamExt;
+use conduwuit::{
+	at,
+	utils::{stream::TryTools, BoolExt},
+	Err, Result,
+};
+use futures::TryStreamExt;
 use ruma::api::client::room::initial_sync::v3::{PaginationChunk, Request, Response};
 
 use crate::Ruma;
@@ -27,10 +31,9 @@ pub(crate) async fn room_initial_sync_route(
 		.rooms
 		.timeline
 		.pdus_rev(None, room_id, None)
-		.await?
-		.take(limit)
-		.collect()
-		.await;
+		.try_take(limit)
+		.try_collect()
+		.await?;
 
 	let state: Vec<_> = services
 		.rooms

src/api/client/sync/mod.rs

@@ -2,10 +2,10 @@ mod v3;
 mod v4;
 
 use conduwuit::{
-	utils::stream::{BroadbandExt, ReadyExt},
+	utils::stream::{BroadbandExt, ReadyExt, TryIgnore},
 	PduCount,
 };
-use futures::StreamExt;
+use futures::{pin_mut, StreamExt};
 use ruma::{RoomId, UserId};
 
 pub(crate) use self::{v3::sync_events_route, v4::sync_events_v4_route};
@@ -29,23 +29,19 @@ async fn load_timeline(
 		return Ok((Vec::new(), false));
 	}
 
-	let mut non_timeline_pdus = services
+	let non_timeline_pdus = services
 		.rooms
 		.timeline
 		.pdus_rev(Some(sender_user), room_id, None)
-		.await?
+		.ignore_err()
 		.ready_skip_while(|&(pducount, _)| pducount > next_batch.unwrap_or_else(PduCount::max))
 		.ready_take_while(|&(pducount, _)| pducount > roomsincecount);
 
 	// Take the last events for the timeline
-	let timeline_pdus: Vec<_> = non_timeline_pdus
-		.by_ref()
-		.take(limit)
-		.collect::<Vec<_>>()
-		.await
-		.into_iter()
-		.rev()
-		.collect();
+	pin_mut!(non_timeline_pdus);
+	let timeline_pdus: Vec<_> = non_timeline_pdus.by_ref().take(limit).collect().await;
+
+	let timeline_pdus: Vec<_> = timeline_pdus.into_iter().rev().collect();
 
 	// They /sync response doesn't always return all messages, so we say the output
 	// is limited unless there are events in non_timeline_pdus