vring: add VringT::get_mut()

Add VringT::get_mut() to get exclusive reference to underlying
VringState object, so the clients could avoid repeatedly lock/unlock
when using VringMutex and VringRwLock.

Signed-off-by: Liu Jiang <[email protected]>

Author: jiangliu

src/lib.rs

@@ -31,7 +31,9 @@ mod handler;
 pub use self::handler::VhostUserHandlerError;
 
 mod vring;
-pub use self::vring::{VringMutex, VringRwLock, VringState, VringT};
+pub use self::vring::{
+    VringMutex, VringRwLock, VringState, VringStateGuard, VringStateMutGuard, VringT,
+};
 
 /// An alias for `GuestMemoryAtomic<GuestMemoryMmap<B>>` to simplify code.
 type GM<B> = GuestMemoryAtomic<GuestMemoryMmap<B>>;

src/vring.rs

@@ -7,7 +7,7 @@
 
 use std::fs::File;
 use std::io;
-use std::ops::Deref;
+use std::ops::{Deref, DerefMut};
 use std::os::unix::io::{FromRawFd, IntoRawFd};
 use std::result::Result;
 use std::sync::{Arc, Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard};
@@ -38,16 +38,37 @@ impl<'a, M: GuestAddressSpace> Deref for VringStateGuard<'a, M> {
     }
 }
 
-/*
-impl<'a, M: GuestAddressSpace> DerefMut for VringStateGuard<'a, M> {
+/// Struct to hold an exclusive reference to the underlying `VringState` object.
+pub enum VringStateMutGuard<'a, M: GuestAddressSpace> {
+    /// A reference to a `VringState` object.
+    StateObject(&'a mut VringState<M>),
+    /// A `MutexGuard` for a `VringState` object.
+    MutexGuard(MutexGuard<'a, VringState<M>>),
+    /// A `WriteGuard` for a `VringState` object.
+    RwLockWriteGuard(RwLockWriteGuard<'a, VringState<M>>),
+}
+
+impl<'a, M: GuestAddressSpace> Deref for VringStateMutGuard<'a, M> {
+    type Target = VringState<M>;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            VringStateMutGuard::StateObject(v) => v,
+            VringStateMutGuard::MutexGuard(v) => v.deref(),
+            VringStateMutGuard::RwLockWriteGuard(v) => v.deref(),
+        }
+    }
+}
+
+impl<'a, M: GuestAddressSpace> DerefMut for VringStateMutGuard<'a, M> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         match self {
-            VringStateGuard::StateObject(v) => v,
-            VringStateGuard::MutexGuard(v) => v.deref_mut(),
+            VringStateMutGuard::StateObject(v) => v,
+            VringStateMutGuard::MutexGuard(v) => v.deref_mut(),
+            VringStateMutGuard::RwLockWriteGuard(v) => v.deref_mut(),
         }
     }
 }
- */
 
 pub trait VringT<M: GuestAddressSpace> {
     /// Create a new instance of Vring.
@@ -56,6 +77,9 @@ pub trait VringT<M: GuestAddressSpace> {
     /// Get an immutable reference to the kick event fd.
     fn get_ref(&self) -> VringStateGuard<M>;
 
+    /// Get a mutable reference to the kick event fd.
+    fn get_mut(&mut self) -> VringStateMutGuard<M>;
+
     /// Add an used descriptor into the used queue.
     fn add_used(&mut self, desc_index: u16, len: u32) -> Result<(), VirtQueError>;
 
@@ -149,6 +173,10 @@ impl<M: GuestAddressSpace> VringT<M> for VringState<M> {
         VringStateGuard::StateObject(self)
     }
 
+    fn get_mut(&mut self) -> VringStateMutGuard<M> {
+        VringStateMutGuard::StateObject(self)
+    }
+
     fn add_used(&mut self, desc_index: u16, len: u32) -> Result<(), VirtQueError> {
         self.queue.add_used(desc_index, len)
     }
@@ -256,6 +284,10 @@ impl<M: GuestAddressSpace> VringT<M> for VringMutex<M> {
         VringStateGuard::MutexGuard(self.state.lock().unwrap())
     }
 
+    fn get_mut(&mut self) -> VringStateMutGuard<M> {
+        VringStateMutGuard::MutexGuard(self.lock())
+    }
+
     fn add_used(&mut self, desc_index: u16, len: u32) -> Result<(), VirtQueError> {
         self.lock().add_used(desc_index, len)
     }
@@ -346,6 +378,10 @@ impl<M: GuestAddressSpace> VringT<M> for VringRwLock<M> {
         VringStateGuard::RwLockReadGuard(self.state.read().unwrap())
     }
 
+    fn get_mut(&mut self) -> VringStateMutGuard<M> {
+        VringStateMutGuard::RwLockWriteGuard(self.write_lock())
+    }
+
     fn add_used(&mut self, desc_index: u16, len: u32) -> Result<(), VirtQueError> {
         self.write_lock().add_used(desc_index, len)
     }