feat(virtq): add reset API

Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>

Author: andreiltd

src/hyperlight_common/src/virtq/consumer.rs

@@ -441,6 +441,12 @@ impl<M: MemOps + Clone, N: Notifier> VirtqConsumer<M, N> {
 
         Ok(Bytes::from(buf))
     }
+
+    /// Reset ring and inflight state to initial values.
+    pub fn reset(&mut self) {
+        self.inner.reset();
+        self.inflight.fill(None);
+    }
 }
 
 /// Parse a descriptor chain into entry/completion buffer elements.
@@ -630,4 +636,48 @@ mod tests {
         assert_eq!(data.as_ref(), b"abc");
         consumer.complete(completion).unwrap();
     }
+
+    #[test]
+    fn test_virtq_consumer_reset() {
+        let ring = make_ring(16);
+        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);
+
+        // Submit and poll (but do not complete)
+        let se = producer.chain().completion(16).build().unwrap();
+        producer.submit(se).unwrap();
+
+        let (_entry, completion) = consumer.poll(1024).unwrap().unwrap();
+        assert!(consumer.inflight.iter().any(|s| s.is_some()));
+
+        // Complete first so we do not leak
+        consumer.complete(completion).unwrap();
+
+        consumer.reset();
+
+        assert!(consumer.inflight.iter().all(|s| s.is_none()));
+        assert_eq!(consumer.inner.num_inflight(), 0);
+    }
+
+    #[test]
+    fn test_virtq_consumer_reset_clears_inflight() {
+        let ring = make_ring(16);
+        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);
+
+        // Submit two entries and poll both
+        let se1 = producer.chain().completion(16).build().unwrap();
+        producer.submit(se1).unwrap();
+        let se2 = producer.chain().completion(16).build().unwrap();
+        producer.submit(se2).unwrap();
+
+        let (_e1, c1) = consumer.poll(1024).unwrap().unwrap();
+        let (_e2, c2) = consumer.poll(1024).unwrap().unwrap();
+        // Complete both before reset
+        consumer.complete(c1).unwrap();
+        consumer.complete(c2).unwrap();
+
+        consumer.reset();
+
+        assert!(consumer.inflight.iter().all(|s| s.is_none()));
+        assert_eq!(consumer.inner.num_inflight(), 0);
+    }
 }

src/hyperlight_common/src/virtq/pool.rs

@@ -516,6 +516,12 @@ impl<const N: usize> Slab<N> {
     pub const fn slot_size() -> usize {
         N
     }
+
+    /// Reset the slab to initial state which is all slots free.
+    pub fn reset(&mut self) {
+        self.used_slots.clear();
+        self.last_free_run = None;
+    }
 }
 
 #[inline]
@@ -547,6 +553,13 @@ impl<const L: usize, const U: usize> BufferPool<L, U> {
             inner: inner.into(),
         })
     }
+
+    /// Reset the pool to initial state
+    pub fn reset(&self) {
+        let mut inner = self.inner.borrow_mut();
+        inner.lower.reset();
+        inner.upper.reset();
+    }
 }
 
 #[cfg(all(test, loom))]
@@ -1171,6 +1184,87 @@ mod tests {
         slab.dealloc(a3).unwrap();
         slab.dealloc(a4).unwrap();
     }
+
+    #[test]
+    fn test_slab_reset_returns_to_initial_state() {
+        let mut slab = make_slab::<256>(4096);
+        let initial_free = slab.free_bytes();
+        let initial_cap = slab.capacity();
+
+        // Allocate some slots
+        let _a1 = slab.alloc(256).unwrap();
+        let _a2 = slab.alloc(512).unwrap();
+        assert!(slab.free_bytes() < initial_free);
+
+        slab.reset();
+
+        assert_eq!(slab.free_bytes(), initial_free);
+        assert_eq!(slab.capacity(), initial_cap);
+        assert!(slab.last_free_run.is_none());
+        assert_eq!(slab.used_slots.count_ones(..), 0);
+
+        // Should be able to allocate the full capacity again
+        let a = slab.alloc(initial_cap).unwrap();
+        assert_eq!(a.len, initial_cap);
+    }
+
+    #[test]
+    fn test_slab_reset_matches_new() {
+        let base = align_up(0x10000, 256) as u64;
+        let region = 4096;
+
+        let fresh = Slab::<256>::new(base, region).unwrap();
+
+        let mut used = Slab::<256>::new(base, region).unwrap();
+        let _a = used.alloc(256).unwrap();
+        let _b = used.alloc(1024).unwrap();
+        used.reset();
+
+        assert_eq!(used.free_bytes(), fresh.free_bytes());
+        assert_eq!(used.capacity(), fresh.capacity());
+        assert_eq!(
+            used.used_slots.count_ones(..),
+            fresh.used_slots.count_ones(..)
+        );
+        assert!(used.last_free_run.is_none());
+        assert!(fresh.last_free_run.is_none());
+    }
+
+    #[test]
+    fn test_buffer_pool_reset_returns_to_initial_state() {
+        let pool = make_pool::<256, 4096>(0x20000);
+
+        // Allocate from both tiers
+        let a1 = pool.inner.borrow_mut().alloc(128).unwrap();
+        let a2 = pool.inner.borrow_mut().alloc(8192).unwrap();
+        assert!(a1.len > 0);
+        assert!(a2.len > 0);
+
+        pool.reset();
+
+        let inner = pool.inner.borrow();
+        assert_eq!(inner.lower.used_slots.count_ones(..), 0);
+        assert_eq!(inner.upper.used_slots.count_ones(..), 0);
+        assert!(inner.lower.last_free_run.is_none());
+        assert!(inner.upper.last_free_run.is_none());
+    }
+
+    #[test]
+    fn test_buffer_pool_reset_allows_reallocation() {
+        let pool = make_pool::<256, 4096>(0x20000);
+
+        // Fill up some allocations
+        let mut allocs = Vec::new();
+        for _ in 0..5 {
+            allocs.push(pool.inner.borrow_mut().alloc(256).unwrap());
+        }
+
+        pool.reset();
+
+        // Should be able to allocate as if fresh
+        let a = pool.inner.borrow_mut().alloc(256).unwrap();
+        assert!(a.len > 0);
+    }
 }
 
 #[cfg(test)]

src/hyperlight_common/src/virtq/producer.rs

@@ -329,6 +329,13 @@ where
         }
         Ok(())
     }
+
+    /// Reset ring and inflight state to initial values.
+    /// Does not reset the buffer pool; call pool.reset() separately if needed.
+    pub fn reset(&mut self) {
+        self.inner.reset();
+        self.inflight.fill(None);
+    }
 }
 
 /// Builder for configuring a descriptor chain's buffer layout.
@@ -787,4 +794,45 @@ mod tests {
         assert_eq!(cqe.token, token);
         assert_eq!(&cqe.data[..], b"response data");
     }
+
+    #[test]
+    fn test_virtq_producer_reset() {
+        let ring = make_ring(16);
+        let (mut producer, mut consumer, _notifier) = make_test_producer(&ring);
+
+        // Submit and complete a round trip
+        let mut se = producer.chain().entry(32).completion(64).build().unwrap();
+        se.write_all(b"hello").unwrap();
+        producer.submit(se).unwrap();
+
+        let (entry, completion) = consumer.poll(1024).unwrap().unwrap();
+        assert_eq!(entry.data().as_ref(), b"hello");
+        consumer.complete(completion).unwrap();
+        let _ = producer.poll().unwrap().unwrap();
+
+        // Now reset
+        producer.reset();
+
+        // All inflight slots should be None
+        assert!(producer.inflight.iter().all(|s| s.is_none()));
+        // Ring state should be back to initial
+        assert_eq!(producer.inner.num_free(), producer.inner.len());
+    }
+
+    #[test]
+    fn test_virtq_producer_reset_clears_inflight() {
+        let ring = make_ring(16);
+        let (mut producer, _consumer, _notifier) = make_test_producer(&ring);
+
+        // Submit without completing
+        let se = producer.chain().completion(64).build().unwrap();
+        producer.submit(se).unwrap();
+
+        assert!(producer.inflight.iter().any(|s| s.is_some()));
+
+        producer.reset();
+
+        assert!(producer.inflight.iter().all(|s| s.is_none()));
+        assert_eq!(producer.inner.num_free(), producer.inner.len());
+    }
 }

src/hyperlight_common/src/virtq/ring.rs

@@ -369,6 +369,12 @@ impl RingCursor {
     pub fn wrap(&self) -> bool {
         self.wrap
     }
+
+    /// Reset cursor to initial state.
+    pub fn reset(&mut self) {
+        self.head = 0;
+        self.wrap = true;
+    }
 }
 
 /// Producer (driver) side of a packed virtqueue.
@@ -817,6 +823,18 @@ impl<M: MemOps> RingProducer<M> {
 
         Ok(should_notify(evt, self.len() as u16, old, new))
     }
+
+    /// Reset to initial state matching a freshly zeroed ring.
+    pub fn reset(&mut self) {
+        let size = self.desc_table.len();
+        self.avail_cursor.reset();
+        self.used_cursor.reset();
+        self.num_free = size;
+        self.id_free.clear();
+        self.id_free.extend(0..size as u16);
+        self.id_num.iter_mut().for_each(|n| *n = 0);
+        self.event_flags_shadow = EventFlags::ENABLE;
+    }
 }
 
 /// Consumer (device) side of a packed virtqueue.
@@ -1164,6 +1182,16 @@ impl<M: MemOps> RingConsumer<M> {
 
         Ok(should_notify(evt, self.desc_table.len() as u16, old, new))
     }
+
+    /// Reset to initial state matching a freshly zeroed ring.
+    /// Does not reallocate internal buffers.
+    pub fn reset(&mut self) {
+        self.avail_cursor.reset();
+        self.used_cursor.reset();
+        self.id_num.iter_mut().for_each(|n| *n = 0);
+        self.num_inflight = 0;
+        self.event_flags_shadow = EventFlags::ENABLE;
+    }
 }
 
 /// Common packed-ring notification decision:
@@ -2974,6 +3002,113 @@ pub(crate) mod tests {
             let (_, _) = consumer.poll_available().unwrap();
         }
     }
+
+    #[test]
+    fn test_ring_cursor_reset() {
+        let mut cursor = RingCursor::new(16);
+        cursor.advance_by(5);
+        assert_eq!(cursor.head(), 5);
+
+        cursor.reset();
+        assert_eq!(cursor, RingCursor::new(16));
+        assert_eq!(cursor.head(), 0);
+        assert!(cursor.wrap());
+    }
+
+    #[test]
+    fn test_ring_cursor_reset_after_wrap() {
+        let mut cursor = RingCursor::new(4);
+        // Advance past the wrap point
+        cursor.advance_by(5);
+        assert_eq!(cursor.head(), 1);
+        assert!(!cursor.wrap());
+
+        cursor.reset();
+        assert_eq!(cursor.head(), 0);
+        assert!(cursor.wrap());
+    }
+
+    #[test]
+    fn test_ring_producer_reset_matches_new() {
+        let ring = make_ring(8);
+        let fresh = make_producer(&ring);
+
+        let mut used = make_producer(&ring);
+        // Mutate state
+        used.submit_one(0x1000, 64, false).unwrap();
+        used.submit_one(0x2000, 128, true).unwrap();
+
+        used.reset();
+
+        assert_eq!(used.avail_cursor, fresh.avail_cursor);
+        assert_eq!(used.used_cursor, fresh.used_cursor);
+        assert_eq!(used.num_free, fresh.num_free);
+        assert_eq!(used.id_free.len(), fresh.id_free.len());
+        assert_eq!(used.id_num.as_slice(), fresh.id_num.as_slice());
+        assert_eq!(used.event_flags_shadow, fresh.event_flags_shadow);
+    }
+
+    #[test]
+    fn test_ring_producer_reset_id_free_complete() {
+        let ring = make_ring(8);
+        let mut producer = make_producer(&ring);
+
+        // Submit and consume several descriptors
+        for i in 0..4u64 {
+            producer.submit_one(0x1000 + i * 0x100, 64, false).unwrap();
+        }
+        assert_eq!(producer.num_free, 4);
+
+        producer.reset();
+
+        assert_eq!(producer.num_free, 8);
+        assert_eq!(producer.id_free.len(), 8);
+        // All IDs 0..8 should be present
+        for id in 0..8u16 {
+            assert!(producer.id_free.contains(&id));
+        }
+    }
+
+    #[test]
+    fn test_ring_consumer_reset_matches_new() {
+        let ring = make_ring(8);
+        let fresh = make_consumer(&ring);
+
+        let mut used = make_consumer(&ring);
+
+        // Submit from producer side so consumer has something to poll
+        let mut producer = make_producer(&ring);
+        producer.submit_one(0x1000, 64, false).unwrap();
+
+        // Consumer polls the available descriptor
+        let (id, _chain) = used.poll_available().unwrap();
+        used.submit_used(id, 64).unwrap();
+
+        used.reset();
+
+        assert_eq!(used.avail_cursor, fresh.avail_cursor);
+        assert_eq!(used.used_cursor, fresh.used_cursor);
+        assert_eq!(used.id_num.as_slice(), fresh.id_num.as_slice());
+        assert_eq!(used.num_inflight, fresh.num_inflight);
+        assert_eq!(used.event_flags_shadow, fresh.event_flags_shadow);
+    }
+
+    #[test]
+    fn test_ring_consumer_reset_clears_inflight() {
+        let ring = make_ring(8);
+        let mut producer = make_producer(&ring);
+        let mut consumer = make_consumer(&ring);
+
+        // Submit and poll two items (consume but do not complete)
+        producer.submit_one(0x1000, 64, false).unwrap();
+        producer.submit_one(0x2000, 64, false).unwrap();
+        let _ = consumer.poll_available().unwrap();
+        let _ = consumer.poll_available().unwrap();
+        assert_eq!(consumer.num_inflight, 2);
+
+        consumer.reset();
+        assert_eq!(consumer.num_inflight, 0);
+    }
 }
 
 #[cfg(test)]