1use std::{
40 cmp::Ordering,
41 collections::{BinaryHeap, HashMap},
42 num::NonZeroUsize,
43 ops::ControlFlow,
44 sync::{Arc, Weak},
45};
46
47use matrix_sdk_base::{RoomState, locks::Mutex, task_monitor::TaskMonitor};
48use matrix_sdk_common::executor::{AbortOnDrop, JoinHandleExt as _, spawn};
49use ruma::OwnedRoomId;
50use tokio::sync::{mpsc, oneshot};
51use tokio_util::sync::{CancellationToken, DropGuard};
52use tracing::{debug, info, instrument, trace, warn};
53
54use super::{EventCacheInner, caches::pagination::BackPaginationOutcome};
55
56pub(crate) const BATCH_SIZE: u16 = 30;
58
59#[allow(dead_code)]
62#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
63pub(crate) enum Priority {
64 Low,
67 Normal,
70 High,
72}
73
74pub(crate) type StopCondition = Box<dyn FnMut(&BackPaginationOutcome) -> ControlFlow<()> + Send>;
82
83pub(crate) struct BackPaginationRequest {
86 pub room_id: OwnedRoomId,
88 pub priority: Priority,
90 pub stop: StopCondition,
92 pub batch_size: u16,
94 pub max_batches: Option<usize>,
96}
97
98#[cfg(not(tarpaulin_include))]
99impl std::fmt::Debug for BackPaginationRequest {
100 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101 f.debug_struct("BackPaginationRequest")
102 .field("room_id", &self.room_id)
103 .field("priority", &self.priority)
104 .field("batch_size", &self.batch_size)
105 .field("max_batches", &self.max_batches)
106 .finish_non_exhaustive()
107 }
108}
109
110#[derive(Clone, Copy, Debug, PartialEq, Eq)]
112pub(crate) enum BackPaginationStopReason {
113 ReachedTimelineStart,
115 StopConditionMet,
117 BatchLimitReached,
120 NoDataAvailable,
123 Failed,
125 Cancelled,
127 Skipped,
130}
131
132#[allow(dead_code)]
134#[derive(Clone, Copy, Debug)]
135pub(crate) struct BackPaginationRunResult {
136 pub reason: BackPaginationStopReason,
138}
139
140type RequestCoalescingKey = (OwnedRoomId, Priority);
142
143pub(crate) struct BackPaginationHandle {
146 guard: Arc<DropGuard>,
149 #[allow(dead_code)]
151 completion: Option<oneshot::Receiver<BackPaginationRunResult>>,
152}
153
154#[cfg(not(tarpaulin_include))]
155impl std::fmt::Debug for BackPaginationHandle {
156 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
157 f.debug_struct("BackPaginationHandle").finish_non_exhaustive()
158 }
159}
160
161impl BackPaginationHandle {
162 pub(crate) fn detach(self) {
168 if let Some(guard) = Arc::into_inner(self.guard) {
169 guard.disarm();
170 }
171 }
172
173 #[allow(dead_code)]
175 pub(crate) async fn join(mut self) -> BackPaginationRunResult {
176 let cancelled = BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
177 match self.completion.take() {
178 Some(completion) => completion.await.unwrap_or(cancelled),
179 None => cancelled,
180 }
181 }
182}
183
184#[derive(Clone)]
186pub struct BackPaginationQueue {
187 inner: Arc<BackPaginationQueueInner>,
188}
189
190struct BackPaginationQueueInner {
191 sender: mpsc::UnboundedSender<SchedulerEvent>,
192 cancellations: Mutex<HashMap<RequestCoalescingKey, SharedCancellation>>,
199 _task: matrix_sdk_base::task_monitor::BackgroundTaskHandle,
200}
201
202struct SharedCancellation {
204 token: CancellationToken,
206 guard: Weak<DropGuard>,
209}
210
211#[cfg(not(tarpaulin_include))]
212impl std::fmt::Debug for BackPaginationQueue {
213 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214 f.debug_struct("BackPaginationQueue").finish_non_exhaustive()
215 }
216}
217
218impl BackPaginationQueue {
219 pub(super) fn new(
221 event_cache: Weak<EventCacheInner>,
222 max_concurrent: NonZeroUsize,
223 task_monitor: &TaskMonitor,
224 ) -> Self {
225 let (sender, receiver) = mpsc::unbounded_channel();
226
227 let task = task_monitor
231 .spawn_infinite_task(
232 "event_cache::back_pagination_queue",
233 scheduler(event_cache, receiver, sender.clone(), max_concurrent.get()),
234 )
235 .abort_on_drop();
236
237 Self {
238 inner: Arc::new(BackPaginationQueueInner {
239 sender,
240 cancellations: Mutex::new(HashMap::new()),
241 _task: task,
242 }),
243 }
244 }
245
246 pub(crate) fn enqueue(
250 &self,
251 request: BackPaginationRequest,
252 ) -> Result<BackPaginationHandle, BackPaginationQueueError> {
253 let key = (request.room_id.clone(), request.priority);
254 let (token, guard) = cancellation_for(&self.inner.cancellations, key);
255
256 let (completion_tx, completion_rx) = oneshot::channel();
257 let submitted = SubmittedRequest { request, token, completion: completion_tx };
258
259 self.inner
260 .sender
261 .send(SchedulerEvent::Submitted(submitted))
262 .map_err(|_| BackPaginationQueueError::ShutDown)?;
263
264 Ok(BackPaginationHandle { guard, completion: Some(completion_rx) })
265 }
266}
267
268fn cancellation_for(
275 cancellations: &Mutex<HashMap<RequestCoalescingKey, SharedCancellation>>,
276 key: RequestCoalescingKey,
277) -> (CancellationToken, Arc<DropGuard>) {
278 let mut cancellations = cancellations.lock();
279
280 if let Some(existing) = cancellations.get(&key)
281 && let Some(guard) = existing.guard.upgrade()
282 {
283 return (existing.token.clone(), guard);
284 }
285
286 cancellations.retain(|_, cancellation| cancellation.guard.strong_count() > 0);
288
289 let token = CancellationToken::new();
290 let guard = Arc::new(token.clone().drop_guard());
291 cancellations
292 .insert(key, SharedCancellation { token: token.clone(), guard: Arc::downgrade(&guard) });
293
294 (token, guard)
295}
296
297#[derive(Debug, thiserror::Error)]
299pub(crate) enum BackPaginationQueueError {
300 #[error("the back-pagination queue executor is not running")]
303 ShutDown,
304}
305
306enum SchedulerEvent {
309 Submitted(SubmittedRequest),
311 Finished(RequestCoalescingKey, BackPaginationRunResult),
313}
314
315struct SubmittedRequest {
318 request: BackPaginationRequest,
319 token: CancellationToken,
320 completion: oneshot::Sender<BackPaginationRunResult>,
321}
322
323struct PendingRequest {
326 request: BackPaginationRequest,
327 seq: u64,
329 token: CancellationToken,
330}
331
332impl PartialEq for PendingRequest {
333 fn eq(&self, other: &Self) -> bool {
334 self.request.priority == other.request.priority && self.seq == other.seq
335 }
336}
337
338impl Eq for PendingRequest {}
339
340impl PartialOrd for PendingRequest {
341 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
342 Some(self.cmp(other))
343 }
344}
345
346impl Ord for PendingRequest {
347 fn cmp(&self, other: &Self) -> Ordering {
348 self.request.priority.cmp(&other.request.priority).then_with(|| other.seq.cmp(&self.seq))
350 }
351}
352
353#[instrument(skip_all)]
356async fn scheduler(
357 event_cache: Weak<EventCacheInner>,
358 mut receiver: mpsc::UnboundedReceiver<SchedulerEvent>,
359 sender: mpsc::UnboundedSender<SchedulerEvent>,
360 max_concurrent: usize,
361) {
362 trace!("Spawning the back-pagination queue executor");
363
364 let mut pending_requests: BinaryHeap<PendingRequest> = BinaryHeap::new();
365 let mut active_requests: HashMap<OwnedRoomId, AbortOnDrop<()>> = HashMap::new();
369 let mut next_seq: u64 = 0;
370
371 let mut waiters: HashMap<RequestCoalescingKey, Vec<oneshot::Sender<BackPaginationRunResult>>> =
377 HashMap::new();
378
379 let mut events = Vec::with_capacity(max_concurrent);
382
383 loop {
384 schedule(
387 &event_cache,
388 &mut pending_requests,
389 &mut active_requests,
390 max_concurrent,
391 &sender,
392 );
393
394 if receiver.recv_many(&mut events, max_concurrent).await == 0 {
397 info!("Back-pagination queue channel closed, exiting");
398 break;
399 }
400
401 for event in events.drain(..) {
402 match event {
403 SchedulerEvent::Submitted(submitted) => {
404 let key = (submitted.request.room_id.clone(), submitted.request.priority);
405
406 if try_coalesce(&mut waiters, &key, submitted.completion) {
407 trace!(
408 room_id = %key.0,
409 priority = ?key.1,
410 "coalesced back-pagination request onto an existing run"
411 );
412 continue;
413 }
414
415 pending_requests.push(PendingRequest {
416 request: submitted.request,
417 seq: next_seq,
418 token: submitted.token,
419 });
420 next_seq += 1;
421 }
422
423 SchedulerEvent::Finished(key, result) => {
424 active_requests.remove(&key.0);
425 if let Some(senders) = waiters.remove(&key) {
428 for waiter in senders {
429 let _ = waiter.send(result);
430 }
431 }
432 }
433 }
434 }
435 }
436}
437
438fn try_coalesce(
447 waiters: &mut HashMap<RequestCoalescingKey, Vec<oneshot::Sender<BackPaginationRunResult>>>,
448 key: &RequestCoalescingKey,
449 completion: oneshot::Sender<BackPaginationRunResult>,
450) -> bool {
451 match waiters.get_mut(key) {
452 Some(existing) => {
453 existing.push(completion);
454 true
455 }
456 None => {
457 waiters.insert(key.to_owned(), vec![completion]);
458 false
459 }
460 }
461}
462
463fn schedule(
465 event_cache: &Weak<EventCacheInner>,
466 pending_requests: &mut BinaryHeap<PendingRequest>,
467 active_requests: &mut HashMap<OwnedRoomId, AbortOnDrop<()>>,
468 max_concurrent: usize,
469 sender: &mpsc::UnboundedSender<SchedulerEvent>,
470) {
471 for request in next_runnable(pending_requests, active_requests, max_concurrent) {
472 let key = (request.request.room_id.clone(), request.request.priority);
473
474 trace!(
475 room_id = %key.0,
476 priority = ?key.1,
477 active = active_requests.len(),
478 queued = pending_requests.len(),
479 "back-pagination scheduled"
480 );
481
482 let room_id = key.0.clone();
483 let event_cache = event_cache.clone();
484 let sender = sender.clone();
485 let task = spawn(async move {
486 let result = run_request(&event_cache, request.request, &request.token).await;
487 let _ = sender.send(SchedulerEvent::Finished(key, result));
490 });
491
492 active_requests.insert(room_id, task.abort_on_drop());
493 }
494}
495
496fn next_runnable<T>(
504 pending_requests: &mut BinaryHeap<PendingRequest>,
505 active_requests: &HashMap<OwnedRoomId, T>,
506 max_concurrent: usize,
507) -> Vec<PendingRequest> {
508 let mut picked: Vec<PendingRequest> = Vec::new();
509 let mut skipped = Vec::new();
510
511 while active_requests.len() + picked.len() < max_concurrent {
512 let Some(request) = pending_requests.pop() else {
513 break;
514 };
515
516 let room_id = &request.request.room_id;
517 if active_requests.contains_key(room_id)
518 || picked.iter().any(|other| other.request.room_id == *room_id)
519 {
520 skipped.push(request);
522 continue;
523 }
524
525 picked.push(request);
526 }
527
528 pending_requests.extend(skipped);
529
530 picked
531}
532
533#[instrument(skip_all, fields(room_id = %request.room_id, priority = ?request.priority))]
536async fn run_request(
537 event_cache: &Weak<EventCacheInner>,
538 mut request: BackPaginationRequest,
539 token: &CancellationToken,
540) -> BackPaginationRunResult {
541 if token.is_cancelled() {
543 return BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
544 }
545
546 let (pagination, _subscriber) = {
554 let Some(inner) = event_cache.upgrade() else {
555 return BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
556 };
557 if let Ok(client) = inner.client()
558 && let Some(room) = client.get_room(&request.room_id)
559 && (room.is_space() || room.state() != RoomState::Joined)
560 {
561 debug!("skipping back-pagination for a space or a non-joined room");
562 return BackPaginationRunResult { reason: BackPaginationStopReason::Skipped };
563 }
564 let room = match inner.all_caches_for_room(&request.room_id).await {
565 Ok(caches) => caches.room().clone(),
566 Err(err) => {
567 warn!("no caches for room while back-paginating: {err}");
568 return BackPaginationRunResult { reason: BackPaginationStopReason::Failed };
569 }
570 };
571 let subscriber = match room.subscribe().await {
572 Ok((_events, subscriber)) => subscriber,
573 Err(err) => {
574 warn!("couldn't subscribe to the room while back-paginating: {err}");
575 return BackPaginationRunResult { reason: BackPaginationStopReason::Failed };
576 }
577 };
578 (room.pagination(), subscriber)
579 };
580
581 let mut batches = 0usize;
582
583 let reason = loop {
584 if token.is_cancelled() {
585 break BackPaginationStopReason::Cancelled;
586 }
587
588 let outcome = match pagination.run_backwards_once(request.batch_size).await {
589 Ok(outcome) => outcome,
590 Err(err) => {
591 warn!("back-pagination failed: {err}");
592 break BackPaginationStopReason::Failed;
593 }
594 };
595
596 if (request.stop)(&outcome).is_break() {
599 break BackPaginationStopReason::StopConditionMet;
600 }
601
602 if outcome.reached_start {
603 break BackPaginationStopReason::ReachedTimelineStart;
604 }
605
606 if outcome.events.is_empty() {
607 break BackPaginationStopReason::NoDataAvailable;
608 }
609
610 batches += 1;
611 if let Some(max) = request.max_batches
612 && batches >= max
613 {
614 break BackPaginationStopReason::BatchLimitReached;
615 }
616 };
617
618 debug!(?reason, "back-pagination run finished");
619
620 BackPaginationRunResult { reason }
621}
622
623#[cfg(all(test, not(target_arch = "wasm32")))]
624mod tests {
625 use std::{collections::HashMap, ops::ControlFlow};
626
627 use matrix_sdk_base::locks::Mutex;
628 use ruma::room_id;
629
630 use super::{
631 BackPaginationRequest, PendingRequest, Priority, cancellation_for, next_runnable,
632 try_coalesce,
633 };
634
635 fn queued(room_id: ruma::OwnedRoomId, priority: Priority, seq: u64) -> PendingRequest {
637 PendingRequest {
638 request: BackPaginationRequest {
639 room_id,
640 priority,
641 stop: Box::new(|_| ControlFlow::Continue(())),
643 batch_size: 10,
644 max_batches: None,
645 },
646 seq,
647 token: tokio_util::sync::CancellationToken::new(),
648 }
649 }
650
651 #[test]
654 fn test_scheduling_priority_and_fifo() {
655 use std::collections::BinaryHeap;
656
657 let (a, b, c, d) = (room_id!("!a:e"), room_id!("!b:e"), room_id!("!c:e"), room_id!("!d:e"));
658
659 let mut pending_requests = BinaryHeap::new();
660 pending_requests.push(queued(a.to_owned(), Priority::Low, 0));
662 pending_requests.push(queued(b.to_owned(), Priority::High, 1));
663 pending_requests.push(queued(c.to_owned(), Priority::Normal, 2));
664 pending_requests.push(queued(d.to_owned(), Priority::High, 3));
665
666 let active_requests: HashMap<_, ()> = HashMap::new();
667 let picked: Vec<_> = next_runnable(&mut pending_requests, &active_requests, 10)
668 .into_iter()
669 .map(|r| r.request.room_id)
670 .collect();
671
672 assert_eq!(picked, vec![b.to_owned(), d.to_owned(), c.to_owned(), a.to_owned()]);
674 }
675
676 #[test]
678 fn test_scheduling_respects_concurrency_cap() {
679 use std::collections::BinaryHeap;
680
681 let mut pending_requests = BinaryHeap::new();
682 for (i, room) in [room_id!("!a:e"), room_id!("!b:e"), room_id!("!c:e")].iter().enumerate() {
683 pending_requests.push(queued((*room).to_owned(), Priority::Normal, i as u64));
684 }
685
686 let active_requests: HashMap<_, ()> = HashMap::new();
687 let picked = next_runnable(&mut pending_requests, &active_requests, 2);
688
689 assert_eq!(picked.len(), 2);
690 assert_eq!(pending_requests.len(), 1);
692 }
693
694 #[test]
697 fn test_scheduling_per_room_single_flight() {
698 use std::collections::BinaryHeap;
699
700 let (a, b) = (room_id!("!a:e"), room_id!("!b:e"));
701
702 let active_requests = HashMap::from([(a.to_owned(), ())]);
704
705 let mut pending_requests = BinaryHeap::new();
706 pending_requests.push(queued(a.to_owned(), Priority::High, 0)); pending_requests.push(queued(a.to_owned(), Priority::High, 1)); pending_requests.push(queued(b.to_owned(), Priority::Low, 2)); let picked: Vec<_> = next_runnable(&mut pending_requests, &active_requests, 10)
711 .into_iter()
712 .map(|r| r.request.room_id)
713 .collect();
714
715 assert_eq!(picked, vec![b.to_owned()]);
717 assert_eq!(pending_requests.len(), 2);
718 }
719
720 #[test]
723 fn test_cancellation_is_shared_per_key() {
724 let cancellations = Mutex::new(HashMap::new());
725 let key = (room_id!("!a:e").to_owned(), Priority::Normal);
726
727 let (first_token, first_guard) = cancellation_for(&cancellations, key.clone());
728 let (second_token, second_guard) = cancellation_for(&cancellations, key.clone());
729
730 assert!(!first_token.is_cancelled());
732 drop(first_guard);
733 assert!(!first_token.is_cancelled());
734 assert!(!second_token.is_cancelled());
735
736 drop(second_guard);
738 assert!(first_token.is_cancelled());
739 assert!(second_token.is_cancelled());
740
741 let (third_token, _third_guard) = cancellation_for(&cancellations, key);
744 assert!(!third_token.is_cancelled());
745
746 let other = (room_id!("!a:e").to_owned(), Priority::High);
748 let (other_token, other_guard) = cancellation_for(&cancellations, other);
749 drop(other_guard);
750 assert!(other_token.is_cancelled());
751 assert!(!third_token.is_cancelled());
752 }
753
754 #[test]
758 fn test_coalescing() {
759 let a = room_id!("!a:e");
760 let mut waiters = HashMap::new();
761
762 let completion = || tokio::sync::oneshot::channel().0;
763 let normal = (a.to_owned(), Priority::Normal);
764 let high = (a.to_owned(), Priority::High);
765
766 assert!(!try_coalesce(&mut waiters, &normal, completion()));
768 assert_eq!(waiters[&normal].len(), 1);
769
770 assert!(try_coalesce(&mut waiters, &normal, completion()));
772 assert_eq!(waiters[&normal].len(), 2);
773
774 assert!(!try_coalesce(&mut waiters, &high, completion()));
776 assert_eq!(waiters.len(), 2);
777 assert_eq!(waiters[&high].len(), 1);
778 }
779}