use std::{
cmp::Ordering,
collections::{BinaryHeap, HashMap},
num::NonZeroUsize,
ops::ControlFlow,
sync::{Arc, Weak},
};
use matrix_sdk_base::{locks::Mutex, task_monitor::TaskMonitor};
use matrix_sdk_common::executor::{AbortOnDrop, JoinHandleExt as _, spawn};
use ruma::OwnedRoomId;
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::{CancellationToken, DropGuard};
use tracing::{debug, info, instrument, trace, warn};
use super::{EventCacheInner, caches::pagination::BackPaginationOutcome};
pub(crate) const BATCH_SIZE: u16 = 30;
#[allow(dead_code)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) enum Priority {
Low,
Normal,
High,
}
pub(crate) type StopCondition = Box<dyn FnMut(&BackPaginationOutcome) -> ControlFlow<()> + Send>;
pub(crate) struct BackPaginationRequest {
pub room_id: OwnedRoomId,
pub priority: Priority,
pub stop: StopCondition,
pub batch_size: u16,
pub max_batches: Option<usize>,
}
#[cfg(not(tarpaulin_include))]
impl std::fmt::Debug for BackPaginationRequest {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BackPaginationRequest")
.field("room_id", &self.room_id)
.field("priority", &self.priority)
.field("batch_size", &self.batch_size)
.field("max_batches", &self.max_batches)
.finish_non_exhaustive()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum BackPaginationStopReason {
ReachedTimelineStart,
StopConditionMet,
BatchLimitReached,
NoDataAvailable,
Failed,
Cancelled,
}
#[allow(dead_code)]
#[derive(Clone, Copy, Debug)]
pub(crate) struct BackPaginationRunResult {
pub reason: BackPaginationStopReason,
}
type RequestCoalescingKey = (OwnedRoomId, Priority);
pub(crate) struct BackPaginationHandle {
guard: Arc<DropGuard>,
#[allow(dead_code)]
completion: Option<oneshot::Receiver<BackPaginationRunResult>>,
}
#[cfg(not(tarpaulin_include))]
impl std::fmt::Debug for BackPaginationHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BackPaginationHandle").finish_non_exhaustive()
}
}
impl BackPaginationHandle {
pub(crate) fn detach(self) {
if let Some(guard) = Arc::into_inner(self.guard) {
guard.disarm();
}
}
#[allow(dead_code)]
pub(crate) async fn join(mut self) -> BackPaginationRunResult {
let cancelled = BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
match self.completion.take() {
Some(completion) => completion.await.unwrap_or(cancelled),
None => cancelled,
}
}
}
#[derive(Clone)]
pub struct BackPaginationQueue {
inner: Arc<BackPaginationQueueInner>,
}
struct BackPaginationQueueInner {
sender: mpsc::UnboundedSender<SchedulerEvent>,
cancellations: Mutex<HashMap<RequestCoalescingKey, SharedCancellation>>,
_task: matrix_sdk_base::task_monitor::BackgroundTaskHandle,
}
struct SharedCancellation {
token: CancellationToken,
guard: Weak<DropGuard>,
}
#[cfg(not(tarpaulin_include))]
impl std::fmt::Debug for BackPaginationQueue {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BackPaginationQueue").finish_non_exhaustive()
}
}
impl BackPaginationQueue {
pub(super) fn new(
event_cache: Weak<EventCacheInner>,
max_concurrent: NonZeroUsize,
task_monitor: &TaskMonitor,
) -> Self {
let (sender, receiver) = mpsc::unbounded_channel();
let task = task_monitor
.spawn_infinite_task(
"event_cache::back_pagination_queue",
scheduler(event_cache, receiver, sender.clone(), max_concurrent.get()),
)
.abort_on_drop();
Self {
inner: Arc::new(BackPaginationQueueInner {
sender,
cancellations: Mutex::new(HashMap::new()),
_task: task,
}),
}
}
pub(crate) fn enqueue(
&self,
request: BackPaginationRequest,
) -> Result<BackPaginationHandle, BackPaginationQueueError> {
let key = (request.room_id.clone(), request.priority);
let (token, guard) = cancellation_for(&self.inner.cancellations, key);
let (completion_tx, completion_rx) = oneshot::channel();
let submitted = SubmittedRequest { request, token, completion: completion_tx };
self.inner
.sender
.send(SchedulerEvent::Submitted(submitted))
.map_err(|_| BackPaginationQueueError::ShutDown)?;
Ok(BackPaginationHandle { guard, completion: Some(completion_rx) })
}
}
fn cancellation_for(
cancellations: &Mutex<HashMap<RequestCoalescingKey, SharedCancellation>>,
key: RequestCoalescingKey,
) -> (CancellationToken, Arc<DropGuard>) {
let mut cancellations = cancellations.lock();
if let Some(existing) = cancellations.get(&key)
&& let Some(guard) = existing.guard.upgrade()
{
return (existing.token.clone(), guard);
}
cancellations.retain(|_, cancellation| cancellation.guard.strong_count() > 0);
let token = CancellationToken::new();
let guard = Arc::new(token.clone().drop_guard());
cancellations
.insert(key, SharedCancellation { token: token.clone(), guard: Arc::downgrade(&guard) });
(token, guard)
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum BackPaginationQueueError {
#[error("the back-pagination queue executor is not running")]
ShutDown,
}
enum SchedulerEvent {
Submitted(SubmittedRequest),
Finished(RequestCoalescingKey, BackPaginationRunResult),
}
struct SubmittedRequest {
request: BackPaginationRequest,
token: CancellationToken,
completion: oneshot::Sender<BackPaginationRunResult>,
}
struct PendingRequest {
request: BackPaginationRequest,
seq: u64,
token: CancellationToken,
}
impl PartialEq for PendingRequest {
fn eq(&self, other: &Self) -> bool {
self.request.priority == other.request.priority && self.seq == other.seq
}
}
impl Eq for PendingRequest {}
impl PartialOrd for PendingRequest {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for PendingRequest {
fn cmp(&self, other: &Self) -> Ordering {
self.request.priority.cmp(&other.request.priority).then_with(|| other.seq.cmp(&self.seq))
}
}
#[instrument(skip_all)]
async fn scheduler(
event_cache: Weak<EventCacheInner>,
mut receiver: mpsc::UnboundedReceiver<SchedulerEvent>,
sender: mpsc::UnboundedSender<SchedulerEvent>,
max_concurrent: usize,
) {
trace!("Spawning the back-pagination queue executor");
let mut pending_requests: BinaryHeap<PendingRequest> = BinaryHeap::new();
let mut active_requests: HashMap<OwnedRoomId, AbortOnDrop<()>> = HashMap::new();
let mut next_seq: u64 = 0;
let mut waiters: HashMap<RequestCoalescingKey, Vec<oneshot::Sender<BackPaginationRunResult>>> =
HashMap::new();
let mut events = Vec::with_capacity(max_concurrent);
loop {
schedule(
&event_cache,
&mut pending_requests,
&mut active_requests,
max_concurrent,
&sender,
);
if receiver.recv_many(&mut events, max_concurrent).await == 0 {
info!("Back-pagination queue channel closed, exiting");
break;
}
for event in events.drain(..) {
match event {
SchedulerEvent::Submitted(submitted) => {
let key = (submitted.request.room_id.clone(), submitted.request.priority);
if try_coalesce(&mut waiters, &key, submitted.completion) {
trace!(
room_id = %key.0,
priority = ?key.1,
"coalesced back-pagination request onto an existing run"
);
continue;
}
pending_requests.push(PendingRequest {
request: submitted.request,
seq: next_seq,
token: submitted.token,
});
next_seq += 1;
}
SchedulerEvent::Finished(key, result) => {
active_requests.remove(&key.0);
if let Some(senders) = waiters.remove(&key) {
for waiter in senders {
let _ = waiter.send(result);
}
}
}
}
}
}
}
fn try_coalesce(
waiters: &mut HashMap<RequestCoalescingKey, Vec<oneshot::Sender<BackPaginationRunResult>>>,
key: &RequestCoalescingKey,
completion: oneshot::Sender<BackPaginationRunResult>,
) -> bool {
match waiters.get_mut(key) {
Some(existing) => {
existing.push(completion);
true
}
None => {
waiters.insert(key.to_owned(), vec![completion]);
false
}
}
}
fn schedule(
event_cache: &Weak<EventCacheInner>,
pending_requests: &mut BinaryHeap<PendingRequest>,
active_requests: &mut HashMap<OwnedRoomId, AbortOnDrop<()>>,
max_concurrent: usize,
sender: &mpsc::UnboundedSender<SchedulerEvent>,
) {
for request in next_runnable(pending_requests, active_requests, max_concurrent) {
let key = (request.request.room_id.clone(), request.request.priority);
trace!(
room_id = %key.0,
priority = ?key.1,
active = active_requests.len(),
queued = pending_requests.len(),
"back-pagination scheduled"
);
let room_id = key.0.clone();
let event_cache = event_cache.clone();
let sender = sender.clone();
let task = spawn(async move {
let result = run_request(&event_cache, request.request, &request.token).await;
let _ = sender.send(SchedulerEvent::Finished(key, result));
});
active_requests.insert(room_id, task.abort_on_drop());
}
}
fn next_runnable<T>(
pending_requests: &mut BinaryHeap<PendingRequest>,
active_requests: &HashMap<OwnedRoomId, T>,
max_concurrent: usize,
) -> Vec<PendingRequest> {
let mut picked: Vec<PendingRequest> = Vec::new();
let mut skipped = Vec::new();
while active_requests.len() + picked.len() < max_concurrent {
let Some(request) = pending_requests.pop() else {
break;
};
let room_id = &request.request.room_id;
if active_requests.contains_key(room_id)
|| picked.iter().any(|other| other.request.room_id == *room_id)
{
skipped.push(request);
continue;
}
picked.push(request);
}
pending_requests.extend(skipped);
picked
}
#[instrument(skip_all, fields(room_id = %request.room_id, priority = ?request.priority))]
async fn run_request(
event_cache: &Weak<EventCacheInner>,
mut request: BackPaginationRequest,
token: &CancellationToken,
) -> BackPaginationRunResult {
if token.is_cancelled() {
return BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
}
let pagination = {
let Some(inner) = event_cache.upgrade() else {
return BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
};
match inner.all_caches_for_room(&request.room_id).await {
Ok(caches) => caches.room().pagination(),
Err(err) => {
warn!("no caches for room while back-paginating: {err}");
return BackPaginationRunResult { reason: BackPaginationStopReason::Failed };
}
}
};
let mut batches = 0usize;
let reason = loop {
if token.is_cancelled() {
break BackPaginationStopReason::Cancelled;
}
let outcome = match pagination.run_backwards_once(request.batch_size).await {
Ok(outcome) => outcome,
Err(err) => {
warn!("back-pagination failed: {err}");
break BackPaginationStopReason::Failed;
}
};
if (request.stop)(&outcome).is_break() {
break BackPaginationStopReason::StopConditionMet;
}
if outcome.reached_start {
break BackPaginationStopReason::ReachedTimelineStart;
}
if outcome.events.is_empty() {
break BackPaginationStopReason::NoDataAvailable;
}
batches += 1;
if let Some(max) = request.max_batches
&& batches >= max
{
break BackPaginationStopReason::BatchLimitReached;
}
};
debug!(?reason, "back-pagination run finished");
BackPaginationRunResult { reason }
}
#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
use std::{collections::HashMap, ops::ControlFlow};
use matrix_sdk_base::locks::Mutex;
use ruma::room_id;
use super::{
BackPaginationRequest, PendingRequest, Priority, cancellation_for, next_runnable,
try_coalesce,
};
fn queued(room_id: ruma::OwnedRoomId, priority: Priority, seq: u64) -> PendingRequest {
PendingRequest {
request: BackPaginationRequest {
room_id,
priority,
stop: Box::new(|_| ControlFlow::Continue(())),
batch_size: 10,
max_batches: None,
},
seq,
token: tokio_util::sync::CancellationToken::new(),
}
}
#[test]
fn test_scheduling_priority_and_fifo() {
use std::collections::BinaryHeap;
let (a, b, c, d) = (room_id!("!a:e"), room_id!("!b:e"), room_id!("!c:e"), room_id!("!d:e"));
let mut pending_requests = BinaryHeap::new();
pending_requests.push(queued(a.to_owned(), Priority::Low, 0));
pending_requests.push(queued(b.to_owned(), Priority::High, 1));
pending_requests.push(queued(c.to_owned(), Priority::Normal, 2));
pending_requests.push(queued(d.to_owned(), Priority::High, 3));
let active_requests: HashMap<_, ()> = HashMap::new();
let picked: Vec<_> = next_runnable(&mut pending_requests, &active_requests, 10)
.into_iter()
.map(|r| r.request.room_id)
.collect();
assert_eq!(picked, vec![b.to_owned(), d.to_owned(), c.to_owned(), a.to_owned()]);
}
#[test]
fn test_scheduling_respects_concurrency_cap() {
use std::collections::BinaryHeap;
let mut pending_requests = BinaryHeap::new();
for (i, room) in [room_id!("!a:e"), room_id!("!b:e"), room_id!("!c:e")].iter().enumerate() {
pending_requests.push(queued((*room).to_owned(), Priority::Normal, i as u64));
}
let active_requests: HashMap<_, ()> = HashMap::new();
let picked = next_runnable(&mut pending_requests, &active_requests, 2);
assert_eq!(picked.len(), 2);
assert_eq!(pending_requests.len(), 1);
}
#[test]
fn test_scheduling_per_room_single_flight() {
use std::collections::BinaryHeap;
let (a, b) = (room_id!("!a:e"), room_id!("!b:e"));
let active_requests = HashMap::from([(a.to_owned(), ())]);
let mut pending_requests = BinaryHeap::new();
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)
.into_iter()
.map(|r| r.request.room_id)
.collect();
assert_eq!(picked, vec![b.to_owned()]);
assert_eq!(pending_requests.len(), 2);
}
#[test]
fn test_cancellation_is_shared_per_key() {
let cancellations = Mutex::new(HashMap::new());
let key = (room_id!("!a:e").to_owned(), Priority::Normal);
let (first_token, first_guard) = cancellation_for(&cancellations, key.clone());
let (second_token, second_guard) = cancellation_for(&cancellations, key.clone());
assert!(!first_token.is_cancelled());
drop(first_guard);
assert!(!first_token.is_cancelled());
assert!(!second_token.is_cancelled());
drop(second_guard);
assert!(first_token.is_cancelled());
assert!(second_token.is_cancelled());
let (third_token, _third_guard) = cancellation_for(&cancellations, key);
assert!(!third_token.is_cancelled());
let other = (room_id!("!a:e").to_owned(), Priority::High);
let (other_token, other_guard) = cancellation_for(&cancellations, other);
drop(other_guard);
assert!(other_token.is_cancelled());
assert!(!third_token.is_cancelled());
}
#[test]
fn test_coalescing() {
let a = room_id!("!a:e");
let mut waiters = HashMap::new();
let completion = || tokio::sync::oneshot::channel().0;
let normal = (a.to_owned(), Priority::Normal);
let high = (a.to_owned(), Priority::High);
assert!(!try_coalesce(&mut waiters, &normal, completion()));
assert_eq!(waiters[&normal].len(), 1);
assert!(try_coalesce(&mut waiters, &normal, completion()));
assert_eq!(waiters[&normal].len(), 2);
assert!(!try_coalesce(&mut waiters, &high, completion()));
assert_eq!(waiters.len(), 2);
assert_eq!(waiters[&high].len(), 1);
}
}