Improve the usability of queue interfaces

crates/virtio-queue/src/iterator.rs

@@ -102,7 +102,7 @@ pub struct AvailIter<'b, M> {
 impl<'b, M> AvailIter<'b, M>
 where
     M: Deref,
-    M::Target: GuestMemory + Sized,
+    M::Target: GuestMemory,
 {
     /// Create a new instance of `AvailInter`.
     ///

crates/virtio-queue/src/lib.rs

@@ -155,7 +155,9 @@ pub trait QueueStateT: for<'a> QueueStateGuard<'a> {
     fn set_event_idx(&mut self, enabled: bool);
 
     /// Read the `idx` field from the available ring.
-    fn avail_idx<M: GuestMemory>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error>;
+    fn avail_idx<M>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error>
+    where
+        M: GuestMemory + ?Sized;
 
     /// Read the `idx` field from the used ring.
     fn used_idx<M: GuestMemory>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error>;
@@ -193,4 +195,31 @@ pub trait QueueStateT: for<'a> QueueStateGuard<'a> {
 
     /// Set the index for the next descriptor in the used ring.
     fn set_next_used(&mut self, next_used: u16);
+
+    /// Pop and return the next available descriptor chain, or `None` when there are no more
+    /// descriptor chains available.
+    ///
+    /// This enables the consumption of available descriptor chains in a "one at a time"
+    /// manner, without having to hold a borrow after the method returns.
+    fn pop_descriptor_chain<M>(&mut self, mem: M) -> Option<DescriptorChain<M>>
+    where
+        M: Clone + Deref,
+        M::Target: GuestMemory;
+}
+
+/// Trait to access and manipulate a Virtio queue that's known to be exclusively accessed
+/// by a single execution thread.
+pub trait QueueStateOwnedT: QueueStateT {
+    /// Get a consuming iterator over all available descriptor chain heads offered by the driver.
+    ///
+    /// # Arguments
+    /// * `mem` - the `GuestMemory` object that can be used to access the queue buffers.
+    fn iter<M>(&mut self, mem: M) -> Result<AvailIter<'_, M>, Error>
+    where
+        M: Deref,
+        M::Target: GuestMemory;
+
+    /// Undo the last advancement of the next available index field by decrementing its
+    /// value by one.
+    fn go_to_previous_position(&mut self);
 }

crates/virtio-queue/src/queue.rs

@@ -16,7 +16,9 @@ use std::sync::atomic::Ordering;
 
 use vm_memory::GuestAddressSpace;
 
-use crate::{AvailIter, Error, QueueGuard, QueueState, QueueStateGuard, QueueStateT};
+use crate::{
+    AvailIter, Error, QueueGuard, QueueState, QueueStateGuard, QueueStateOwnedT, QueueStateT,
+};
 
 /// A convenient wrapper struct for a virtio queue, with associated `GuestMemory` object.
 ///

crates/virtio-queue/src/queue_guard.rs

@@ -8,7 +8,7 @@ use std::sync::atomic::Ordering;
 
 use vm_memory::GuestMemory;
 
-use crate::{AvailIter, Error, QueueState, QueueStateT};
+use crate::{AvailIter, DescriptorChain, Error, QueueState, QueueStateOwnedT, QueueStateT};
 
 /// A guard object to exclusively access an `Queue` object.
 ///
@@ -192,10 +192,21 @@ where
         self.state.set_next_used(next_used);
     }
 
+    /// Pop and return the next available descriptor chain, or `None` when there are no more
+    /// descriptor chains available.
+    pub fn pop_descriptor_chain(&mut self) -> Option<DescriptorChain<M>> {
+        self.state.pop_descriptor_chain(self.mem.clone())
+    }
+
     /// Get a consuming iterator over all available descriptor chain heads offered by the driver.
     pub fn iter(&mut self) -> Result<AvailIter<'_, M>, Error> {
         self.state.deref_mut().iter(self.mem.clone())
     }
+
+    /// Decrement the value of the next available index by one position.
+    pub fn go_to_previous_position(&mut self) {
+        self.state.go_to_previous_position();
+    }
 }
 
 #[cfg(test)]
@@ -259,7 +270,27 @@ mod tests {
                 break;
             }
         }
+
+        {
+            g.go_to_previous_position();
+            let mut chain = g.pop_descriptor_chain().unwrap();
+            // The descriptor index of the head descriptor from the last available chain
+            // defined above is equal to 5. The chain has two descriptors, the second with
+            // the index equal to 6.
+            assert_eq!(chain.head_index(), 5);
+
+            let desc = chain.next().unwrap();
+            assert!(desc.has_next());
+            assert_eq!(desc.next(), 6);
+
+            let desc = chain.next().unwrap();
+            assert!(!desc.has_next());
+
+            assert!(chain.next().is_none());
+        }
+
         // The next chain that can be consumed should have index 3.
+
         assert_eq!(g.next_avail(), 3);
         assert_eq!(g.avail_idx(Ordering::Acquire).unwrap(), Wrapping(3));
         assert_eq!(g.next_used(), 3);

crates/virtio-queue/src/state.rs

@@ -19,7 +19,10 @@ use crate::defs::{
     VIRTQ_USED_ELEMENT_SIZE, VIRTQ_USED_F_NO_NOTIFY, VIRTQ_USED_RING_HEADER_SIZE,
     VIRTQ_USED_RING_META_SIZE,
 };
-use crate::{error, AvailIter, Descriptor, Error, QueueStateGuard, QueueStateT, VirtqUsedElem};
+use crate::{
+    error, AvailIter, Descriptor, DescriptorChain, Error, QueueStateGuard, QueueStateOwnedT,
+    QueueStateT, VirtqUsedElem,
+};
 
 /// Struct to maintain information and manipulate state of a virtio queue.
 ///
@@ -64,19 +67,6 @@ pub struct QueueState {
 }
 
 impl QueueState {
-    /// Get a consuming iterator over all available descriptor chain heads offered by the driver.
-    ///
-    /// # Arguments
-    /// * `mem` - the `GuestMemory` object that can be used to access the queue buffers.
-    pub fn iter<M>(&mut self, mem: M) -> Result<AvailIter<'_, M>, Error>
-    where
-        M: Deref,
-        M::Target: GuestMemory + Sized,
-    {
-        self.avail_idx(mem.deref(), Ordering::Acquire)
-            .map(move |idx| AvailIter::new(mem, idx, self))
-    }
-
     // Helper method that writes `val` to the `avail_event` field of the used ring, using
     // the provided ordering.
     fn set_avail_event<M: GuestMemory>(
@@ -311,7 +301,10 @@ impl QueueStateT for QueueState {
         self.event_idx_enabled = enabled;
     }
 
-    fn avail_idx<M: GuestMemory>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error> {
+    fn avail_idx<M>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error>
+    where
+        M: GuestMemory + ?Sized,
+    {
         let addr = self
             .avail_ring
             .checked_add(2)
@@ -450,4 +443,34 @@ impl QueueStateT for QueueState {
     fn set_next_used(&mut self, next_used: u16) {
         self.next_used = Wrapping(next_used);
     }
+
+    fn pop_descriptor_chain<M>(&mut self, mem: M) -> Option<DescriptorChain<M>>
+    where
+        M: Clone + Deref,
+        M::Target: GuestMemory,
+    {
+        // Default, iter-based impl. Will be subsequently improved.
+        match self.iter(mem) {
+            Ok(mut iter) => iter.next(),
+            Err(e) => {
+                error!("Iterator error {}", e);
+                None
+            }
+        }
+    }
+}
+
+impl QueueStateOwnedT for QueueState {
+    fn iter<M>(&mut self, mem: M) -> Result<AvailIter<'_, M>, Error>
+    where
+        M: Deref,
+        M::Target: GuestMemory,
+    {
+        self.avail_idx(mem.deref(), Ordering::Acquire)
+            .map(move |idx| AvailIter::new(mem, idx, self))
+    }
+
+    fn go_to_previous_position(&mut self) {
+        self.next_avail -= Wrapping(1);
+    }
 }

crates/virtio-queue/src/state_sync.rs

@@ -3,12 +3,13 @@
 // SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
 
 use std::num::Wrapping;
+use std::ops::Deref;
 use std::sync::atomic::Ordering;
 use std::sync::{Arc, Mutex, MutexGuard};
 
 use vm_memory::GuestMemory;
 
-use crate::{Error, QueueState, QueueStateGuard, QueueStateT};
+use crate::{DescriptorChain, Error, QueueState, QueueStateGuard, QueueStateT};
 
 /// Struct to maintain information and manipulate state of a virtio queue for multi-threaded
 /// context.
@@ -107,7 +108,10 @@ impl QueueStateT for QueueStateSync {
         self.lock_state().set_event_idx(enabled);
     }
 
-    fn avail_idx<M: GuestMemory>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error> {
+    fn avail_idx<M>(&self, mem: &M, order: Ordering) -> Result<Wrapping<u16>, Error>
+    where
+        M: GuestMemory + ?Sized,
+    {
         self.lock_state().avail_idx(mem, order)
     }
 
@@ -151,6 +155,14 @@ impl QueueStateT for QueueStateSync {
     fn set_next_used(&mut self, next_used: u16) {
         self.lock_state().set_next_used(next_used);
     }
+
+    fn pop_descriptor_chain<M>(&mut self, mem: M) -> Option<DescriptorChain<M>>
+    where
+        M: Clone + Deref,
+        M::Target: GuestMemory,
+    {
+        self.lock_state().pop_descriptor_chain(mem)
+    }
 }
 
 #[cfg(test)]