qmux 0.5.0

//! Per-stream priority scheduling for outbound STREAM data.
//!
//! Replaces the bounded `mpsc` data channel the session used to interleave
//! STREAM frames purely by arrival order. Frames are bucketed by stream
//! priority (a `u8`, higher = sent first, matching the W3C `sendOrder`
//! convention) so a high-priority stream can jump ahead of a low-priority
//! backlog. Re-prioritization is retroactive and cheap: only the stream's
//! scheduling pointer moves between bands, never its queued frames, so the
//! per-stream FIFO order the receiver relies on (it appends STREAM data by
//! arrival, ignoring the wire offset) is always preserved.
//!
//! Control frames are *not* routed here — the session keeps its separate
//! unbounded `outbound_priority` channel and a `biased` select so control
//! always precedes any data this queue yields.

use std::{
    collections::{BTreeMap, HashMap, VecDeque},
    sync::{Arc, Mutex},
};

use tokio::sync::Notify;

use crate::{Error, Frame, StreamId};

/// Per-stream FIFO of pending frames plus the stream's current priority band.
struct StreamSlot {
    priority: u8,
    frames: VecDeque<Frame>,
}

struct Inner {
    /// Ready streams bucketed by priority. The MAX key is served first.
    /// A `StreamId` appears in at most one band at a time, and only while its
    /// slot has at least one queued frame.
    bands: BTreeMap<u8, VecDeque<StreamId>>,
    /// Per-stream queued frames + current priority.
    streams: HashMap<StreamId, StreamSlot>,
    /// Total queued frames across all streams (the capacity bound).
    len: usize,
    /// Set on session teardown; unblocks producers and the consumer.
    closed: bool,
}

impl Inner {
    /// Schedule `id` in `band` if it isn't already scheduled somewhere.
    ///
    /// Caller guarantees the slot exists and has at least one frame.
    fn arm(&mut self, id: StreamId, band: u8) {
        // A StreamId lives in at most one band. The slot's `priority` is the
        // single source of truth for which band it would be in, so checking the
        // band's queue for membership is unnecessary as long as callers only arm
        // a stream that was previously unscheduled (empty slot just created, or
        // a slot drained then refilled).
        self.bands.entry(band).or_default().push_back(id);
    }
}

/// A bounded, priority-aware queue of outbound STREAM frames shared between the
/// session's writer loop (consumer) and its `SendStream`s (producers).
///
/// Cloning shares the same underlying queue.
#[derive(Clone)]
pub struct PriorityQueue {
    inner: Arc<Mutex<Inner>>,
    /// Notified when a frame becomes available (or the queue is closed).
    non_empty: Arc<Notify>,
    /// Notified when capacity frees up (or the queue is closed).
    has_space: Arc<Notify>,
    capacity: usize,
}

impl PriorityQueue {
    /// Create a queue holding at most `capacity` frames total before `reserve`
    /// blocks. Outstanding [`Permit`]s count against the bound.
    pub fn new(capacity: usize) -> Self {
        Self {
            inner: Arc::new(Mutex::new(Inner {
                bands: BTreeMap::new(),
                streams: HashMap::new(),
                len: 0,
                closed: false,
            })),
            non_empty: Arc::new(Notify::new()),
            has_space: Arc::new(Notify::new()),
            capacity,
        }
    }

    /// Reserve one slot, blocking while the queue is full.
    ///
    /// Await this *before* building the frame: the returned [`Permit`] enqueues
    /// synchronously, so a caller can move bytes out of a source buffer and into
    /// the queue with no await in between. Awaiting capacity and committing the
    /// frame in one step instead (the obvious `push(frame).await`) leaves a
    /// window where a cancelled future has already taken the caller's bytes but
    /// not queued them, silently punching a hole in the stream.
    ///
    /// Cancel-safe: the slot is claimed synchronously right before returning, so
    /// dropping this future reserves nothing.
    pub async fn reserve(&self) -> Result<Permit, Error> {
        loop {
            // Register interest *before* checking, so a `notify_one` that fires
            // between our check and `.await` isn't lost.
            let notified = self.has_space.notified();
            {
                let mut inner = self.inner.lock().unwrap();
                if inner.closed {
                    return Err(Error::Closed);
                }
                if inner.len < self.capacity {
                    // Hold the slot against the capacity bound until the permit
                    // is used or dropped.
                    inner.len += 1;
                    return Ok(Permit {
                        queue: self.clone(),
                        armed: true,
                    });
                }
            }
            notified.await;
        }
    }

    /// Enqueue a frame synchronously, bypassing the capacity bound — for small,
    /// must-not-be-dropped control markers like a stream FIN, which would
    /// otherwise have to either block (impossible from a sync caller) or be
    /// detached to a task (racing reset/teardown). The frame still lands in the
    /// stream's band, after its data. Fails only if the queue is closed.
    pub fn push_now(&self, priority: u8, id: StreamId, frame: Frame) -> Result<(), Error> {
        let mut inner = self.inner.lock().unwrap();
        if inner.closed {
            return Err(Error::Closed);
        }
        self.push_locked(&mut inner, priority, id, frame);
        Ok(())
    }

    fn push_locked(&self, inner: &mut Inner, priority: u8, id: StreamId, frame: Frame) {
        self.enqueue_locked(inner, priority, id, frame);
        inner.len += 1;
    }

    /// Append `frame` to `id`'s FIFO without touching `len`. Callers own the
    /// capacity accounting: `push_now` bumps it after the fact, while a
    /// [`Permit`] already reserved its slot in `reserve`.
    fn enqueue_locked(&self, inner: &mut Inner, priority: u8, id: StreamId, frame: Frame) {
        match inner.streams.get_mut(&id) {
            Some(slot) => {
                // Already scheduled (its band points at `id`); just append.
                slot.frames.push_back(frame);
            }
            None => {
                let mut frames = VecDeque::new();
                frames.push_back(frame);
                inner.streams.insert(id, StreamSlot { priority, frames });
                inner.arm(id, priority);
            }
        }
        self.non_empty.notify_one();
    }

    /// Pop the next frame to send, honoring priority then round-robin fairness
    /// among equal-priority streams. Blocks until a frame is available or the
    /// queue is closed (returns `None` once closed and drained).
    pub async fn pop(&self) -> Option<Frame> {
        loop {
            let notified = self.non_empty.notified();
            {
                let mut inner = self.inner.lock().unwrap();
                if let Some(frame) = self.pop_locked(&mut inner) {
                    return Some(frame);
                }
                if inner.closed {
                    return None;
                }
            }
            notified.await;
        }
    }

    fn pop_locked(&self, inner: &mut Inner) -> Option<Frame> {
        // Highest band first.
        let (&band, queue) = inner.bands.iter_mut().next_back()?;
        let id = queue.pop_front().expect("scheduled band must be non-empty");
        if queue.is_empty() {
            inner.bands.remove(&band);
        }

        let slot = inner
            .streams
            .get_mut(&id)
            .expect("scheduled stream must have a slot");
        let frame = slot
            .frames
            .pop_front()
            .expect("scheduled slot must be non-empty");

        if slot.frames.is_empty() {
            // Drained: drop the slot so a future push re-arms it cleanly.
            inner.streams.remove(&id);
        } else {
            // Round-robin: re-arm at the back of its *current* band (which may
            // differ from `band` if set_priority moved it; but since it was at
            // the head of the max band, its priority equals `band`).
            let priority = slot.priority;
            inner.arm(id, priority);
        }

        inner.len -= 1;
        self.has_space.notify_one();
        Some(frame)
    }

    /// Retroactively re-prioritize a stream. Frames don't move — only the
    /// scheduling pointer relocates from the old band to `new`. No-op if the
    /// stream has no queued frames.
    pub fn set_priority(&self, id: StreamId, new: u8) {
        let mut inner = self.inner.lock().unwrap();
        let old = match inner.streams.get(&id) {
            Some(slot) => slot.priority,
            None => return,
        };
        if old == new {
            return;
        }

        // Relocate the scheduling pointer if the stream is currently scheduled.
        // It is scheduled iff it has frames (invariant), which it does here.
        if let Some(queue) = inner.bands.get_mut(&old) {
            if let Some(pos) = queue.iter().position(|&s| s == id) {
                queue.remove(pos);
                if queue.is_empty() {
                    inner.bands.remove(&old);
                }
                inner.bands.entry(new).or_default().push_back(id);
            }
        }

        if let Some(slot) = inner.streams.get_mut(&id) {
            slot.priority = new;
        }
    }

    /// Drop every queued frame for `id`, unscheduling it and freeing the
    /// capacity they held. Returns the number of queued STREAM payload bytes
    /// removed so the caller can return flow-control credit that never reached
    /// the wire. Used when a stream is reset: buffered STREAM data must not trail
    /// RESET_STREAM. Returns zero if the stream has nothing queued.
    pub fn remove(&self, id: StreamId) -> u64 {
        let mut inner = self.inner.lock().unwrap();
        let Some(slot) = inner.streams.remove(&id) else {
            return 0;
        };
        let removed = slot.frames.len();
        let removed_bytes = slot
            .frames
            .iter()
            .map(|frame| match frame {
                Frame::Stream(stream) => stream.data.len() as u64,
                _ => 0,
            })
            .sum();

        // Unschedule it from its band (it's scheduled iff it had frames, which it
        // did). `slot.priority` is the single source of truth for which band.
        if let Some(queue) = inner.bands.get_mut(&slot.priority) {
            if let Some(pos) = queue.iter().position(|&s| s == id) {
                queue.remove(pos);
            }
            if queue.is_empty() {
                inner.bands.remove(&slot.priority);
            }
        }

        inner.len -= removed;
        drop(inner);
        // Freed `removed` slots; wake that many blocked producers (mirrors the
        // per-slot notify in `pop_locked`).
        for _ in 0..removed {
            self.has_space.notify_one();
        }
        removed_bytes
    }

    /// Close the queue, unblocking any blocked producers and the consumer.
    pub fn close(&self) {
        {
            let mut inner = self.inner.lock().unwrap();
            inner.closed = true;
        }
        self.non_empty.notify_waiters();
        self.has_space.notify_waiters();
    }
}

/// One reserved slot in a [`PriorityQueue`], from [`PriorityQueue::reserve`].
///
/// Holds its slot against the capacity bound until [`Permit::send`] fills it or
/// it is dropped, which returns the slot. Separating the (awaited) reservation
/// from the (synchronous) send is what lets a caller commit bytes to the queue
/// without an await in between, so a cancelled write can't lose them.
pub struct Permit {
    queue: PriorityQueue,
    /// Cleared by `send`; a permit still armed at drop returns its slot.
    armed: bool,
}

impl Permit {
    /// Append `frame` to `id`'s FIFO using the reserved slot. Never blocks.
    ///
    /// Consumes the permit, so the reservation can't be reused or double-sent.
    /// Fails if the queue closed while the permit was outstanding: the consumer
    /// has already stopped popping, so enqueuing here would strand the frame and
    /// report success for data that will never be sent.
    pub fn send(mut self, priority: u8, id: StreamId, frame: Frame) -> Result<(), Error> {
        let mut inner = self.queue.inner.lock().unwrap();
        if inner.closed {
            // Release the lock before returning: `self` still counts as armed, so
            // its `Drop` re-locks to hand the slot back.
            drop(inner);
            return Err(Error::Closed);
        }
        self.armed = false;
        // `reserve` already counted this frame against `len`.
        self.queue.enqueue_locked(&mut inner, priority, id, frame);
        Ok(())
    }
}

impl Drop for Permit {
    fn drop(&mut self) {
        if !self.armed {
            return;
        }
        {
            let mut inner = self.queue.inner.lock().unwrap();
            inner.len -= 1;
        }
        self.queue.has_space.notify_one();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;

    use crate::proto::Stream;
    use crate::{StreamDir, StreamId};

    fn sid(index: u64) -> StreamId {
        StreamId::new(index, StreamDir::Uni, false)
    }

    fn frame(id: StreamId, tag: u8) -> Frame {
        Frame::Stream(Stream {
            id,
            offset: 0,
            data: Bytes::copy_from_slice(&[tag]),
            fin: false,
        })
    }

    fn tag_of(frame: &Frame) -> u8 {
        match frame {
            Frame::Stream(s) => s.data[0],
            _ => panic!("expected stream frame"),
        }
    }

    fn id_of(frame: &Frame) -> StreamId {
        match frame {
            Frame::Stream(s) => s.id,
            _ => panic!("expected stream frame"),
        }
    }

    /// Reserve and send in one step. The scheduling tests care about ordering and
    /// capacity, not the cancellation split, so they read better without it.
    async fn push(
        q: &PriorityQueue,
        priority: u8,
        id: StreamId,
        frame: Frame,
    ) -> Result<(), Error> {
        q.reserve().await?.send(priority, id, frame)
    }

    #[tokio::test]
    async fn higher_priority_first() {
        let q = PriorityQueue::new(8);
        let lo = sid(0);
        let hi = sid(1);

        push(&q, 10, lo, frame(lo, b'l')).await.unwrap();
        push(&q, 200, hi, frame(hi, b'h')).await.unwrap();

        // High priority drains first.
        assert_eq!(tag_of(&q.pop().await.unwrap()), b'h');
        assert_eq!(tag_of(&q.pop().await.unwrap()), b'l');
    }

    #[tokio::test]
    async fn equal_priority_round_robin() {
        let q = PriorityQueue::new(8);
        let a = sid(0);
        let b = sid(1);

        push(&q, 5, a, frame(a, 1)).await.unwrap();
        push(&q, 5, a, frame(a, 2)).await.unwrap();
        push(&q, 5, b, frame(b, 1)).await.unwrap();
        push(&q, 5, b, frame(b, 2)).await.unwrap();

        // Round-robin between equal-priority streams: a, b, a, b.
        assert_eq!(id_of(&q.pop().await.unwrap()), a);
        assert_eq!(id_of(&q.pop().await.unwrap()), b);
        assert_eq!(id_of(&q.pop().await.unwrap()), a);
        assert_eq!(id_of(&q.pop().await.unwrap()), b);
    }

    #[tokio::test]
    async fn per_stream_fifo_preserved() {
        let q = PriorityQueue::new(8);
        let a = sid(0);

        for i in 0..4u8 {
            push(&q, 5, a, frame(a, i)).await.unwrap();
        }
        for i in 0..4u8 {
            assert_eq!(tag_of(&q.pop().await.unwrap()), i);
        }
    }

    #[tokio::test]
    async fn set_priority_moves_pointer_not_frames() {
        let q = PriorityQueue::new(8);
        let lo = sid(0);
        let hi = sid(1);

        // Two frames each, lo enqueued first.
        push(&q, 10, lo, frame(lo, 1)).await.unwrap();
        push(&q, 10, lo, frame(lo, 2)).await.unwrap();
        push(&q, 20, hi, frame(hi, 1)).await.unwrap();

        // Promote lo above hi; its frames keep their order.
        q.set_priority(lo, 100);

        assert_eq!(id_of(&q.pop().await.unwrap()), lo);
        assert_eq!(tag_of(&q.pop().await.unwrap()), 2); // lo's second frame, in order
        assert_eq!(id_of(&q.pop().await.unwrap()), hi);
    }

    #[tokio::test]
    async fn set_priority_unknown_stream_is_noop() {
        let q = PriorityQueue::new(8);
        q.set_priority(sid(99), 50); // must not panic
    }

    #[tokio::test]
    async fn close_unblocks_pop() {
        let q = PriorityQueue::new(8);
        let q2 = q.clone();
        let handle = tokio::spawn(async move { q2.pop().await });
        tokio::task::yield_now().await;
        q.close();
        assert!(handle.await.unwrap().is_none());
    }

    #[tokio::test]
    async fn close_unblocks_push() {
        let q = PriorityQueue::new(1);
        let a = sid(0);
        push(&q, 5, a, frame(a, 1)).await.unwrap();

        let q2 = q.clone();
        let handle = tokio::spawn(async move { push(&q2, 5, sid(1), frame(sid(1), 2)).await });
        tokio::task::yield_now().await;
        q.close();
        assert!(matches!(handle.await.unwrap(), Err(Error::Closed)));
    }

    #[tokio::test]
    async fn backpressure_blocks_at_capacity() {
        let q = PriorityQueue::new(2);
        let a = sid(0);
        push(&q, 5, a, frame(a, 1)).await.unwrap();
        push(&q, 5, a, frame(a, 2)).await.unwrap();

        let q2 = q.clone();
        let pushing = tokio::spawn(async move { push(&q2, 5, a, frame(a, 3)).await });
        tokio::task::yield_now().await;
        assert!(!pushing.is_finished(), "push should block while full");

        // Draining one frame makes room.
        q.pop().await.unwrap();
        pushing.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn remove_drops_a_streams_queued_frames() {
        let q = PriorityQueue::new(8);
        let a = sid(0);
        let b = sid(1);

        push(&q, 5, a, frame(a, 1)).await.unwrap();
        push(&q, 5, a, frame(a, 2)).await.unwrap();
        push(&q, 5, b, frame(b, 9)).await.unwrap();

        // Reset `a`: its queued frames vanish, `b`'s survive.
        assert_eq!(q.remove(a), 2);

        assert_eq!(id_of(&q.pop().await.unwrap()), b);
        // Nothing left: the next pop blocks, so a closed-drain returns None.
        q.close();
        assert!(q.pop().await.is_none());
    }

    #[tokio::test]
    async fn remove_frees_capacity_for_blocked_producers() {
        let q = PriorityQueue::new(2);
        let a = sid(0);
        let b = sid(1);
        push(&q, 5, a, frame(a, 1)).await.unwrap();
        push(&q, 5, a, frame(a, 2)).await.unwrap();

        // Queue is full; a push for `b` blocks.
        let q2 = q.clone();
        let pushing = tokio::spawn(async move { push(&q2, 5, b, frame(b, 1)).await });
        tokio::task::yield_now().await;
        assert!(!pushing.is_finished(), "push should block while full");

        // Removing `a` frees both slots and wakes the blocked producer.
        assert_eq!(q.remove(a), 2);
        pushing.await.unwrap().unwrap();
        assert_eq!(id_of(&q.pop().await.unwrap()), b);
    }

    #[tokio::test]
    async fn remove_unknown_stream_is_noop() {
        let q = PriorityQueue::new(8);
        assert_eq!(q.remove(sid(99)), 0); // must not panic or underflow
    }

    #[tokio::test]
    async fn reserved_permit_holds_capacity() {
        let q = PriorityQueue::new(1);
        let a = sid(0);

        // An un-sent permit still occupies the queue's only slot.
        let permit = q.reserve().await.unwrap();

        let q2 = q.clone();
        let blocked = tokio::spawn(async move { q2.reserve().await.map(|_| ()) });
        tokio::task::yield_now().await;
        assert!(!blocked.is_finished(), "reserve should block while full");

        permit.send(5, a, frame(a, 1)).unwrap();
        // Still full: the permit became a real frame rather than freeing the slot.
        tokio::task::yield_now().await;
        assert!(
            !blocked.is_finished(),
            "sending a permit must not free its slot"
        );

        q.pop().await.unwrap();
        blocked.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn dropped_permit_returns_capacity() {
        let q = PriorityQueue::new(1);

        let q2 = q.clone();
        let blocked = tokio::spawn(async move { q2.reserve().await.map(|_| ()) });

        {
            let _permit = q.reserve().await.unwrap();
            tokio::task::yield_now().await;
            assert!(!blocked.is_finished(), "reserve should block while full");
            // Dropped here: the caller changed its mind (or was cancelled).
        }

        // The freed slot must wake the waiter, or a cancelled write wedges the
        // queue forever.
        blocked.await.unwrap().unwrap();
    }

    #[tokio::test]
    async fn cancelled_reserve_holds_nothing() {
        let q = PriorityQueue::new(1);
        let a = sid(0);

        // Fill the queue, then abandon a blocked reserve.
        push(&q, 5, a, frame(a, 1)).await.unwrap();
        tokio::select! {
            _ = q.reserve() => panic!("reserve should not succeed while full"),
            _ = tokio::task::yield_now() => {}
        }

        // The abandoned reserve must not have leaked the slot it was waiting for.
        q.pop().await.unwrap();
        push(&q, 5, a, frame(a, 2)).await.unwrap();
        assert_eq!(tag_of(&q.pop().await.unwrap()), 2);
    }

    /// A permit reserved before `close` must not enqueue after it: `pop` has
    /// already reported the queue drained and the consumer is gone, so the frame
    /// would sit there forever while the writer was told it was sent.
    #[tokio::test]
    async fn send_after_close_fails() {
        let q = PriorityQueue::new(4);
        let a = sid(0);

        let permit = q.reserve().await.unwrap();
        q.close();

        // The consumer has already given up on this queue.
        assert!(q.pop().await.is_none());
        assert!(matches!(permit.send(5, a, frame(a, 1)), Err(Error::Closed)));
        assert!(
            q.pop().await.is_none(),
            "a rejected send must queue nothing"
        );
    }

    /// Dropping a permit after close must not panic or underflow `len`.
    #[tokio::test]
    async fn drop_after_close_is_clean() {
        let q = PriorityQueue::new(4);
        let permit = q.reserve().await.unwrap();
        q.close();
        drop(permit);
        assert!(q.pop().await.is_none());
    }

    #[tokio::test]
    async fn close_unblocks_reserve() {
        let q = PriorityQueue::new(1);
        let a = sid(0);
        push(&q, 5, a, frame(a, 1)).await.unwrap();

        let q2 = q.clone();
        let handle = tokio::spawn(async move { q2.reserve().await.map(|_| ()) });
        tokio::task::yield_now().await;
        q.close();
        assert!(matches!(handle.await.unwrap(), Err(Error::Closed)));
    }
}