Skip to main content

matrix_sdk/event_cache/
back_pagination_queue.rs

1// Copyright 2026 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A single component that owns one background task and executes
16//! back-pagination requests coming from various components (search,
17//! latest event, read receipt).
18//!
19//! Callers enqueue a [`BackPaginationRequest`] describing which room to
20//! back-paginate, at which priority, and until when and they get a
21//! [`BackPaginationHandle`] to await. Each consumer builds its own requests, in
22//! its own module; this one only schedules and runs them.
23//!
24//! The executor:
25//! - runs at most [`EventCacheConfig::max_concurrent_back_paginations`]
26//!   requests concurrently
27//! - schedules by [`Priority`], higher first, FIFO within a priority
28//! - never runs two requests for the same room concurrently
29//! - deduplicates by room *and* priority: a request for a room already queued
30//!   or running at the same priority is coalesced onto that run, so both
31//!   callers await and share its result rather than paginating the same history
32//!   twice.
33//!
34//! A running request is never preempted, so a queued request only starts once
35//! the current run for its room ends. Consumers bound their own runs with a
36//! [`BackPaginationRequest::max_batches`] cap or a stop predicate that fires
37//! once they have what they want.
38
39use 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
56/// Number of events requested per background pagination batch.
57pub(crate) const BATCH_SIZE: u16 = 30;
58
59/// Priority of a [`BackPaginationRequest`], relative to the others in the
60/// queue.
61#[allow(dead_code)]
62#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
63pub(crate) enum Priority {
64    /// Lowest priority: will run slowly when no higher priority requests are
65    /// pending.
66    Low,
67    /// Default priority. Higher than [`Self::Low`] and lower than
68    /// [`Self::High`].
69    Normal,
70    /// Highest priority: will run before any other pending request.
71    High,
72}
73
74/// When a single room's back-pagination run should stop: a predicate over each
75/// freshly loaded batch.
76///
77/// Each consumer supplies its own (a suitable latest-event candidate is loaded,
78/// a read receipt's target event id shows up, the batch is old enough for the
79/// search backfill). A predicate that never fires leaves the run bounded only
80/// by `max_batches`, the start of the timeline, or cancellation.
81pub(crate) type StopCondition = Box<dyn FnMut(&BackPaginationOutcome) -> ControlFlow<()> + Send>;
82
83/// A request to back-paginate one room, enqueued on the
84/// [`BackPaginationQueue`].
85pub(crate) struct BackPaginationRequest {
86    /// The room to back-paginate.
87    pub room_id: OwnedRoomId,
88    /// Scheduling priority.
89    pub priority: Priority,
90    /// When to stop.
91    pub stop: StopCondition,
92    /// Number of events to request per pagination.
93    pub batch_size: u16,
94    /// Maximum number of paginations for this request (`None` = unbounded).
95    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/// Why a single room's back-pagination run ended.
111#[derive(Clone, Copy, Debug, PartialEq, Eq)]
112pub(crate) enum BackPaginationStopReason {
113    /// Reached the start of the room's timeline; nothing more to load.
114    ReachedTimelineStart,
115    /// The request's [`StopCondition`] was met.
116    StopConditionMet,
117    /// Hit `max_batches` without satisfying the stop condition or reaching the
118    /// start of the timeline. More work likely remains; safe to retry later.
119    BatchLimitReached,
120    /// A pagination returned no events (e.g. a gap with no token to resolve it
121    /// yet). Not an error; there's simply nothing to load right now.
122    NoDataAvailable,
123    /// Setting up or running the pagination failed.
124    Failed,
125    /// The request was cancelled.
126    Cancelled,
127    /// The room isn't worth back-paginating: it's a space, which has no
128    /// messages to load, or the user isn't joined to it.
129    Skipped,
130}
131
132/// The result of running a single [`BackPaginationRequest`] to completion.
133#[allow(dead_code)]
134#[derive(Clone, Copy, Debug)]
135pub(crate) struct BackPaginationRunResult {
136    /// Why the run ended.
137    pub reason: BackPaginationStopReason,
138}
139
140/// Identifies a coalescable run i.e. a room back-paginated at a given priority.
141type RequestCoalescingKey = (OwnedRoomId, Priority);
142
143/// A handle to an enqueued [`BackPaginationRequest`]. Dropping the last handle
144/// for a request cancels it.
145pub(crate) struct BackPaginationHandle {
146    /// Cancels the request once every handle sharing it is dropped, unless
147    /// disarmed by [`BackPaginationHandle::detach`].
148    guard: Arc<DropGuard>,
149    // Only read by `join`, which the consumers awaiting a result are landing with.
150    #[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    /// Let the request run to completion instead of cancelling it, for callers
163    /// that don't care about its result.
164    ///
165    /// Only the last handle for a request can disarm its cancellation; while
166    /// other handles are alive they keep the request running anyway.
167    pub(crate) fn detach(self) {
168        if let Some(guard) = Arc::into_inner(self.guard) {
169            guard.disarm();
170        }
171    }
172
173    /// Await the request's completion, returning why it ended.
174    #[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/// A bounded queue of back-pagination requests, ordered by priority.
185#[derive(Clone)]
186pub struct BackPaginationQueue {
187    inner: Arc<BackPaginationQueueInner>,
188}
189
190struct BackPaginationQueueInner {
191    sender: mpsc::UnboundedSender<SchedulerEvent>,
192    /// The cancellation shared by all the handles of a coalescing key, so a run
193    /// is only cancelled once every caller waiting on it has dropped its
194    /// handle.
195    ///
196    /// This only tracks handle lifetimes; the scheduler stays authoritative for
197    /// whether a request actually coalesces onto an existing run.
198    cancellations: Mutex<HashMap<RequestCoalescingKey, SharedCancellation>>,
199    _task: matrix_sdk_base::task_monitor::BackgroundTaskHandle,
200}
201
202/// The cancellation of a single coalescable run.
203struct SharedCancellation {
204    /// Handed to the run, cancelled when the last `guard` below is dropped.
205    token: CancellationToken,
206    /// Weak, because the guard is owned by the handles: once they're all gone
207    /// the token fires and this entry is stale.
208    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    /// Create the queue and spawn its executor task.
220    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        // The scheduler holds a sender of its own, to be handed to the runs it
228        // spawns, so the channel never closes on its own: the task runs until
229        // the queue is dropped, which aborts it.
230        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    /// Enqueue a new request returning a handle to await it. A request for a
247    /// room already queued or running at the same priority is coalesced onto
248    /// that run rather than starting a second one.
249    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
268/// The cancellation shared by every handle for `key`, created if this is the
269/// first live one.
270///
271/// Requests that coalesce onto the same run must not be able to cancel each
272/// other, so they all hold a clone of one guard and the run is only cancelled
273/// once the last of them is dropped.
274fn 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    // Forget the keys whose handles are all gone, while we're holding the lock.
287    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/// An error happening while interacting with the [`BackPaginationQueue`].
298#[derive(Debug, thiserror::Error)]
299pub(crate) enum BackPaginationQueueError {
300    /// The queue's executor isn't running anymore, so no new request can be
301    /// enqueued.
302    #[error("the back-pagination queue executor is not running")]
303    ShutDown,
304}
305
306/// Everything the scheduler reacts to, on a single channel so it can wait on
307/// one `recv_many` rather than selecting over two receivers.
308enum SchedulerEvent {
309    /// A caller enqueued a new request.
310    Submitted(SubmittedRequest),
311    /// A run finished: free its room and hand its result to every waiter.
312    Finished(RequestCoalescingKey, BackPaginationRunResult),
313}
314
315/// A request as it arrives on the queue's channel, before the scheduler has
316/// assigned it a sequence number or decided whether to coalesce it.
317struct SubmittedRequest {
318    request: BackPaginationRequest,
319    token: CancellationToken,
320    completion: oneshot::Sender<BackPaginationRunResult>,
321}
322
323/// A [`BackPaginationRequest`] admitted to the scheduler, holding the sequence
324/// number and request details.
325struct PendingRequest {
326    request: BackPaginationRequest,
327    /// Insertion order, assigned by the scheduler, for FIFO within a priority.
328    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        // Higher priority first, then earlier `seq` first
349        self.request.priority.cmp(&other.request.priority).then_with(|| other.seq.cmp(&self.seq))
350    }
351}
352
353/// The executor schedules requests by priority, bounded concurrency at one run
354/// per room at a time.
355#[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    // The tasks of the currently running requests, keyed by room: also the set
366    // of rooms that can't take another run right now. Dropping the scheduler
367    // aborts all of them.
368    let mut active_requests: HashMap<OwnedRoomId, AbortOnDrop<()>> = HashMap::new();
369    let mut next_seq: u64 = 0;
370
371    // Completion senders for every request the scheduler knows about (queued or
372    // running), keyed by room and priority. A duplicate request coalesces onto
373    // the existing run by adding its completion sender here rather than
374    // starting a second run. When the run finishes every waiter for the key
375    // receives the same result.
376    let mut waiters: HashMap<RequestCoalescingKey, Vec<oneshot::Sender<BackPaginationRunResult>>> =
377        HashMap::new();
378
379    // At most `max_concurrent` runs are in flight, so there's never a reason to
380    // drain more than that in one go.
381    let mut events = Vec::with_capacity(max_concurrent);
382
383    loop {
384        // Schedule as many pending requests as the concurrency budget allows,
385        // never starting a second run for a room that's already running one.
386        schedule(
387            &event_cache,
388            &mut pending_requests,
389            &mut active_requests,
390            max_concurrent,
391            &sender,
392        );
393
394        // Unreachable while this task holds `sender`, but guards against a hot
395        // loop if that ever stops being true.
396        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                    // Fan the single run's result out to every coalesced
426                    // waiter.
427                    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
438/// Coalesce a new caller's `completion` onto an existing run for `key`, or
439/// admit it as the first waiter of a new run.
440///
441/// Requests sharing a [`RequestCoalescingKey`] are functionally
442/// interchangeable, so only one run is needed; extra callers wait on the same
443/// result. Different priorities for the same room have different keys, so they
444/// never coalesce. A busy room must still let a higher-priority request wait
445/// its turn.
446fn 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
463/// Pop and spawn every currently-schedulable request.
464fn 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            // The scheduler owns the completion senders (for coalescing), so
488            // hand it the result to fan out to every waiter for this key.
489            let _ = sender.send(SchedulerEvent::Finished(key, result));
490        });
491
492        active_requests.insert(room_id, task.abort_on_drop());
493    }
494}
495
496/// Pick the requests that can start right now, highest priority first: bounded
497/// by `max_concurrent` total in flight, and never two runs for the same room.
498///
499/// Requests popped but not yet runnable (their room is busy) are pushed back
500/// onto the heap.
501// Generic over the map's value type so the scheduling tests don't need a
502// runtime to build a task handle; only the keys matter here.
503fn 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            // This room is busy, try it again next round.
521            skipped.push(request);
522            continue;
523        }
524
525        picked.push(request);
526    }
527
528    pending_requests.extend(skipped);
529
530    picked
531}
532
533/// Back-paginate one room until its [`StopCondition`], the start of the
534/// timeline, the batch budget, or cancellation.
535#[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    // Cancelled while still queued, nothing to do.
542    if token.is_cancelled() {
543        return BackPaginationRunResult { reason: BackPaginationStopReason::Cancelled };
544    }
545
546    // Grab an owned `RoomPagination`, dropping the caches guard immediately so
547    // we don't hold the room lock across network paginations.
548    //
549    // Also keep a subscription to the room for the whole run: dropping it when
550    // this function returns triggers an auto-shrink of the room's linked chunk
551    // if nobody else is subscribed, so backfilled events don't linger in
552    // memory.
553    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        // Reaching the start of the timeline can still come with a last batch
597        // of events, so let the stop condition see it before ending the run.
598        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    /// Build a queued request for a room, at a priority, with an insertion seq.
636    fn queued(room_id: ruma::OwnedRoomId, priority: Priority, seq: u64) -> PendingRequest {
637        PendingRequest {
638            request: BackPaginationRequest {
639                room_id,
640                priority,
641                // Scheduling tests never evaluate the stop condition.
642                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    /// `next_runnable` serves highest priority first, then FIFO within a
652    /// priority.
653    #[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        // Push out of priority order, with monotonic seqs.
661        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        // High first (b before d by FIFO), then Normal, then Low.
673        assert_eq!(picked, vec![b.to_owned(), d.to_owned(), c.to_owned(), a.to_owned()]);
674    }
675
676    /// `next_runnable` never returns more than `max_concurrent`.
677    #[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        // The third request stays queued.
691        assert_eq!(pending_requests.len(), 1);
692    }
693
694    /// `next_runnable` won't start a room that's already active, nor two runs
695    /// for the same room in one pass.
696    #[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        // `a` is already running.
703        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)); // same room as active
707        pending_requests.push(queued(a.to_owned(), Priority::High, 1)); // and again
708        pending_requests.push(queued(b.to_owned(), Priority::Low, 2)); // a different room
709
710        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        // Only `b` runs; both `a` requests stay queued (a is busy).
716        assert_eq!(picked, vec![b.to_owned()]);
717        assert_eq!(pending_requests.len(), 2);
718    }
719
720    /// Every handle for a room + priority shares one cancellation, so a caller
721    /// dropping its handle can't cancel a run others are still waiting on.
722    #[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        // One run, so one token for both callers.
731        assert!(!first_token.is_cancelled());
732        drop(first_guard);
733        assert!(!first_token.is_cancelled());
734        assert!(!second_token.is_cancelled());
735
736        // The last handle to go cancels the run.
737        drop(second_guard);
738        assert!(first_token.is_cancelled());
739        assert!(second_token.is_cancelled());
740
741        // Once they're all gone the key is stale, so the next caller gets a
742        // fresh token rather than an already-cancelled one.
743        let (third_token, _third_guard) = cancellation_for(&cancellations, key);
744        assert!(!third_token.is_cancelled());
745
746        // A different priority for the same room is a different run.
747        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    /// The first request for a room + priority opens a new run; a later one at
755    /// the same key coalesces onto it (shares its waiter list); a different
756    /// priority for the same room opens its own run.
757    #[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        // First request at (a, Normal): opens a new run.
767        assert!(!try_coalesce(&mut waiters, &normal, completion()));
768        assert_eq!(waiters[&normal].len(), 1);
769
770        // Second request at the same key: coalesces onto it.
771        assert!(try_coalesce(&mut waiters, &normal, completion()));
772        assert_eq!(waiters[&normal].len(), 2);
773
774        // Same room, different priority: a separate run, not a coalesce.
775        assert!(!try_coalesce(&mut waiters, &high, completion()));
776        assert_eq!(waiters.len(), 2);
777        assert_eq!(waiters[&high].len(), 1);
778    }
779}