From dc302799a0811d5078882552c2990ab3857a5c40 Mon Sep 17 00:00:00 2001
From: timedout <git@nexy7574.co.uk>
Date: Fri, 28 Aug 2026 15:41:26 +0100
Subject: [PATCH] feat: Add pdu pull backoff

# Conflicts:
#	src/service/rooms/event_handler/mod.rs
---
 src/service/rooms/event_handler/mod.rs | 64 ++++++++++++++++++++++++--
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/src/service/rooms/event_handler/mod.rs b/src/service/rooms/event_handler/mod.rs
index 6163e3972..11318af76 100644
--- a/src/service/rooms/event_handler/mod.rs
+++ b/src/service/rooms/event_handler/mod.rs
@@ -14,24 +14,32 @@ mod state_at_incoming;
 mod upgrade_outlier_pdu;
 
 use std::{collections::HashMap, fmt::Write, sync::Arc, time::Instant};
-
+use std::time::Duration;
+use assign::assign;
 use async_trait::async_trait;
-use conduwuit::{Err, Event, PduEvent, Result, Server, SyncRwLock, utils::MutexMap};
+use http::StatusCode;
+use conduwuit::{utils::MutexMap, Err, Error, Event, PduEvent, Result, Server, SyncRwLock};
 pub use fetch_and_handle_outliers::{
 	DagBuilderTree, GET_MISSING_EVENTS_MAX_BATCH_SIZE, build_local_dag,
 };
 use ruma::{
-	OwnedEventId, OwnedRoomId, OwnedServerName, events::room::create::RoomCreateEventContent,
+	EventId, OwnedEventId, OwnedRoomId, OwnedServerName,
+	api::error::{ErrorKind, LimitExceededErrorData, RetryAfter},
+	events::room::create::RoomCreateEventContent,
 	room_version_rules::RoomVersionRules,
 };
 use serde_json::value::RawValue as RawJsonValue;
 use tokio::sync::{Notify, mpsc};
 
 use crate::{Dep, globals, rooms, sending, server_keys};
+
+pub type FailedPDUPull = (u32, Instant);  // (pull count, last retry)
+
 pub struct Service {
 	pub mutex_federation: RoomMutexMap,
 	pub federation_handletime: SyncRwLock<HandleTimeMap>,
 	pub extremity_squashers: SyncRwLock<HashMap<OwnedRoomId, mpsc::Sender<(usize, bool)>>>,
+	pub failed_pdu_pulls: SyncRwLock<HashMap<OwnedEventId, FailedPDUPull>>,
 	joining_rooms: SyncRwLock<HashMap<OwnedRoomId, mpsc::Sender<PendingJoinPdu>>>,
 	services: Services,
 	server_shutdown: Notify,
@@ -71,6 +79,7 @@ impl crate::Service for Service {
 			mutex_federation: RoomMutexMap::new(),
 			federation_handletime: HandleTimeMap::new().into(),
 			extremity_squashers: SyncRwLock::new(HashMap::new()),
+			failed_pdu_pulls: SyncRwLock::new(HashMap::new()),
 			joining_rooms: HashMap::new().into(),
 			services: Services {
 				globals: args.depend::<globals::Service>("globals"),
@@ -108,7 +117,10 @@ impl crate::Service for Service {
 
 	fn interrupt(&self) { self.server_shutdown.notify_waiters(); }
 
-	async fn clear_cache(&self) {}
+	async fn clear_cache(&self) {
+		let mut failed_pdu_pulls = self.failed_pdu_pulls.write();
+		failed_pdu_pulls.clear();
+	}
 }
 
 impl Service {
@@ -122,6 +134,50 @@ impl Service {
 	async fn event_fetch(&self, event_id: OwnedEventId) -> Option<PduEvent> {
 		self.services.timeline.get_pdu(&event_id).await.ok()
 	}
+
+	/// Returns a rate-limit error if the requested event had a recent failed pull attempt.
+	///
+	/// If the event is not being backed off from, `Ok(())` is returned.
+	pub(super) fn ensure_can_pull_event(&self, event_id: &EventId) -> Result<()> {
+		let map = self.failed_pdu_pulls.read();
+		let Some((retry_count, last_retry)) = map.get(event_id) else {
+			return Ok(());
+		};
+		let min = Duration::from_secs(self.services.server.config.sender_retry_backoff_base);
+		let max = Duration::from_secs(self.services.server.config.sender_retry_backoff_limit);
+		let next_retry = min.saturating_mul(*retry_count).min(max);
+		if last_retry.elapsed() >= next_retry {
+			Ok(())
+		} else {
+			Err(Error::Request(
+				ErrorKind::LimitExceeded(assign!(LimitExceededErrorData::new(), {
+					retry_after: Some(RetryAfter::Delay(next_retry)),
+				})),
+				format!("Event {event_id} failed to be fetched recently, backing off").into(),
+				StatusCode::TOO_MANY_REQUESTS,
+			))
+		}
+	}
+
+	/// Marks a PDU as having a failed pull attempt, preventing it from being immediately re-fetched without a short cooldown.
+	pub(super) fn hit_failed_pdu_pull(&self, event_id: OwnedEventId) {
+		let now = Instant::now();
+		let mut map = self.failed_pdu_pulls.write();
+		map.entry(event_id).and_modify(|(retries, last_retry)| {
+			if *last_retry > now {
+				return;
+			}
+			*retries = retries.saturating_add(1);
+			*last_retry = Instant::now();
+		}).or_insert_with(|| (1, Instant::now()));
+	}
+
+	/// Removes a PDU from the failed pulls map, allowing it to be re-fetched in future if needed.
+	pub(super) fn clear_failed_pdu(&self, event_id: &EventId) {
+		// NOTE: this is a bit pointless since we're unlikely to try pulling an event again if we fetch it successfully, but it doesn't hurt to have to prevent the map growing indefinitely.
+		let mut map = self.failed_pdu_pulls.write();
+		map.remove(event_id);
+	}
 }
 
 fn get_room_version_rules<Pdu: Event>(create_event: &Pdu) -> Result<RoomVersionRules> {