refactor: use slab for user_data instead of box

Use slab crate instead of box to avoid `unsafe`.
Move generic type from push/pop to IoUring struct.

Signed-off-by: ihciah <[email protected]>
Signed-off-by: Sudan Landge <[email protected]>

Author: Sudan Landge

Cargo.lock

@@ -37,6 +37,7 @@ bincode = "1.2.1"
 micro_http = { git = "https://github.com/firecracker-microvm/micro-http" }
 log-instrument = { path = "../log-instrument", optional = true }
 crc64 = "2.0.0"
+slab = "0.4.7"
 
 seccompiler = { path = "../seccompiler" }
 utils = { path = "../utils" }

src/vmm/Cargo.toml

@@ -37,7 +37,7 @@ pub enum AsyncIoError {
 #[derive(Debug)]
 pub struct AsyncFileEngine<T> {
     file: File,
-    ring: IoUring,
+    ring: IoUring<WrappedUserData<T>>,
     completion_evt: EventFd,
     phantom: PhantomData<T>,
 }
@@ -73,7 +73,10 @@ impl<T: Debug> WrappedUserData<T> {
 }
 
 impl<T: Debug> AsyncFileEngine<T> {
-    fn new_ring(file: &File, completion_fd: RawFd) -> Result<IoUring, io_uring::IoUringError> {
+    fn new_ring(
+        file: &File,
+        completion_fd: RawFd,
+    ) -> Result<IoUring<WrappedUserData<T>>, io_uring::IoUringError> {
         IoUring::new(
             u32::from(IO_URING_NUM_ENTRIES),
             vec![file],
@@ -142,21 +145,18 @@ impl<T: Debug> AsyncFileEngine<T> {
 
         let wrapped_user_data = WrappedUserData::new_with_dirty_tracking(addr, user_data);
 
-        // SAFETY: Safe because we trust that the host kernel will pass us back a completed entry
-        // with this same `user_data`, so that the value will not be leaked.
-        unsafe {
-            self.ring.push(Operation::read(
+        self.ring
+            .push(Operation::read(
                 0,
                 buf as usize,
                 count,
                 offset,
                 wrapped_user_data,
             ))
-        }
-        .map_err(|err_tuple| UserDataError {
-            user_data: err_tuple.1.user_data,
-            error: AsyncIoError::IoUring(err_tuple.0),
-        })
+            .map_err(|err_tuple| UserDataError {
+                user_data: err_tuple.1.user_data,
+                error: AsyncIoError::IoUring(err_tuple.0),
+            })
     }
 
     pub fn push_write(
@@ -179,34 +179,29 @@ impl<T: Debug> AsyncFileEngine<T> {
 
         let wrapped_user_data = WrappedUserData::new(user_data);
 
-        // SAFETY: Safe because we trust that the host kernel will pass us back a completed entry
-        // with this same `user_data`, so that the value will not be leaked.
-        unsafe {
-            self.ring.push(Operation::write(
+        self.ring
+            .push(Operation::write(
                 0,
                 buf as usize,
                 count,
                 offset,
                 wrapped_user_data,
             ))
-        }
-        .map_err(|err_tuple| UserDataError {
-            user_data: err_tuple.1.user_data,
-            error: AsyncIoError::IoUring(err_tuple.0),
-        })
+            .map_err(|err_tuple| UserDataError {
+                user_data: err_tuple.1.user_data,
+                error: AsyncIoError::IoUring(err_tuple.0),
+            })
     }
 
     pub fn push_flush(&mut self, user_data: T) -> Result<(), UserDataError<T, AsyncIoError>> {
         let wrapped_user_data = WrappedUserData::new(user_data);
 
-        // SAFETY: Safe because we trust that the host kernel will pass us back a completed entry
-        // with this same `user_data`, so that the value will not be leaked.
-        unsafe { self.ring.push(Operation::fsync(0, wrapped_user_data)) }.map_err(|err_tuple| {
-            UserDataError {
+        self.ring
+            .push(Operation::fsync(0, wrapped_user_data))
+            .map_err(|err_tuple| UserDataError {
                 user_data: err_tuple.1.user_data,
                 error: AsyncIoError::IoUring(err_tuple.0),
-            }
-        })
+            })
     }
 
     pub fn kick_submission_queue(&mut self) -> Result<(), AsyncIoError> {
@@ -242,11 +237,7 @@ impl<T: Debug> AsyncFileEngine<T> {
     }
 
     fn do_pop(&mut self) -> Result<Option<Cqe<WrappedUserData<T>>>, AsyncIoError> {
-        // SAFETY: We trust that the host kernel did not touch the operation's `user_data` field.
-        // The `T` type is the same one used for `push`ing and since the kernel made this entry
-        // available in the completion queue, we now have full ownership of it.
-
-        unsafe { self.ring.pop::<WrappedUserData<T>>() }.map_err(AsyncIoError::IoUring)
+        self.ring.pop().map_err(AsyncIoError::IoUring)
     }
 
     pub fn pop(&mut self, mem: &GuestMemoryMmap) -> Result<Option<Cqe<T>>, AsyncIoError> {

src/vmm/src/devices/virtio/virtio_block/io/async_io.rs

@@ -78,7 +78,7 @@ impl IoUringError {
 
 /// Main object representing an io_uring instance.
 #[derive(Debug)]
-pub struct IoUring {
+pub struct IoUring<T> {
     registered_fds_count: u32,
     squeue: SubmissionQueue,
     cqueue: CompletionQueue,
@@ -91,9 +91,10 @@ pub struct IoUring {
     // The total number of ops. These includes the ops on the submission queue, the in-flight ops
     // and the ops that are in the CQ, but haven't been popped yet.
     num_ops: u32,
+    slab: slab::Slab<T>,
 }
 
-impl IoUring {
+impl<T: Debug> IoUring<T> {
     /// Create a new instance.
     ///
     /// # Arguments
@@ -136,13 +137,16 @@ impl IoUring {
 
         let squeue = SubmissionQueue::new(fd, &params).map_err(IoUringError::SQueue)?;
         let cqueue = CompletionQueue::new(fd, &params).map_err(IoUringError::CQueue)?;
+        let slab =
+            slab::Slab::with_capacity(params.sq_entries as usize + params.cq_entries as usize);
 
         let mut instance = Self {
             squeue,
             cqueue,
             fd: file,
             registered_fds_count: 0,
             num_ops: 0,
+            slab,
         };
 
         instance.check_operations()?;
@@ -161,11 +165,7 @@ impl IoUring {
     }
 
     /// Push an [`Operation`](operation/struct.Operation.html) onto the submission queue.
-    ///
-    /// # Safety
-    /// Unsafe because we pass a raw user_data pointer to the kernel.
-    /// It's up to the caller to make sure that this value is ever freed (not leaked).
-    pub unsafe fn push<T: Debug>(&mut self, op: Operation<T>) -> Result<(), (IoUringError, T)> {
+    pub fn push(&mut self, op: Operation<T>) -> Result<(), (IoUringError, T)> {
         // validate that we actually did register fds
         let fd = op.fd();
         match self.registered_fds_count {
@@ -176,29 +176,38 @@ impl IoUring {
                     return Err((IoUringError::FullCQueue, op.user_data()));
                 }
                 self.squeue
-                    .push(op.into_sqe())
+                    .push(op.into_sqe(&mut self.slab))
                     .map(|res| {
                         // This is safe since self.num_ops < IORING_MAX_CQ_ENTRIES (65536)
                         self.num_ops += 1;
                         res
                     })
-                    .map_err(|err_tuple: (SQueueError, T)| -> (IoUringError, T) {
-                        (IoUringError::SQueue(err_tuple.0), err_tuple.1)
+                    .map_err(|(sqe_err, user_data_key)| -> (IoUringError, T) {
+                        (
+                            IoUringError::SQueue(sqe_err),
+                            // We don't use slab.try_remove here for 2 reasons:
+                            // 1. user_data was insertde in slab with step `op.into_sqe` just
+                            //    before the push op so the user_data key should be valid and if
+                            //    key is valid then `slab.remove()` will not fail.
+                            // 2. If we use `slab.try_remove()` we'll have to find a way to return
+                            //    a default value for the generic type T which is difficult because
+                            //    it expands to more crates which don't make it easy to define a
+                            //    default/clone type for type T.
+                            // So beleiving that `slab.remove` won't fail we don't use
+                            // the `slab.try_remove` method.
+                            #[allow(clippy::cast_possible_truncation)]
+                            self.slab.remove(user_data_key as usize),
+                        )
                     })
             }
         }
     }
 
     /// Pop a completed entry off the completion queue. Returns `Ok(None)` if there are no entries.
     /// The type `T` must be the same as the `user_data` type used for `push`-ing the operation.
-    ///
-    /// # Safety
-    /// Unsafe because we reconstruct the `user_data` from a raw pointer passed by the kernel.
-    /// It's up to the caller to make sure that `T` is the correct type of the `user_data`, that
-    /// the raw pointer is valid and that we have full ownership of that address.
-    pub unsafe fn pop<T: Debug>(&mut self) -> Result<Option<Cqe<T>>, IoUringError> {
+    pub fn pop(&mut self) -> Result<Option<Cqe<T>>, IoUringError> {
         self.cqueue
-            .pop()
+            .pop(&mut self.slab)
             .map(|maybe_cqe| {
                 maybe_cqe.map(|cqe| {
                     // This is safe since the pop-ed CQEs have been previously pushed. However
@@ -391,8 +400,8 @@ mod tests {
     use super::*;
     use crate::vstate::memory::{Bytes, MmapRegion};
 
-    fn drain_cqueue(ring: &mut IoUring) {
-        while let Some(entry) = unsafe { ring.pop::<u32>().unwrap() } {
+    fn drain_cqueue(ring: &mut IoUring<u32>) {
+        while let Some(entry) = ring.pop().unwrap() {
             entry.result().unwrap();
 
             // Assert that there were no partial writes.
@@ -581,9 +590,7 @@ mod tests {
                             ring.submit_and_wait_all().unwrap();
                             drain_cqueue(&mut ring);
                         }
-                        unsafe {
-                            ring.push(operation).unwrap();
-                        }
+                        ring.push(operation).unwrap();
                     }
 
                     // Submit any left async ops and wait.

src/vmm/src/io_uring/mod.rs

@@ -13,22 +13,13 @@ unsafe impl ByteValued for io_uring_cqe {}
 #[derive(Debug)]
 pub struct Cqe<T> {
     res: i32,
-    user_data: Box<T>,
+    user_data: T,
 }
 
 impl<T: Debug> Cqe<T> {
-    /// Construct a [`Cqe`] object from a raw `io_uring_cqe`.
-    ///
-    /// # Safety
-    /// Unsafe because we assume full ownership of the inner.user_data address.
-    /// We assume that it points to a valid address created with a `Box<T>`,
-    /// with the correct type `T`, and that ownership of that address is passed
-    /// to this function.
-    pub(crate) unsafe fn new(inner: io_uring_cqe) -> Self {
-        Self {
-            res: inner.res,
-            user_data: Box::from_raw(inner.user_data as *mut T),
-        }
+    /// Construct a Cqe object from a raw `io_uring_cqe`.
+    pub(crate) fn new(res: i32, user_data: T) -> Self {
+        Self { res, user_data }
     }
 
     /// Return the number of bytes successfully transferred by this operation.
@@ -51,13 +42,13 @@ impl<T: Debug> Cqe<T> {
     pub fn map_user_data<U: Debug, F: FnOnce(T) -> U>(self, op: F) -> Cqe<U> {
         Cqe {
             res: self.res,
-            user_data: Box::new(op(self.user_data())),
+            user_data: op(self.user_data()),
         }
     }
 
     /// Consume the object and return the user_data.
     pub fn user_data(self) -> T {
-        *self.user_data
+        self.user_data
     }
 }
 
@@ -69,15 +60,8 @@ mod tests {
     fn test_result() {
         // Check that `result()` returns an `Error` when `res` is negative.
         {
-            let user_data = Box::new(10u8);
-
-            let cqe: Cqe<u8> = unsafe {
-                Cqe::new(io_uring_cqe {
-                    user_data: Box::into_raw(user_data) as u64,
-                    res: -22,
-                    flags: 0,
-                })
-            };
+            let user_data = 10_u8;
+            let cqe: Cqe<u8> = Cqe::new(-22, user_data);
 
             assert_eq!(
                 cqe.result().unwrap_err().kind(),
@@ -87,32 +71,27 @@ mod tests {
 
         // Check that `result()` returns Ok() when `res` is positive.
         {
-            let user_data = Box::new(10u8);
-
-            let cqe: Cqe<u8> = unsafe {
-                Cqe::new(io_uring_cqe {
-                    user_data: Box::into_raw(user_data) as u64,
-                    res: 128,
-                    flags: 0,
-                })
-            };
+            let user_data = 10_u8;
+            let cqe: Cqe<u8> = Cqe::new(128, user_data);
 
             assert_eq!(cqe.result().unwrap(), 128);
         }
     }
 
     #[test]
     fn test_user_data() {
-        let user_data = Box::new(10u8);
-
-        let cqe: Cqe<u8> = unsafe {
-            Cqe::new(io_uring_cqe {
-                user_data: Box::into_raw(user_data) as u64,
-                res: 0,
-                flags: 0,
-            })
-        };
+        let user_data = 10_u8;
+        let cqe: Cqe<u8> = Cqe::new(0, user_data);
 
         assert_eq!(cqe.user_data(), 10);
     }
+
+    #[test]
+    fn test_map_user_data() {
+        let user_data = 10_u8;
+        let cqe: Cqe<u8> = Cqe::new(0, user_data);
+        let cqe = cqe.map_user_data(|x| x + 1);
+
+        assert_eq!(cqe.user_data(), 11);
+    }
 }