hyperlight_common/virtq/

mod.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 The Hyperlight Authors.

//! Packed Virtqueue Implementation
//!
//! This module provides a high-level API for virtio packed virtqueues, built on top of
//! the lower-level ring primitives. It implements the VIRTIO 1.1+ packed ring format
//! with proper memory ordering and event suppression support.
//!
//! # Architecture
//!
//! The implementation is split into layers:
//!
//! - **High-level API** ([`VirtqProducer`], [`VirtqConsumer`]): Manages buffer allocation,
//!   chain lifecycle, and notification decisions. This is the recommended API
//!   for most use cases.
//!
//! - **Ring primitives** ([`RingProducer`], [`RingConsumer`]): Low-level descriptor ring
//!   operations with explicit buffer chain management. Use this when you need full control
//!   over buffer layouts or custom allocation strategies.
//!
//! - **Descriptor and event types** ([`Descriptor`], [`EventSuppression`]): Raw virtio
//!   data structures for direct memory manipulation.
//!
//! # Quick Start
//!
//! ## Single Readable/Writable Chain
//!
//! ```ignore
//! // Producer (driver) side - build and submit a send chain
//! let mut chain = producer.chain()
//!     .readable(64)
//!     .writable(128)
//!     .build()?;
//! chain.write_all(b"request data")?;
//! let token = producer.submit(chain)?;
//! // ... wait for notification ...
//! if let Some(used) = producer.poll()? {
//!     match used {
//!         UsedChain::Data(_, segments) => process(segments),
//!         UsedChain::Ack(_) => {}
//!     }
//! }
//!
//! // Consumer (device) side - receive a chain and reply/ack it
//! if let Some((chain, reply)) = consumer.poll(max_recv_len)? {
//!     let request = chain.to_bytes();
//!     match reply {
//!         ReplyChain::Writable(mut wc) => {
//!             let response = handle(request);
//!             wc.write_all(&response)?;
//!             consumer.complete(wc)?;
//!         }
//!         ReplyChain::Ack(ack) => {
//!             consumer.complete(ack)?;
//!         }
//!     }
//! }
//!
//! // Multiple pending completions (no borrow on consumer)
//! let mut pending = Vec::new();
//! while let Some((chain, reply)) = consumer.poll(max_recv_len)? {
//!     pending.push((process(chain), reply));
//! }
//! for (result, reply) in pending {
//!     consumer.complete(reply)?;
//! }
//! ```
//!
//! ## Multiple Chains
//!
//! Each submit checks event suppression and notifies independently. Use
//! [`VirtqProducer::batch`] when a higher-level protocol wants to publish
//! multiple chains and kick the queue once.
//!
//! ```ignore
//! let mut batch = producer.batch();
//! for data in entries {
//!     let mut chain = batch.chain()
//!         .readable(data.len())
//!         .writable(64)
//!         .build()?;
//!     chain.write_all(data)?;
//!     batch.submit(chain)?;
//! }
//! batch.finish()?;
//! ```
//!
//! ## Completion Batching with Event Suppression
//!
//! To receive a single notification when multiple requests complete:
//!
//! ```ignore
//! // Submit chains
//! for data in entries {
//!     let mut chain = producer.chain()
//!         .readable(data.len())
//!         .writable(64)
//!         .build()?;
//!     chain.write_all(data)?;
//!     producer.submit(chain)?;
//! }
//!
//! // Tell device: "notify me only after completing past this cursor"
//! let cursor = producer.used_cursor();
//! producer.set_used_suppression(SuppressionKind::Descriptor(cursor))?;
//!
//! // Wait for single notification, then drain all responses
//! producer.drain(|used| {
//!     if let UsedChain::Data(token, data) = used {
//!         handle_response(token, data);
//!     }
//! })?;
//! ```
//!
//! # Event Suppression
//!
//! Both sides can control when they want to be notified using [`SuppressionKind`]:
//!
//! - [`SuppressionKind::Enable`]: Always notify (default, lowest latency)
//! - [`SuppressionKind::Disable`]: Never notify (polling mode, lowest overhead)
//! - [`SuppressionKind::Descriptor`]: Notify at specific ring position (batching)
//!
//! See [`VirtqProducer::set_used_suppression`] and [`VirtqConsumer::set_avail_suppression`].
//!
//! # Low-Level API
//!
//! For advanced use cases, the ring module exposes lower-level primitives:
//!
//! - [`RingProducer`] / [`RingConsumer`]: Direct ring access with [`BufferChain`] submission
//! - [`BufferChainBuilder`]: Construct scatter-gather buffer lists
//! - [`RingCursor`]: Track ring positions for event suppression
//!
//! Example using low-level API:
//!
//! ```ignore
//! let chain = BufferChainBuilder::new()
//!     .readable(header_addr, header_len)
//!     .readable(data_addr, data_len)
//!     .writable(response_addr, response_len)
//!     .build()?;
//!
//! let result = ring_producer.submit_available_with_notify(&chain)?;
//! if result.notify {
//!     kick_device();
//! }
//! ```

mod access;
mod buffer;
mod consumer;
mod desc;
mod event;
pub mod msg;
mod pool;
mod producer;
mod ring;

#[cfg(all(test, loom))]
mod concurrency;

use core::num::NonZeroU16;

pub use access::*;
pub use buffer::*;
pub use consumer::*;
pub use desc::*;
pub use event::*;
pub use pool::*;
pub use producer::*;
pub use ring::*;
use thiserror::Error;

/// A trait for notifying the consumer about virtqueue events.
pub trait Notifier {
    fn notify(&self, stats: QueueStats);
}

/// Errors that can occur in the virtqueue operations.
#[derive(Error, Debug)]
pub enum VirtqError {
    #[error("Ring error: {0}")]
    RingError(RingError),
    #[error("Allocation error: {0}")]
    Alloc(AllocError),
    #[error("Ring or pool temporarily full")]
    Backpressure,
    #[error("Allocation exceeds pool capacity")]
    OutOfMemory,
    #[error("Invalid chain received")]
    BadChain,
    #[error("Payload data too large: received {recv} bytes, limit {limit} bytes")]
    PayloadTooLarge { recv: usize, limit: usize },
    #[error("Reply data too large for allocated buffer")]
    ReplyTooLarge,
    #[error("Internal state error")]
    InvalidState,
    #[error("Memory write error")]
    MemoryWriteError,
    #[error("Memory read error")]
    MemoryReadError,
    #[error("No payload segment in this chain")]
    NoPayloadSegment,
}

impl VirtqError {
    /// Check if this error is transient or unrecoverable.
    #[inline(always)]
    pub fn is_transient(&self) -> bool {
        matches!(self, Self::Backpressure)
    }
}

impl From<RingError> for VirtqError {
    fn from(e: RingError) -> Self {
        match e {
            RingError::WouldBlock => Self::Backpressure,
            other => Self::RingError(other),
        }
    }
}

impl From<AllocError> for VirtqError {
    fn from(e: AllocError) -> Self {
        match e {
            AllocError::NoSpace => Self::Backpressure,
            AllocError::OutOfMemory => Self::OutOfMemory,
            other => Self::Alloc(other),
        }
    }
}

/// Layout of a packed virtqueue ring in shared memory.
///
/// Describes the memory addresses for the descriptor table and event suppression
/// structures. Use [`from_base`](Self::from_base) to compute the layout from a
/// base address, or [`query_size`](Self::query_size) to determine memory requirements.
///
/// # Memory Layout
///
/// The packed ring consists of:
/// 1. Descriptor table: `num_descs` × 16 bytes, aligned to 16 bytes
/// 2. Driver event suppression: 4 bytes, aligned to 4 bytes
/// 3. Device event suppression: 4 bytes, aligned to 4 bytes
#[derive(Clone, Copy, Debug)]
pub struct Layout {
    /// Packed ring descriptor table base in shared memory.
    desc_table_addr: u64,
    /// Number of descriptors (ring size, must be power of 2).
    desc_table_len: u16,
    /// Driver-written event suppression area in shared memory.
    drv_evt_addr: u64,
    /// Device-written event suppression area in shared memory.
    dev_evt_addr: u64,
}

#[inline]
const fn align_up(val: usize, align: usize) -> usize {
    val.next_multiple_of(align)
}

impl Layout {
    /// Create a Layout from a base address and number of descriptors.
    ///
    /// The base address must be aligned to `Descriptor::ALIGN`.
    /// The number of descriptors must be a power of 2.
    /// The memory region starting at `base` must be at least `Layout::query_size(num_descs)` bytes.
    ///
    /// # Safety
    /// - `base` must be valid for `Layout::query_size(num_descs)` bytes.
    /// - `base` must be aligned to `Descriptor::ALIGN`.
    /// - Memory must remain valid for the lifetime of the ring.
    pub const unsafe fn from_base(base: u64, num_descs: NonZeroU16) -> Result<Self, RingError> {
        let num_descs = num_descs.get() as usize;
        if !num_descs.is_power_of_two() {
            return Err(RingError::InvalidLayout);
        }

        if !base.is_multiple_of(Descriptor::ALIGN as u64) {
            return Err(RingError::InvalidLayout);
        }

        if base
            .checked_add(Layout::query_size(num_descs) as u64)
            .is_none()
        {
            return Err(RingError::InvalidLayout);
        }

        let desc_size = num_descs * Descriptor::SIZE;
        let event_size = EventSuppression::SIZE;
        let event_align = EventSuppression::ALIGN;

        let drv_evt_offset = align_up(desc_size, event_align);
        let dev_evt_offset = align_up(drv_evt_offset + event_size, event_align);

        Ok(Self {
            desc_table_addr: base,
            desc_table_len: num_descs as u16,
            drv_evt_addr: base + drv_evt_offset as u64,
            dev_evt_addr: base + dev_evt_offset as u64,
        })
    }

    /// Packed ring descriptor table base in shared memory.
    pub const fn desc_table_addr(&self) -> u64 {
        self.desc_table_addr
    }

    /// Number of descriptors in the ring.
    pub const fn desc_table_len(&self) -> u16 {
        self.desc_table_len
    }

    /// Driver-written event suppression area in shared memory.
    pub const fn drv_evt_addr(&self) -> u64 {
        self.drv_evt_addr
    }

    /// Device-written event suppression area in shared memory.
    pub const fn dev_evt_addr(&self) -> u64 {
        self.dev_evt_addr
    }

    /// Calculate the memory size needed for a ring with `num_descs` descriptors,
    /// accounting for alignment requirements.
    pub const fn query_size(num_descs: usize) -> usize {
        let desc_size = num_descs * Descriptor::SIZE;
        let event_size = EventSuppression::SIZE;
        let event_align = EventSuppression::ALIGN;

        // desc table at offset 0, then aligned events
        let drv_evt_offset = align_up(desc_size, event_align);
        let dev_evt_offset = align_up(drv_evt_offset + event_size, event_align);

        dev_evt_offset + event_size
    }
}

/// Statistics about the current virtqueue state.
///
/// Provided to the [`Notifier`] when sending notifications, allowing
/// the notifier to make decisions based on queue pressure.
#[derive(Debug, Clone, Copy)]
pub struct QueueStats {
    /// Number of free descriptor slots available.
    pub num_free: usize,
    /// Number of descriptors currently in-flight (submitted but not completed).
    pub num_inflight: usize,
}

/// Event suppression mode for controlling when notifications are sent.
///
/// This configures when the other side should signal (interrupt/kick) us
/// about new data. Used to optimize batching and reduce interrupt overhead.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SuppressionKind {
    /// Always signal after each operation (default behavior).
    Enable,
    /// Never signal.
    Disable,
    /// Signal only when reaching a specific descriptor position.
    Descriptor(RingCursor),
}

/// A token representing a sent chain in the virtqueue.
///
/// Tokens uniquely identify in-flight requests and are used to correlate
/// requests with their responses.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Token {
    /// Monotonically increasing generation counter.
    pub seq: u32,
    /// Descriptor ID this token maps to.
    pub id: u16,
}

impl From<BufferElement> for Allocation {
    fn from(value: BufferElement) -> Self {
        Allocation {
            addr: value.addr,
            len: value.len as usize,
        }
    }
}

const _: () = {
    #[allow(clippy::panic)]
    #[allow(clippy::unwrap_used)]
    const fn verify_layout(num_descs: usize) {
        let base = 0x1000u64;

        // Safety: base is aligned and we're only checking layout math
        let layout =
            match unsafe { Layout::from_base(base, NonZeroU16::new(num_descs as u16).unwrap()) } {
                Ok(l) => l,
                Err(_) => panic!("from_base failed"),
            };

        let expected_size = Layout::query_size(num_descs);

        assert!(layout.desc_table_addr() == base);
        assert!(layout.desc_table_len() as usize == num_descs);
        assert!(
            layout
                .drv_evt_addr()
                .is_multiple_of(EventSuppression::ALIGN as u64)
        );
        assert!(
            layout
                .dev_evt_addr()
                .is_multiple_of(EventSuppression::ALIGN as u64)
        );

        // Events don't overlap with descriptor table
        let desc_end = base + (num_descs * Descriptor::SIZE) as u64;
        assert!(layout.drv_evt_addr() >= desc_end);
        assert!(layout.dev_evt_addr() >= layout.drv_evt_addr() + EventSuppression::SIZE as u64);

        // Total size from query_size covers entire layout
        let layout_end = layout.dev_evt_addr() + EventSuppression::SIZE as u64;
        assert!(base + expected_size as u64 == layout_end);
    }

    unsafe {
        assert!(Layout::from_base(u64::MAX, NonZeroU16::new(1).unwrap()).is_err());
    }

    verify_layout(1);
    verify_layout(2);
    verify_layout(4);
    verify_layout(8);
    verify_layout(16);
    verify_layout(32);
    verify_layout(64);
    verify_layout(128);
    verify_layout(256);
    verify_layout(512);
    verify_layout(1024);
};

/// Shared test utilities for virtqueue tests.
#[cfg(test)]
pub(crate) mod test_utils {
    use alloc::collections::BTreeMap;
    use alloc::sync::Arc;
    use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
    use std::sync::Mutex;

    use super::*;
    use crate::virtq::ring::tests::{OwnedRing, TestMem};

    /// Simple notifier that tracks notification count.
    #[derive(Debug, Clone)]
    pub(crate) struct TestNotifier {
        pub(crate) count: Arc<AtomicUsize>,
    }

    impl TestNotifier {
        pub(crate) fn new() -> Self {
            Self {
                count: Arc::new(AtomicUsize::new(0)),
            }
        }

        pub(crate) fn notification_count(&self) -> usize {
            self.count.load(Ordering::Relaxed)
        }
    }

    impl Notifier for TestNotifier {
        fn notify(&self, _stats: QueueStats) {
            self.count.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Simple test buffer pool that allocates from a range.
    #[derive(Clone)]
    pub(crate) struct TestPool {
        base: u64,
        next: Arc<AtomicU64>,
        size: usize,
        max_alloc_len: usize,
        allocations: Arc<Mutex<BTreeMap<u64, usize>>>,
    }

    impl TestPool {
        pub(crate) fn new(base: u64, size: usize) -> Self {
            Self {
                base,
                next: Arc::new(AtomicU64::new(base)),
                size,
                max_alloc_len: usize::MAX,
                allocations: Arc::new(Mutex::new(BTreeMap::new())),
            }
        }

        pub(crate) fn new_with_max_alloc_len(base: u64, size: usize, max_alloc_len: usize) -> Self {
            Self {
                base,
                next: Arc::new(AtomicU64::new(base)),
                size,
                max_alloc_len,
                allocations: Arc::new(Mutex::new(BTreeMap::new())),
            }
        }
    }

    impl BufferProvider for TestPool {
        fn max_alloc_len(&self) -> usize {
            self.max_alloc_len
        }

        fn alloc(&self, len: usize) -> Result<Allocation, AllocError> {
            if len == 0 {
                return Err(AllocError::InvalidArg);
            }

            let addr = self.next.fetch_add(len as u64, Ordering::Relaxed);
            let end = addr + len as u64;
            if end > self.base + self.size as u64 {
                return Err(AllocError::NoSpace);
            }
            self.allocations
                .lock()
                .expect("poisoned mutex")
                .insert(addr, len);
            Ok(Allocation { addr, len })
        }

        fn dealloc(&self, addr: u64) -> Result<(), AllocError> {
            self.allocations
                .lock()
                .expect("poisoned mutex")
                .remove(&addr)
                .map(|_| ())
                .ok_or(AllocError::InvalidFree(addr, 0))
        }
    }

    type TestProducer = VirtqProducer<TestMem, TestNotifier, TestPool>;
    type TestConsumer = VirtqConsumer<TestMem, TestNotifier>;

    /// Create test infrastructure: a producer, consumer, and notifier backed
    /// by the supplied [`OwnedRing`].
    pub(crate) fn make_test_producer(
        ring: &OwnedRing,
    ) -> (TestProducer, TestConsumer, TestNotifier) {
        let layout = ring.layout();
        let mem = ring.mem();

        // Pool needs to be in memory accessible via mem - use memory after ring layout
        let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100;
        let pool = TestPool::new(pool_base, 0x8000);
        let notifier = TestNotifier::new();

        let producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool);
        let consumer = VirtqConsumer::new(layout, mem, notifier.clone());

        (producer, consumer, notifier)
    }
}

#[cfg(test)]
mod tests {
    use alloc::sync::Arc;
    use core::sync::atomic::{AtomicUsize, Ordering};

    use super::*;
    use crate::virtq::ring::tests::{TestMem, make_ring};
    use crate::virtq::test_utils::*;

    /// Helper: build and submit a readable+writable chain using the chain() builder.
    fn send_readwrite(
        producer: &mut VirtqProducer<TestMem, TestNotifier, TestPool>,
        entry_data: &[u8],
        used_cap: usize,
    ) -> Token {
        let mut se = producer
            .chain()
            .readable(entry_data.len())
            .writable(used_cap)
            .build()
            .unwrap();
        se.write_all(entry_data).unwrap();
        producer.submit(se).unwrap()
    }

    fn poll_received(
        consumer: &mut VirtqConsumer<TestMem, TestNotifier>,
    ) -> (RecvChain, ReplyChain<TestMem>) {
        consumer.poll(1024).unwrap().unwrap()
    }

    #[test]
    fn test_submit_notifies() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, notifier) = make_test_producer(&ring);

        let initial_count = notifier.notification_count();

        let token = send_readwrite(&mut producer, b"hello", 64);
        assert!(notifier.notification_count() > initial_count);

        let (recv, _reply) = poll_received(&mut consumer);
        assert_eq!(recv.token(), token);
    }

    #[test]
    fn test_multiple_submits() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        let tok1 = send_readwrite(&mut producer, b"request1", 64);
        let tok2 = send_readwrite(&mut producer, b"request2", 64);
        let tok3 = send_readwrite(&mut producer, b"request3", 64);

        // Consumer sees all requests
        for _ in 0..3 {
            let (_recv, reply) = poll_received(&mut consumer);
            consumer.complete(reply).unwrap();
        }

        // All completions available
        let used1 = producer.poll().unwrap().unwrap();
        let used2 = producer.poll().unwrap().unwrap();
        let used3 = producer.poll().unwrap().unwrap();
        assert!(
            [used1.token(), used2.token(), used3.token()].contains(&tok1)
                && [used1.token(), used2.token(), used3.token()].contains(&tok2)
                && [used1.token(), used2.token(), used3.token()].contains(&tok3)
        );
    }

    #[test]
    fn test_completion_batching_with_suppression() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        // Submit entries
        let tok1 = send_readwrite(&mut producer, b"req1", 64);
        let tok2 = send_readwrite(&mut producer, b"req2", 64);
        let tok3 = send_readwrite(&mut producer, b"req3", 64);

        // Set up reply batching via used suppression
        let cursor = producer.used_cursor();
        producer
            .set_used_suppression(SuppressionKind::Descriptor(cursor))
            .unwrap();

        // Consumer processes requests
        for _ in 0..3 {
            let (_recv, reply) = poll_received(&mut consumer);
            let ReplyChain::Writable(mut wc) = reply else {
                panic!("expected writable reply");
            };
            wc.write_all(b"used-data").unwrap();
            consumer.complete(wc).unwrap();
        }

        // Producer can drain all responses
        let mut responses = Vec::new();
        producer
            .drain(|reply| {
                responses.push(reply.token());
            })
            .unwrap();

        assert_eq!(responses.len(), 3);
        assert!(responses.contains(&tok1));
        assert!(responses.contains(&tok2));
        assert!(responses.contains(&tok3));
    }

    #[test]
    fn test_notifier_receives_context() {
        #[derive(Debug, Clone)]
        struct CtxNotifier {
            last_num_free: Arc<AtomicUsize>,
            last_num_inflight: Arc<AtomicUsize>,
            count: Arc<AtomicUsize>,
        }

        impl Notifier for CtxNotifier {
            fn notify(&self, stats: QueueStats) {
                self.last_num_free.store(stats.num_free, Ordering::Relaxed);
                self.last_num_inflight
                    .store(stats.num_inflight, Ordering::Relaxed);
                self.count.fetch_add(1, Ordering::Relaxed);
            }
        }

        let ring = make_ring(16);
        let layout = ring.layout();
        let mem = ring.mem();
        let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100;
        let pool = TestPool::new(pool_base, 0x8000);
        let notifier = CtxNotifier {
            last_num_free: Arc::new(AtomicUsize::new(0)),
            last_num_inflight: Arc::new(AtomicUsize::new(0)),
            count: Arc::new(AtomicUsize::new(0)),
        };

        let mut producer = VirtqProducer::new(layout, mem, notifier.clone(), pool);

        let mut se = producer.chain().readable(4).writable(32).build().unwrap();
        se.write_all(b"test").unwrap();
        producer.submit(se).unwrap();
        assert_eq!(notifier.count.load(Ordering::Relaxed), 1);
        assert!(notifier.last_num_inflight.load(Ordering::Relaxed) > 0);
    }

    #[test]
    fn test_chain_batch() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, notifier) = make_test_producer(&ring);

        let initial_count = notifier.notification_count();

        // First readable chain
        let mut se1 = producer.chain().readable(64).writable(128).build().unwrap();
        se1.write_all(b"first-ent").unwrap();
        let _tok1 = producer.submit(se1).unwrap();

        // Write-based recv
        let mut se2 = producer.chain().readable(64).writable(64).build().unwrap();
        se2.write_all(b"copy-ent").unwrap();
        let _tok2 = producer.submit(se2).unwrap();

        // Completion-only chain
        let se3 = producer.chain().writable(32).build().unwrap();
        let tok3 = producer.submit(se3).unwrap();

        // Each submit may notify independently
        assert!(notifier.notification_count() > initial_count);

        // Consumer sees all three entries
        let (recv1, reply1) = poll_received(&mut consumer);
        assert_eq!(recv1.to_bytes().as_ref(), b"first-ent");
        consumer.complete(reply1).unwrap();

        let (recv2, reply2) = poll_received(&mut consumer);
        assert_eq!(recv2.to_bytes().as_ref(), b"copy-ent");
        consumer.complete(reply2).unwrap();

        let (_recv3, reply3) = poll_received(&mut consumer);
        let ReplyChain::Writable(mut wc) = reply3 else {
            panic!("expected writable reply");
        };
        wc.write_all(b"resp").unwrap();
        consumer.complete(wc).unwrap();

        // Drain completions
        let _ = producer.poll().unwrap().unwrap();
        let _ = producer.poll().unwrap().unwrap();

        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), tok3);
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"resp");
    }

    #[test]
    fn test_chain_write_send() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        let mut se = producer.chain().readable(64).writable(128).build().unwrap();
        se.write_all(b"hello").unwrap();
        let token = producer.submit(se).unwrap();

        // Consumer sees the data
        let (recv, reply) = poll_received(&mut consumer);
        assert_eq!(recv.token(), token);
        assert_eq!(recv.to_bytes().as_ref(), b"hello");

        // Write response
        let ReplyChain::Writable(mut wc) = reply else {
            panic!("expected writable reply");
        };
        wc.write_all(b"world").unwrap();
        consumer.complete(wc).unwrap();
        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"world");
    }

    #[test]
    fn test_full_round_trip() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        // Send an recv
        let token = send_readwrite(&mut producer, b"round-trip-recv", 128);

        // Consumer receives and responds
        let (recv, reply) = poll_received(&mut consumer);
        assert_eq!(recv.token(), token);
        assert_eq!(recv.to_bytes().as_ref(), b"round-trip-recv");

        let ReplyChain::Writable(mut wc) = reply else {
            panic!("expected writable reply");
        };
        assert!(wc.capacity() >= 128);
        wc.write_all(b"round-trip-rsp").unwrap();
        consumer.complete(wc).unwrap();

        // Producer gets the reply
        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), token);
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"round-trip-rsp");
    }

    #[test]
    fn test_cancel_submits_zero_length() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        let token = send_readwrite(&mut producer, b"recv-data", 64);

        let (_recv, reply) = poll_received(&mut consumer);
        consumer.complete(reply).unwrap();

        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), token);
        assert_eq!(used.to_bytes().unwrap().len(), 0);
        assert!(used.to_bytes().unwrap().is_empty());
    }

    #[test]
    fn test_hold_reply_and_complete() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        let token = send_readwrite(&mut producer, b"deferred", 64);

        // Poll and hold the reply
        let (recv, reply) = poll_received(&mut consumer);
        assert_eq!(recv.token(), token);
        assert_eq!(recv.to_bytes().as_ref(), b"deferred");

        let ReplyChain::Writable(mut wc) = reply else {
            panic!("expected writable reply");
        };
        wc.write_all(b"deferred-used").unwrap();
        consumer.complete(wc).unwrap();

        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), token);
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"deferred-used");
    }

    #[test]
    fn test_concurrent_pending_replies() {
        let ring = make_ring(16);
        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);

        let tok1 = send_readwrite(&mut producer, b"first", 64);
        let tok2 = send_readwrite(&mut producer, b"second", 64);

        // Poll both
        let (recv1, reply1) = poll_received(&mut consumer);
        assert_eq!(recv1.token(), tok1);
        assert_eq!(recv1.to_bytes().as_ref(), b"first");

        let (recv2, reply2) = poll_received(&mut consumer);
        assert_eq!(recv2.token(), tok2);
        assert_eq!(recv2.to_bytes().as_ref(), b"second");

        // Complete second first (out of order)
        let ReplyChain::Writable(mut wc2) = reply2 else {
            panic!("expected writable");
        };
        wc2.write_all(b"resp2").unwrap();
        consumer.complete(wc2).unwrap();

        let ReplyChain::Writable(mut wc1) = reply1 else {
            panic!("expected writable");
        };
        wc1.write_all(b"resp1").unwrap();
        consumer.complete(wc1).unwrap();

        let used1 = producer.poll().unwrap().unwrap();
        let used2 = producer.poll().unwrap().unwrap();
        let mut responses: Vec<_> = vec![
            (used1.token(), used1.to_bytes().unwrap().to_vec()),
            (used2.token(), used2.to_bytes().unwrap().to_vec()),
        ];
        responses.sort_by_key(|(t, _)| t.seq);

        let expected_first = responses.iter().find(|(t, _)| *t == tok1).unwrap();
        let expected_second = responses.iter().find(|(t, _)| *t == tok2).unwrap();
        assert_eq!(&expected_first.1[..], b"resp1");
        assert_eq!(&expected_second.1[..], b"resp2");
    }

    /// Helper: submit a read-only chain (readable data, no writable reply).
    fn send_readonly(
        producer: &mut VirtqProducer<TestMem, TestNotifier, TestPool>,
        entry_data: &[u8],
    ) -> Token {
        let mut se = producer.chain().readable(entry_data.len()).build().unwrap();
        se.write_all(entry_data).unwrap();
        producer.submit(se).unwrap()
    }

    #[test]
    fn test_reclaim_frees_ring_slots() {
        let ring = make_ring(4);
        let (mut producer, mut consumer, _) = make_test_producer(&ring);

        // Fill the ring with ReadOnly entries
        send_readonly(&mut producer, b"a");
        send_readonly(&mut producer, b"b");
        send_readonly(&mut producer, b"c");
        send_readonly(&mut producer, b"d");

        // Ring is now full - next submit should fail with Backpressure
        let mut se = producer.chain().readable(1).build().unwrap();
        se.write_all(b"e").unwrap();
        let res = producer.submit(se);
        assert!(
            matches!(res, Err(VirtqError::Backpressure)),
            "expected Backpressure from full ring"
        );

        // Consumer acks all entries
        while let Some(result) = consumer.poll(1024).unwrap() {
            let (_, reply) = result;
            consumer.complete(reply).unwrap();
        }

        // Reclaim should free ring slots without losing data
        let count = producer.reclaim().unwrap();
        assert_eq!(count, 4, "expected 4 reclaimed entries");

        // Ring should have space now
        send_readonly(&mut producer, b"e");
    }

    #[test]
    fn test_reclaim_buffers_rw_completions() {
        let ring = make_ring(4);
        let (mut producer, mut consumer, _) = make_test_producer(&ring);

        // Submit a ReadWrite recv
        let tok = send_readwrite(&mut producer, b"request", 64);

        // Consumer processes and writes response
        let (_, reply) = poll_received(&mut consumer);
        let ReplyChain::Writable(mut wc) = reply else {
            panic!("expected writable");
        };
        wc.write_all(b"response-data").unwrap();
        consumer.complete(wc).unwrap();

        // Reclaim buffers the reply (doesn't discard it)
        let count = producer.reclaim().unwrap();
        assert_eq!(count, 1);

        // poll() should return the buffered reply
        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), tok);
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"response-data");
    }

    #[test]
    fn test_reclaim_discards_readonly_completions() {
        let ring = make_ring(8);
        let (mut producer, mut consumer, _) = make_test_producer(&ring);

        // Submit 3 entries: RO, RW, RO
        let _tok_ro1 = send_readonly(&mut producer, b"log1");
        let tok_rw = send_readwrite(&mut producer, b"call", 64);
        let _tok_ro2 = send_readonly(&mut producer, b"log2");

        // Consumer processes all 3
        let (_, reply1) = poll_received(&mut consumer);
        consumer.complete(reply1).unwrap(); // ack RO

        let (_, reply2) = poll_received(&mut consumer);
        let ReplyChain::Writable(mut wc) = reply2 else {
            panic!("expected writable");
        };
        wc.write_all(b"result").unwrap();
        consumer.complete(wc).unwrap(); // complete RW

        let (_, reply3) = poll_received(&mut consumer);
        consumer.complete(reply3).unwrap(); // ack RO

        // Reclaim all 3 - RO completions are discarded, only RW is buffered
        let count = producer.reclaim().unwrap();
        assert_eq!(count, 3);

        // poll() returns only the RW reply
        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), tok_rw);
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"result");

        // No more - RO completions were discarded
        assert!(producer.poll().unwrap().is_none());
    }

    #[test]
    fn test_reclaim_mixed_with_poll() {
        let ring = make_ring(8);
        let (mut producer, mut consumer, _) = make_test_producer(&ring);

        // Submit and complete 2 entries
        send_readonly(&mut producer, b"x");
        let tok_rw = send_readwrite(&mut producer, b"y", 64);

        let (_, reply1) = poll_received(&mut consumer);
        consumer.complete(reply1).unwrap();

        let (_, reply2) = poll_received(&mut consumer);
        let ReplyChain::Writable(mut wc) = reply2 else {
            panic!("expected writable");
        };
        wc.write_all(b"reply").unwrap();
        consumer.complete(wc).unwrap();

        // poll() consumes first recv directly from ring
        let used1 = producer.poll().unwrap().unwrap();
        assert!(matches!(used1, UsedChain::Ack(_)));

        // reclaim() buffers second recv
        let count = producer.reclaim().unwrap();
        assert_eq!(count, 1);

        // poll() returns the buffered one
        let used2 = producer.poll().unwrap().unwrap();
        assert_eq!(used2.token(), tok_rw);
        assert_eq!(used2.to_bytes().unwrap().as_ref(), b"reply");
    }

    /// reclaim + submit must not cause token collisions.
    #[test]
    fn test_reclaim_submit_no_token_collision() {
        let ring = make_ring(8);
        let (mut producer, mut consumer, _) = make_test_producer(&ring);

        // Submit and complete a ReadOnly recv
        let tok_old = send_readonly(&mut producer, b"log");

        let (_, reply) = poll_received(&mut consumer);
        consumer.complete(reply).unwrap();

        let count = producer.reclaim().unwrap();
        assert_eq!(count, 1);

        // Submit a new ReadWrite recv - may reuse the same descriptor ID
        let tok_new = send_readwrite(&mut producer, b"call", 64);

        // Tokens must differ even if the descriptor ID was recycled
        assert_ne!(
            tok_old, tok_new,
            "tokens must be unique across reclaim/submit cycles"
        );

        // Complete the ReadWrite recv
        let (_, reply) = poll_received(&mut consumer);
        let ReplyChain::Writable(mut wc) = reply else {
            panic!("expected writable");
        };
        wc.write_all(b"result").unwrap();
        consumer.complete(wc).unwrap();

        // Poll returns only the RW reply (RO was discarded by reclaim)
        let used = producer.poll().unwrap().unwrap();
        assert_eq!(used.token(), tok_new);
        assert_eq!(used.to_bytes().unwrap().as_ref(), b"result");

        // No stale RO reply in the queue
        assert!(producer.poll().unwrap().is_none());
    }

    /// Verify that repeated oneshot submit/reclaim cycles do not accumulate pending completions.
    #[test]
    fn test_reclaim_readonly_does_not_leak_pending() {
        let ring = make_ring(4);
        let (mut producer, mut consumer, _) = make_test_producer(&ring);

        for _ in 0..10 {
            // Fill the ring
            for _ in 0..4 {
                send_readonly(&mut producer, b"msg");
            }

            // Consumer acks all
            while let Some(result) = consumer.poll(1024).unwrap() {
                let (_, reply) = result;
                consumer.complete(reply).unwrap();
            }

            // Reclaim frees ring slots; empty completions are discarded
            let count = producer.reclaim().unwrap();
            assert_eq!(count, 4);

            // No completions should be buffered in pending
            assert!(
                producer.poll().unwrap().is_none(),
                "pending should be empty after reclaiming RO entries"
            );
        }
    }
}