virtio-devices: balloon: Factorization and cleanup

Improving the existing code for better readability and in anticipation
for adding an additional virtqueue for the free page reporting feature.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>

virtio-devices/src/balloon.rs

@@ -12,41 +12,43 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use super::{
-    ActivateError, ActivateResult, EpollHelper, EpollHelperError, EpollHelperHandler, VirtioCommon,
-    VirtioDevice, VirtioDeviceType, EPOLL_HELPER_EVENT_LAST, VIRTIO_F_VERSION_1,
+use crate::{
+    seccomp_filters::Thread, thread_helper::spawn_virtio_thread, ActivateError, ActivateResult,
+    EpollHelper, EpollHelperError, EpollHelperHandler, GuestMemoryMmap, VirtioCommon, VirtioDevice,
+    VirtioDeviceType, VirtioInterrupt, VirtioInterruptType, EPOLL_HELPER_EVENT_LAST,
+    VIRTIO_F_VERSION_1,
 };
-use crate::seccomp_filters::Thread;
-use crate::thread_helper::spawn_virtio_thread;
-use crate::GuestMemoryMmap;
-use crate::{VirtioInterrupt, VirtioInterruptType};
 use libc::EFD_NONBLOCK;
 use seccompiler::SeccompAction;
 use std::io;
 use std::mem::size_of;
 use std::os::unix::io::AsRawFd;
 use std::result;
-use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
-use std::sync::mpsc;
-use std::sync::{Arc, Barrier, Mutex};
+use std::sync::{
+    atomic::{AtomicBool, AtomicU64, Ordering},
+    mpsc, Arc, Barrier, Mutex,
+};
 use versionize::{VersionMap, Versionize, VersionizeResult};
 use versionize_derive::Versionize;
 use virtio_queue::Queue;
-use vm_memory::GuestMemory;
-use vm_memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemoryAtomic, GuestMemoryError};
-use vm_migration::VersionMapped;
-use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
+use vm_memory::{
+    Address, ByteValued, Bytes, GuestAddress, GuestMemory, GuestMemoryAtomic, GuestMemoryError,
+    GuestMemoryRegion,
+};
+use vm_migration::{
+    Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable, VersionMapped,
+};
 use vmm_sys_util::eventfd::EventFd;
 
 const QUEUE_SIZE: u16 = 128;
 const NUM_QUEUES: usize = 2;
 const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE; NUM_QUEUES];
 
-// Get resize event.
+// Resize event.
 const RESIZE_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 1;
-// New descriptors are pending on the virtio queue.
+// Inflate virtio queue event.
 const INFLATE_QUEUE_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 2;
-// New descriptors are pending on the virtio queue.
+// Deflate virtio queue event.
 const DEFLATE_QUEUE_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 3;
 
 // Size of a PFN in the balloon interface.
@@ -75,8 +77,8 @@ pub enum Error {
     MpscRecvFail(mpsc::RecvError),
     // Resize invalid argument
     ResizeInval(String),
-    // process_queue got wrong ev_type
-    ProcessQueueWrongEvType(u16),
+    // Invalid queue index
+    InvalidQueueIndex(usize),
     // Fail tp signal
     FailedSignal(io::Error),
     /// Descriptor chain is too short
@@ -172,23 +174,78 @@ impl BalloonEpollHandler {
         })
     }
 
-    fn process_queue(&mut self, ev_type: u16) -> result::Result<(), Error> {
-        let queue_index = match ev_type {
-            INFLATE_QUEUE_EVENT => 0,
-            DEFLATE_QUEUE_EVENT => 1,
-            _ => return Err(Error::ProcessQueueWrongEvType(ev_type)),
-        };
+    fn advise_memory_range(
+        memory: &GuestMemoryMmap,
+        range_base: GuestAddress,
+        range_len: usize,
+        advice: libc::c_int,
+    ) -> result::Result<(), Error> {
+        let hva = memory
+            .get_host_address(range_base)
+            .map_err(Error::GuestMemory)?;
+        // Need unsafe to do syscall madvise
+        let res =
+            unsafe { libc::madvise(hva as *mut libc::c_void, range_len as libc::size_t, advice) };
+        if res != 0 {
+            return Err(Error::MadviseFail(io::Error::last_os_error()));
+        }
+        Ok(())
+    }
 
-        let mut used_desc_heads = [0; QUEUE_SIZE as usize];
-        let mut used_count = 0;
+    fn release_memory_range(
+        memory: &GuestMemoryMmap,
+        range_base: GuestAddress,
+        range_len: usize,
+    ) -> result::Result<(), Error> {
+        let region = memory.find_region(range_base).ok_or(Error::GuestMemory(
+            GuestMemoryError::InvalidGuestAddress(range_base),
+        ))?;
+        if let Some(f_off) = region.file_offset() {
+            let offset = range_base.0 - region.start_addr().0;
+            let res = unsafe {
+                libc::fallocate64(
+                    f_off.file().as_raw_fd(),
+                    libc::FALLOC_FL_PUNCH_HOLE | libc::FALLOC_FL_KEEP_SIZE,
+                    (offset as u64 + f_off.start()) as libc::off64_t,
+                    range_len as libc::off64_t,
+                )
+            };
+
+            if res != 0 {
+                return Err(Error::FallocateFail(io::Error::last_os_error()));
+            }
+        }
+
+        Self::advise_memory_range(memory, range_base, range_len, libc::MADV_DONTNEED)
+    }
+
+    fn notify_queue(
+        &mut self,
+        queue_index: usize,
+        used_descs: Vec<(u16, u32)>,
+    ) -> result::Result<(), Error> {
+        for (desc_index, len) in used_descs.iter() {
+            self.queues[queue_index]
+                .add_used(*desc_index, *len)
+                .map_err(Error::QueueAddUsed)?;
+        }
+
+        if !used_descs.is_empty() {
+            self.signal(VirtioInterruptType::Queue(queue_index as u16))?;
+        }
+
+        Ok(())
+    }
+
+    fn process_queue(&mut self, queue_index: usize) -> result::Result<(), Error> {
+        let mut used_descs = Vec::new();
         for mut desc_chain in self.queues[queue_index]
             .iter()
             .map_err(Error::QueueIterator)?
         {
             let desc = desc_chain.next().ok_or(Error::DescriptorChainTooShort)?;
 
-            used_desc_heads[used_count] = desc_chain.head_index();
-            used_count += 1;
+            used_descs.push((desc_chain.head_index(), desc.len()));
 
             let data_chunk_size = size_of::<u32>();
 
@@ -211,63 +268,27 @@ impl BalloonEpollHandler {
                     .map_err(Error::GuestMemory)?;
                 offset += data_chunk_size as u64;
 
-                let gpa = (pfn as u64) << VIRTIO_BALLOON_PFN_SHIFT;
-                if let Ok(hva) = desc_chain.memory().get_host_address(GuestAddress(gpa)) {
-                    let advice = match ev_type {
-                        INFLATE_QUEUE_EVENT => {
-                            let region = desc_chain.memory().find_region(GuestAddress(gpa)).ok_or(
-                                Error::GuestMemory(GuestMemoryError::InvalidGuestAddress(
-                                    GuestAddress(gpa),
-                                )),
-                            )?;
-                            if let Some(f_off) = region.file_offset() {
-                                let offset = hva as usize - region.as_ptr() as usize;
-                                let res = unsafe {
-                                    libc::fallocate64(
-                                        f_off.file().as_raw_fd(),
-                                        libc::FALLOC_FL_PUNCH_HOLE | libc::FALLOC_FL_KEEP_SIZE,
-                                        (offset as u64 + f_off.start()) as libc::off64_t,
-                                        (1 << VIRTIO_BALLOON_PFN_SHIFT) as libc::off64_t,
-                                    )
-                                };
+                let range_base = GuestAddress((pfn as u64) << VIRTIO_BALLOON_PFN_SHIFT);
+                let range_len = 1 << VIRTIO_BALLOON_PFN_SHIFT;
 
-                                if res != 0 {
-                                    return Err(Error::FallocateFail(io::Error::last_os_error()));
-                                }
-                            }
-                            libc::MADV_DONTNEED
-                        }
-                        DEFLATE_QUEUE_EVENT => libc::MADV_WILLNEED,
-                        _ => return Err(Error::ProcessQueueWrongEvType(ev_type)),
-                    };
-                    // Need unsafe to do syscall madvise
-                    let res = unsafe {
-                        libc::madvise(
-                            hva as *mut libc::c_void,
-                            (1 << VIRTIO_BALLOON_PFN_SHIFT) as libc::size_t,
-                            advice,
-                        )
-                    };
-                    if res != 0 {
-                        return Err(Error::MadviseFail(io::Error::last_os_error()));
+                match queue_index {
+                    0 => {
+                        Self::release_memory_range(desc_chain.memory(), range_base, range_len)?;
                     }
-                } else {
-                    error!("Address 0x{:x} is not available", gpa);
-                    return Err(Error::InvalidRequest);
+                    1 => {
+                        Self::advise_memory_range(
+                            desc_chain.memory(),
+                            range_base,
+                            range_len,
+                            libc::MADV_WILLNEED,
+                        )?;
+                    }
+                    _ => return Err(Error::InvalidQueueIndex(queue_index)),
                 }
             }
         }
 
-        for &desc_index in &used_desc_heads[..used_count] {
-            self.queues[queue_index]
-                .add_used(desc_index, 0)
-                .map_err(Error::QueueAddUsed)?;
-        }
-        if used_count > 0 {
-            self.signal(VirtioInterruptType::Queue(queue_index as u16))?;
-        }
-
-        Ok(())
+        self.notify_queue(queue_index, used_descs)
     }
 
     fn run(
@@ -322,7 +343,7 @@ impl EpollHelperHandler for BalloonEpollHandler {
                 if let Err(e) = self.inflate_queue_evt.read() {
                     error!("Failed to get inflate queue event: {:?}", e);
                     return true;
-                } else if let Err(e) = self.process_queue(ev_type) {
+                } else if let Err(e) = self.process_queue(0) {
                     error!("Failed to signal used inflate queue: {:?}", e);
                     return true;
                 }
@@ -331,7 +352,7 @@ impl EpollHelperHandler for BalloonEpollHandler {
                 if let Err(e) = self.deflate_queue_evt.read() {
                     error!("Failed to get deflate queue event: {:?}", e);
                     return true;
-                } else if let Err(e) = self.process_queue(ev_type) {
+                } else if let Err(e) = self.process_queue(1) {
                     error!("Failed to signal used deflate queue: {:?}", e);
                     return true;
                 }
@@ -479,6 +500,9 @@ impl VirtioDevice for Balloon {
         self.common.activate(&queues, &queue_evts, &interrupt_cb)?;
         let (kill_evt, pause_evt) = self.common.dup_eventfds();
 
+        let inflate_queue_evt = queue_evts.remove(0);
+        let deflate_queue_evt = queue_evts.remove(0);
+
         let mut handler = BalloonEpollHandler {
             config: self.config.clone(),
             resize_receiver: self.resize.get_receiver().map_err(|e| {
@@ -487,8 +511,8 @@ impl VirtioDevice for Balloon {
             })?,
             queues,
             interrupt_cb,
-            inflate_queue_evt: queue_evts.remove(0),
-            deflate_queue_evt: queue_evts.remove(0),
+            inflate_queue_evt,
+            deflate_queue_evt,
             kill_evt,
             pause_evt,
         };