virtio: add ring buffer type for describing guest memory

Add a ring buffer type that is tailored for holding `struct iovec`
objects that point to guest memory for IO. The `struct iovec` objects
represent the memory that the guest passed to us as `Descriptors` in a
VirtIO queue for performing some I/O operation.

We plan to use this type to describe the guest memory we have available
for doing network RX. This should facilitate us in optimizing the
reception of data from the TAP device using `readv`, thus avoiding a
memory copy.

Co-authored-by: Egor Lazarchuk <[email protected]>
Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

src/vmm/src/devices/virtio/iov_deque.rs

@@ -0,0 +1,393 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::os::fd::AsRawFd;
+
+use libc::{c_int, c_void, iovec, off_t, size_t};
+use memfd;
+
+use super::queue::FIRECRACKER_MAX_QUEUE_SIZE;
+use crate::arch::PAGE_SIZE;
+
+#[derive(Debug, thiserror::Error, displaydoc::Display)]
+pub enum IovDequeError {
+    /// Error with memfd: {0}
+    Memfd(#[from] memfd::Error),
+    /// Error while resizing memfd: {0}
+    MemfdResize(std::io::Error),
+    /// Error calling mmap: {0}
+    Mmap(std::io::Error),
+}
+
+/// ['IovDeque'] is a ring buffer tailored for `struct iovec` objects.
+///
+/// From the point of view of API, [`IovDeque`] is a typical ring buffer that allows us to push
+/// `struct iovec` objects at the end of the buffer and pop them from its beginning.
+///
+/// It is tailored to store `struct iovec` objects that described memory that was passed to us from
+/// the guest via a VirtIO queue. This allows us to assume the maximum size of a ring buffer (the
+/// negotiated size of the queue).
+// An important feature of the data structure is that it can give us a slice of all `struct iovec`
+// objects in the queue, so that we can use this `&mut [iovec]` to perform operations such as
+// `readv`. A typical implementation of a ring buffer allows for entries to wrap around the end of
+// the underlying buffer. For example, a ring buffer with a capacity of 10 elements which
+// currently holds 4 elements can look like this:
+//
+//                      tail                        head
+//                       |                           |
+//                       v                           v
+//                 +---+---+---+---+---+---+---+---+---+---+
+// ring buffer:    | C | D |   |   |   |   |   |   | A | B |
+//                 +---+---+---+---+---+---+---+---+---+---+
+//
+// When getting a slice for this data we should get something like that: &[A, B, C, D], which
+// would require copies in order to make the elements continuous in memory.
+//
+// In order to avoid that and make the operation of getting a slice more efficient, we implement
+// the optimization described in the "Optimization" section of the "Circular buffer" wikipedia
+// entry: https://en.wikipedia.org/wiki/Circular_buffer. The optimization consists of allocating
+// double the size of the virtual memory required for the buffer and map both parts on the same
+// physical address. Looking at the same example as before, we should get, this picture:
+//
+//                                    head   |    tail
+//                                     |     |     |
+//                                     v     |     v
+//   +---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+
+//   | C | D |   |   |   |   |   |   | A | B | C | D |   |   |   |   |   |   | A | B |
+//   +---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+
+//            First virtual page             |       Second virtual page
+//                                           |
+//                                           |
+//
+//                                     Virtual memory
+// ---------------------------------------------------------------------------------------
+//                                    Physical memory
+//
+//                      +---+---+---+---+---+---+---+---+---+---+
+//                      | C | D |   |   |   |   |   |   | A | B |
+//                      +---+---+---+---+---+---+---+---+---+---+
+//
+// Like that, the elements stored in the buffer are always laid out in contiguous virtual memory,
+// so making a slice out of them does not require any copies.
+#[derive(Debug)]
+pub struct IovDeque {
+    iov: *mut libc::iovec,
+    start: u16,
+    len: u16,
+}
+
+// SAFETY: This is `Send`. We hold sole ownership of the underlying buffer.
+unsafe impl Send for IovDeque {}
+
+impl IovDeque {
+    /// Create a [`memfd`] object that represents a single physical page
+    fn create_memfd() -> Result<memfd::Memfd, IovDequeError> {
+        // Create a sealable memfd.
+        let opts = memfd::MemfdOptions::default().allow_sealing(true);
+        let mfd = opts.create("sized-1K")?;
+
+        // Resize to system page size.
+        mfd.as_file()
+            .set_len(PAGE_SIZE.try_into().unwrap())
+            .map_err(IovDequeError::MemfdResize)?;
+
+        // Add seals to prevent further resizing.
+        mfd.add_seals(&[memfd::FileSeal::SealShrink, memfd::FileSeal::SealGrow])?;
+
+        // Prevent further sealing changes.
+        mfd.add_seal(memfd::FileSeal::SealSeal)?;
+
+        Ok(mfd)
+    }
+
+    /// A safe wrapper on top of libc's `mmap` system call
+    fn mmap(
+        addr: *mut c_void,
+        len: size_t,
+        prot: c_int,
+        flags: c_int,
+        fd: c_int,
+        offset: off_t,
+    ) -> Result<*mut c_void, IovDequeError> {
+        // SAFETY: We are calling the system call with valid arguments and properly checking its
+        // return value
+        let ptr = unsafe { libc::mmap(addr, len, prot, flags, fd, offset) };
+        if ptr == libc::MAP_FAILED {
+            return Err(IovDequeError::Mmap(std::io::Error::last_os_error()));
+        }
+
+        Ok(ptr)
+    }
+
+    /// Allocate memory for our ring buffer
+    ///
+    /// This will allocate exactly two pages of virtual memory. In order to implement the
+    /// optimization that allows us to always have elements in contiguous memory we need
+    /// allocations at the granularity of `PAGE_SIZE`. Now, our queues are at maximum 256
+    /// descriptors long and `struct iovec` looks like this:
+    ///
+    /// ```Rust
+    /// pub struct iovec {
+    ///    pub iov_base: *mut ::c_void,
+    ///    pub iov_len: ::size_t,
+    /// }
+    /// ```
+    ///
+    /// so, it's 16 bytes long. As a result, we need a single page for holding the actual data of
+    /// our buffer.
+    fn allocate_ring_buffer_memory() -> Result<*mut c_void, IovDequeError> {
+        // The fact that we allocate two pages is due to the size of `struct iovec` times our queue
+        // size equals the page size. Add here a debug assertion to reflect that and ensure that we
+        // will adapt our logic if the assumption changes in the future.
+        debug_assert_eq!(
+            std::mem::size_of::<iovec>() * usize::from(FIRECRACKER_MAX_QUEUE_SIZE),
+            PAGE_SIZE
+        );
+
+        Self::mmap(
+            std::ptr::null_mut(),
+            PAGE_SIZE * 2,
+            libc::PROT_NONE,
+            libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
+            -1,
+            0,
+        )
+    }
+
+    /// Create a new [`IovDeque`] that can hold memory described by a single VirtIO queue.
+    pub fn new() -> Result<Self, IovDequeError> {
+        let memfd = Self::create_memfd()?;
+        let raw_memfd = memfd.as_file().as_raw_fd();
+        let buffer = Self::allocate_ring_buffer_memory()?;
+
+        // Map the first page of virtual memory to the physical page described by the memfd object
+        let _ = Self::mmap(
+            buffer,
+            PAGE_SIZE,
+            libc::PROT_READ | libc::PROT_WRITE,
+            libc::MAP_SHARED | libc::MAP_FIXED,
+            raw_memfd,
+            0,
+        )?;
+
+        // Map the second page of virtual memory to the physical page described by the memfd object
+        // SAFETY: safe because `Self::allocate_ring_buffer_memory` allocates exactly two pages for
+        // us
+        let next_page = unsafe { buffer.add(PAGE_SIZE) };
+        let _ = Self::mmap(
+            next_page,
+            PAGE_SIZE,
+            libc::PROT_READ | libc::PROT_WRITE,
+            libc::MAP_SHARED | libc::MAP_FIXED,
+            raw_memfd,
+            0,
+        )?;
+
+        Ok(Self {
+            iov: buffer.cast(),
+            start: 0,
+            len: 0,
+        })
+    }
+
+    /// Returns the number of `iovec` objects currently in the [`IovDeque`]
+    #[inline(always)]
+    pub fn len(&self) -> u16 {
+        self.len
+    }
+
+    /// Returns `true` if the [`IovDeque`] is full, `false` otherwise
+    #[inline(always)]
+    fn is_full(&self) -> bool {
+        self.len() == FIRECRACKER_MAX_QUEUE_SIZE
+    }
+
+    /// Resets the queue, dropping all its elements.
+    #[inline(always)]
+    pub fn clear(&mut self) {
+        self.start = 0;
+        self.len = 0;
+    }
+
+    /// Adds an `iovec` in the ring buffer.
+    ///
+    /// Returns an `IovDequeError::Full` error if the buffer is full.
+    pub fn push_back(&mut self, iov: iovec) {
+        // This should NEVER happen, since our ring buffer is as big as the maximum queue size.
+        // We also check for the sanity of the VirtIO queues, in queue.rs, which means that if we
+        // ever try to add something in a full ring buffer, there is an internal bug in the device
+        // emulation logic. Panic here because the device is hopelessly broken.
+        assert!(
+            !self.is_full(),
+            "The number of `iovec` objects is bigger than the available space"
+        );
+
+        // SAFETY: self.iov is a valid pointer and `self.start + self.len` is within range (we
+        // asserted before that the buffer is not full).
+        unsafe {
+            self.iov
+                .add((self.start + self.len) as usize)
+                .write_volatile(iov)
+        };
+        self.len += 1;
+    }
+
+    /// Returns the `iovec` at position `pos` if `pos` < self.len(), otherwise it will return
+    /// `None`.
+    pub fn peek(&self, pos: u16) -> Option<iovec> {
+        if pos < self.len() {
+            // SAFETY: Safe, because `self.iov` is a valid pointer and we checked that `pos` points
+            // to an existing item.
+            Some(unsafe { self.iov.add((self.start + pos) as usize).read_volatile() })
+        } else {
+            None
+        }
+    }
+
+    /// Pops the first `nr_iovecs` iovecs from the buffer.
+    ///
+    /// Returns the total number of bytes of all the popped iovecs. This will panic if we are asked
+    /// to pop more iovecs than what is currently available in the buffer.
+    pub fn pop_front(&mut self, nr_iovecs: u16) -> usize {
+        if self.len() < nr_iovecs {
+            panic!("Internal bug! Trying to drop more iovec objects than what is available");
+        }
+
+        let mut bytes = 0;
+
+        for i in 0..nr_iovecs {
+            let iov = self.peek(i).unwrap();
+            bytes += iov.iov_len;
+        }
+
+        self.start += nr_iovecs;
+        self.len -= nr_iovecs;
+        if self.start >= FIRECRACKER_MAX_QUEUE_SIZE {
+            self.start -= FIRECRACKER_MAX_QUEUE_SIZE;
+        }
+
+        bytes
+    }
+
+    /// Gets a mutable pointer of the underlying `iovec` objects currently in the buffer.
+    pub fn as_mut_ptr(&self) -> *mut iovec {
+        // SAFETY: This is safe because `self.iov` is a valid `*mut iovec` and `self.start` is
+        // always within range
+        unsafe { self.iov.add(self.start as usize) }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use libc::iovec;
+
+    use super::IovDeque;
+
+    #[test]
+    fn test_new() {
+        let iov = IovDeque::new().unwrap();
+        assert_eq!(iov.len(), 0);
+    }
+
+    fn make_iovec(id: u16, len: u16) -> iovec {
+        iovec {
+            iov_base: id as *mut libc::c_void,
+            iov_len: len as usize,
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_push_back_too_many() {
+        let mut iov = IovDeque::new().unwrap();
+        assert_eq!(iov.len(), 0);
+
+        for i in 0u16..256 {
+            iov.push_back(make_iovec(i, i));
+            assert_eq!(iov.len(), i + 1);
+        }
+
+        iov.push_back(make_iovec(0, 0));
+    }
+
+    #[test]
+    fn test_pop() {
+        let mut deque = IovDeque::new().unwrap();
+        assert_eq!(deque.len(), 0);
+        assert!(!deque.is_full());
+        assert_eq!(deque.pop_front(0), 0);
+
+        for i in 0u16..256 {
+            deque.push_back(make_iovec(i, i));
+            assert_eq!(deque.len(), i + 1);
+        }
+
+        assert!(deque.is_full());
+        assert!(deque.len() != 0);
+
+        for i in 0u16..256 {
+            let bytes = deque.pop_front(1);
+            assert_eq!(bytes, i as usize);
+        }
+    }
+
+    #[test]
+    fn test_pop_many() {
+        let mut deque = IovDeque::new().unwrap();
+
+        for i in 0u16..256 {
+            deque.push_back(make_iovec(i, i));
+        }
+
+        assert_eq!(deque.pop_front(1), 0);
+        assert_eq!(deque.len(), 255);
+        assert_eq!(deque.pop_front(2), 3);
+        assert_eq!(deque.len(), 253);
+        assert_eq!(deque.pop_front(4), (3..7).sum::<usize>());
+        assert_eq!(deque.len(), 249);
+        assert_eq!(deque.pop_front(8), (7..15).sum::<usize>());
+        assert_eq!(deque.len(), 241);
+        assert_eq!(deque.pop_front(16), (15..31).sum::<usize>());
+        assert_eq!(deque.len(), 225);
+        assert_eq!(deque.pop_front(32), (31..63).sum::<usize>());
+        assert_eq!(deque.len(), 193);
+        assert_eq!(deque.pop_front(64), (63..127).sum::<usize>());
+        assert_eq!(deque.len(), 129);
+        assert_eq!(deque.pop_front(128), (127..255).sum::<usize>());
+        assert_eq!(deque.len(), 1);
+    }
+
+    #[test]
+    fn test_peek() {
+        let mut deque = IovDeque::new().unwrap();
+
+        for i in 0u16..256 {
+            deque.push_back(make_iovec(i, i));
+        }
+
+        for i in 0u16..256 {
+            assert_eq!(make_iovec(i, i), deque.peek(i).unwrap());
+        }
+
+        assert!(deque.peek(256).is_none());
+    }
+
+    #[test]
+    fn test_as_mut_ptr() {
+        let mut deque = IovDeque::new().unwrap();
+        let mut copy: Vec<libc::iovec> = vec![];
+        let ptr = deque.as_mut_ptr();
+        assert!(!ptr.is_null());
+        assert!(ptr.is_aligned());
+
+        for i in 0..256 {
+            deque.push_back(make_iovec(i, 100));
+            copy.push(deque.peek(i).unwrap());
+        }
+
+        assert_eq!(copy.len(), deque.len() as usize);
+        for copied_iov in copy {
+            assert_eq!(copied_iov.iov_len, deque.pop_front(1));
+        }
+    }
+}