diff --git a/src/service/rooms/event_handler/handle_incoming_pdu.rs b/src/service/rooms/event_handler/handle_incoming_pdu.rs index 0c9974ee8..ae7fcb5e4 100644 --- a/src/service/rooms/event_handler/handle_incoming_pdu.rs +++ b/src/service/rooms/event_handler/handle_incoming_pdu.rs @@ -18,6 +18,7 @@ room::member::{MembershipState, RoomMemberEventContent}, }, }; +use tokio::sync::mpsc; use crate::rooms::timeline::{RawPduId, pdu_fits}; @@ -267,6 +268,8 @@ pub async fn handle_incoming_pdu<'a>( error!("Failed to fetch and persist incoming event's prev_events: {e:?}"); })?; + let is_dummy_event = incoming_pdu.event_type().to_string() == "org.matrix.dummy_event"; + // Done with prev events, now handling the incoming event let pdu_id = Box::pin(self.upgrade_outlier_to_timeline_pdu( incoming_pdu, @@ -284,24 +287,108 @@ pub async fn handle_incoming_pdu<'a>( .count() .await; - if extremities_count >= self.services.server.config.dummy_event_threshold.into() { - self.squash_extremities(room_id, extremities_count).await; - } + self.maybe_squash_extremities(room_id, extremities_count, is_dummy_event) + .await; Ok(pdu_id) } - async fn squash_extremities(&self, room_id: &RoomId, count: usize) { - let last_squash = { - let squash_timings = self.last_extremity_squash.read(); - squash_timings.get(room_id).copied() + async fn maybe_squash_extremities( + &self, + room_id: &RoomId, + extremities_count: usize, + is_dummy_event: bool, + ) { + let (tx, fut) = { + if let Some(tx) = self.extremity_squashers.read().get(room_id) + && !tx.is_closed() + { + (tx.clone(), None) + } else { + let mut map = self.extremity_squashers.upgradable_read(); + + if let Some(tx) = map.get(room_id) + && !tx.is_closed() + { + (tx.clone(), None) + } else { + let (tx, rx) = mpsc::channel(100); + map.with_upgraded(|map| map.insert(room_id.to_owned(), tx.clone())); + + (tx, Some(self.spawn_squasher(room_id, rx))) + } + } }; - if last_squash.is_some_and(|s| s.elapsed() < Duration::from_mins(1)) { - // Avoid sending more than one squash per minute to avoid flooding rooms. - return; + + if let Some(fut) = fut { + fut.await; } + let _ = tx.try_send((extremities_count, is_dummy_event)); + } + + async fn spawn_squasher(&self, room_id: &RoomId, mut rx: mpsc::Receiver<(usize, bool)>) { + let Some(service) = self.me.upgrade() else { + return; + }; + let room_id = room_id.to_owned(); + + self.services.server.runtime().spawn(async move { + let mut latest_extremity_count = None; + let mut non_dummy_event = false; + + let mut closing = false; + + let waker = tokio::time::sleep(Duration::from_mins(2)); + tokio::pin!(waker); + + loop { + tokio::select! { + msg = rx.recv() => { + if let Some((extremities_count, is_dummy_event)) = msg { + latest_extremity_count = Some(extremities_count); + non_dummy_event = non_dummy_event || !is_dummy_event; + #[allow(clippy::arithmetic_side_effects)] + waker.as_mut().reset(tokio::time::Instant::now() + Duration::from_mins(1)); + } else { + {let mut map = service.extremity_squashers.write(); + if let Some(tx) = map.get(&room_id) && tx.is_closed() { + map.remove(&room_id); + }} + + if let Some(count) = latest_extremity_count { + if non_dummy_event && count >= service.services.server.config.dummy_event_threshold.into() { + Self::squash_extremities(&service, &room_id, count).await; + } + } + break; + } + } + () = &mut waker, if !closing => { + if let Some(count) = latest_extremity_count { + if non_dummy_event && count >= service.services.server.config.dummy_event_threshold.into() { + Self::squash_extremities(&service, &room_id, count).await; + } + latest_extremity_count = None; + non_dummy_event = false; + #[allow(clippy::arithmetic_side_effects)] + waker.as_mut().reset(tokio::time::Instant::now() + Duration::from_mins(2)); + } else { + rx.close(); + closing = true; + } + } + () = service.server_shutdown.notified(), if !closing => { + rx.close(); + closing = true; + } + } + } + }); + } + + async fn squash_extremities(&self, room_id: &RoomId, extremities_count: usize) { debug_warn!( - %count, + %extremities_count, threshold=%self.services.server.config.dummy_event_threshold, "Attempting to squash extremities after upgrading pdu" ); @@ -340,7 +427,5 @@ async fn squash_extremities(&self, room_id: &RoomId, count: usize) { break; } } - let mut squash_timings = self.last_extremity_squash.write(); - squash_timings.insert(room_id.to_owned(), Instant::now()); } } diff --git a/src/service/rooms/event_handler/mod.rs b/src/service/rooms/event_handler/mod.rs index 616dae3a4..3eae2f477 100644 --- a/src/service/rooms/event_handler/mod.rs +++ b/src/service/rooms/event_handler/mod.rs @@ -19,15 +19,16 @@ OwnedEventId, OwnedRoomId, RoomId, events::room::create::RoomCreateEventContent, room_version_rules::RoomVersionRules, }; -use tokio::sync::Notify; +use tokio::sync::{Notify, mpsc}; use crate::{Dep, globals, rooms, sending, server_keys}; pub struct Service { pub mutex_federation: RoomMutexMap, pub federation_handletime: SyncRwLock, - pub last_extremity_squash: SyncRwLock>, + pub extremity_squashers: SyncRwLock>>, services: Services, server_shutdown: Notify, + me: std::sync::Weak, } struct Services { @@ -53,10 +54,11 @@ struct Services { #[async_trait] impl crate::Service for Service { fn build(args: crate::Args<'_>) -> Result> { - Ok(Arc::new(Self { + Ok(Arc::new_cyclic(|s| Self { + me: s.clone(), mutex_federation: RoomMutexMap::new(), federation_handletime: HandleTimeMap::new().into(), - last_extremity_squash: SyncRwLock::new(HashMap::new()), + extremity_squashers: SyncRwLock::new(HashMap::new()), services: Services { globals: args.depend::("globals"), sending: args.depend::("sending"), @@ -93,10 +95,7 @@ fn name(&self) -> &str { crate::service::make_name(std::module_path!()) } fn interrupt(&self) { self.server_shutdown.notify_waiters(); } - async fn clear_cache(&self) { - let mut squashes = self.last_extremity_squash.write(); - squashes.clear(); - } + async fn clear_cache(&self) {} } impl Service {