virtio-devices: Move guest_memory to VhostUserCommon

Rather than each device implementation holding a reference to the guest
memory move this to VhostUserCommon. This refactoring simplifies the
function signatures but also allows for methods that act on
VhostUserCommon that don't have the memory available to them.

Signed-off-by: Rob Bradford <rbradford@meta.com>

virtio-devices/src/vhost_user/blk.rs

@@ -17,7 +17,7 @@ use virtio_bindings::virtio_blk::{
     VIRTIO_BLK_F_GEOMETRY, VIRTIO_BLK_F_MQ, VIRTIO_BLK_F_RO, VIRTIO_BLK_F_SEG_MAX,
     VIRTIO_BLK_F_SIZE_MAX, VIRTIO_BLK_F_TOPOLOGY, VIRTIO_BLK_F_WRITE_ZEROES,
 };
-use vm_memory::{ByteValued, GuestMemoryAtomic};
+use vm_memory::ByteValued;
 use vm_migration::protocol::MemoryRangeTable;
 use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
 use vmm_sys_util::eventfd::EventFd;
@@ -27,7 +27,7 @@ use super::vu_common_ctrl::{VhostUserConfig, VhostUserHandle};
 use super::{DEFAULT_VIRTIO_FEATURES, Error, Result};
 use crate::seccomp_filters::Thread;
 use crate::vhost_user::{VhostUserCommon, VhostUserState};
-use crate::{GuestMemoryMmap, GuestRegionMmap, VIRTIO_F_ACCESS_PLATFORM};
+use crate::{GuestRegionMmap, VIRTIO_F_ACCESS_PLATFORM};
 
 const DEFAULT_QUEUE_NUMBER: usize = 1;
 
@@ -40,7 +40,6 @@ pub struct Blk {
     vu_common: VhostUserCommon,
     id: String,
     config: VirtioBlockConfig,
-    guest_memory: Option<GuestMemoryAtomic<GuestMemoryMmap>>,
     seccomp_action: SeccompAction,
     exit_evt: EventFd,
     access_platform_enabled: bool,
@@ -191,7 +190,6 @@ impl Blk {
             },
             id,
             config,
-            guest_memory: None,
             seccomp_action,
             exit_evt,
             access_platform_enabled,
@@ -275,7 +273,6 @@ impl VirtioDevice for Blk {
         self.vu_common
             .virtio_common
             .activate(&queues, interrupt_cb.clone())?;
-        self.guest_memory = Some(mem.clone());
 
         let backend_req_handler: Option<FrontendReqHandler<BackendReqHandler>> = None;
 
@@ -321,7 +318,7 @@ impl VirtioDevice for Blk {
         &mut self,
         region: &Arc<GuestRegionMmap>,
     ) -> std::result::Result<(), crate::Error> {
-        self.vu_common.add_memory_region(&self.guest_memory, region)
+        self.vu_common.add_memory_region(region)
     }
 }
 
@@ -350,7 +347,7 @@ impl Transportable for Blk {}
 
 impl Migratable for Blk {
     fn start_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
-        self.vu_common.start_dirty_log(&self.guest_memory)
+        self.vu_common.start_dirty_log()
     }
 
     fn stop_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
@@ -358,7 +355,7 @@ impl Migratable for Blk {
     }
 
     fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        self.vu_common.dirty_log(&self.guest_memory)
+        self.vu_common.dirty_log()
     }
 
     fn start_migration(&mut self) -> std::result::Result<(), MigratableError> {

virtio-devices/src/vhost_user/fs.rs

@@ -13,7 +13,7 @@ use serde_with::{Bytes, serde_as};
 use vhost::vhost_user::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};
 use vhost::vhost_user::{FrontendReqHandler, VhostUserFrontend, VhostUserFrontendReqHandler};
 use vm_device::UserspaceMapping;
-use vm_memory::{ByteValued, GuestMemoryAtomic};
+use vm_memory::ByteValued;
 use vm_migration::protocol::MemoryRangeTable;
 use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
 use vmm_sys_util::eventfd::EventFd;
@@ -23,8 +23,8 @@ use super::{DEFAULT_VIRTIO_FEATURES, Error, Result};
 use crate::seccomp_filters::Thread;
 use crate::vhost_user::{VhostUserCommon, VhostUserState};
 use crate::{
-    ActivateResult, GuestMemoryMmap, GuestRegionMmap, MmapRegion, VIRTIO_F_ACCESS_PLATFORM,
+    ActivateResult, GuestRegionMmap, MmapRegion, VIRTIO_F_ACCESS_PLATFORM, VirtioCommon,
-    VirtioCommon, VirtioDevice, VirtioDeviceType, VirtioSharedMemoryList,
+    VirtioDevice, VirtioDeviceType, VirtioSharedMemoryList,
 };
 
 const NUM_QUEUE_OFFSET: usize = 1;
@@ -65,7 +65,6 @@ pub struct Fs {
     // which will be automatically dropped when the device is dropped
     cache: Option<(VirtioSharedMemoryList, MmapRegion)>,
     seccomp_action: SeccompAction,
-    guest_memory: Option<GuestMemoryAtomic<GuestMemoryMmap>>,
     exit_evt: EventFd,
     access_platform_enabled: bool,
 }
@@ -198,7 +197,6 @@ impl Fs {
             config,
             cache,
             seccomp_action,
-            guest_memory: None,
             exit_evt,
             access_platform_enabled,
         })
@@ -250,7 +248,6 @@ impl VirtioDevice for Fs {
         self.vu_common
             .virtio_common
             .activate(&queues, interrupt_cb.clone())?;
-        self.guest_memory = Some(mem.clone());
 
         let backend_req_handler: Option<FrontendReqHandler<BackendReqHandler>> = None;
         // Run a dedicated thread for handling potential reconnections with
@@ -312,7 +309,7 @@ impl VirtioDevice for Fs {
         &mut self,
         region: &Arc<GuestRegionMmap>,
     ) -> std::result::Result<(), crate::Error> {
-        self.vu_common.add_memory_region(&self.guest_memory, region)
+        self.vu_common.add_memory_region(region)
     }
 
     fn userspace_mappings(&self) -> Vec<UserspaceMapping> {
@@ -355,7 +352,7 @@ impl Transportable for Fs {}
 
 impl Migratable for Fs {
     fn start_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
-        self.vu_common.start_dirty_log(&self.guest_memory)
+        self.vu_common.start_dirty_log()
     }
 
     fn stop_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
@@ -363,7 +360,7 @@ impl Migratable for Fs {
     }
 
     fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        self.vu_common.dirty_log(&self.guest_memory)
+        self.vu_common.dirty_log()
     }
 
     fn start_migration(&mut self) -> std::result::Result<(), MigratableError> {

virtio-devices/src/vhost_user/generic_vhost_user.rs

@@ -14,7 +14,6 @@ use vhost::vhost_user::message::{
 };
 use vhost::vhost_user::{FrontendReqHandler, VhostUserFrontend, VhostUserFrontendReqHandler};
 use vm_device::UserspaceMapping;
-use vm_memory::GuestMemoryAtomic;
 use vm_migration::protocol::MemoryRangeTable;
 use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
 use vmm_sys_util::eventfd::EventFd;
@@ -24,8 +23,8 @@ use super::{Error, Result};
 use crate::seccomp_filters::Thread;
 use crate::vhost_user::{VhostUserCommon, VhostUserState};
 use crate::{
-    ActivateResult, GuestMemoryMmap, GuestRegionMmap, MmapRegion, VIRTIO_F_ACCESS_PLATFORM,
+    ActivateResult, GuestRegionMmap, MmapRegion, VIRTIO_F_ACCESS_PLATFORM, VirtioCommon,
-    VirtioCommon, VirtioDevice, VirtioInterrupt, VirtioInterruptType, VirtioSharedMemoryList,
+    VirtioDevice, VirtioInterrupt, VirtioInterruptType, VirtioSharedMemoryList,
 };
 
 pub type State = VhostUserState<()>;
@@ -53,7 +52,6 @@ pub struct GenericVhostUser {
     // which will be automatically dropped when the device is dropped
     cache: Option<(VirtioSharedMemoryList, MmapRegion)>,
     seccomp_action: SeccompAction,
-    guest_memory: Option<GuestMemoryAtomic<GuestMemoryMmap>>,
     exit_evt: EventFd,
     access_platform_enabled: bool,
     cfg_warning: AtomicBool,
@@ -172,7 +170,6 @@ since the backend only supports {backend_num_queues}\n",
             id,
             cache,
             seccomp_action,
-            guest_memory: None,
             exit_evt,
             access_platform_enabled,
             cfg_warning: AtomicBool::new(false),
@@ -288,7 +285,6 @@ impl VirtioDevice for GenericVhostUser {
         self.vu_common
             .virtio_common
             .activate(&queues, interrupt_cb.clone())?;
-        self.guest_memory = Some(mem.clone());
 
         let has_backend_req = self.vu_common.acked_protocol_features
             & VhostUserProtocolFeatures::BACKEND_REQ.bits()
@@ -374,7 +370,7 @@ impl VirtioDevice for GenericVhostUser {
         &mut self,
         region: &Arc<GuestRegionMmap>,
     ) -> std::result::Result<(), crate::Error> {
-        self.vu_common.add_memory_region(&self.guest_memory, region)
+        self.vu_common.add_memory_region(region)
     }
 
     fn userspace_mappings(&self) -> Vec<UserspaceMapping> {
@@ -417,7 +413,7 @@ impl Transportable for GenericVhostUser {}
 
 impl Migratable for GenericVhostUser {
     fn start_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
-        self.vu_common.start_dirty_log(&self.guest_memory)
+        self.vu_common.start_dirty_log()
     }
 
     fn stop_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
@@ -425,7 +421,7 @@ impl Migratable for GenericVhostUser {
     }
 
     fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        self.vu_common.dirty_log(&self.guest_memory)
+        self.vu_common.dirty_log()
     }
 
     fn start_migration(&mut self) -> std::result::Result<(), MigratableError> {

virtio-devices/src/vhost_user/mod.rs

@@ -459,6 +459,7 @@ impl<C> VhostUserState<C> {
 pub struct VhostUserCommon {
     pub virtio_common: VirtioCommon,
     pub vu: Option<Arc<Mutex<VhostUserHandle>>>,
+    pub guest_memory: Option<GuestMemoryAtomic<GuestMemoryMmap>>,
     pub acked_protocol_features: u64,
     pub socket_path: String,
     pub vu_num_queues: usize,
@@ -481,6 +482,8 @@ impl VhostUserCommon {
         kill_evt: EventFd,
         pause_evt: EventFd,
     ) -> std::result::Result<VhostUserEpollHandler<T>, ActivateError> {
+        self.guest_memory = Some(mem.clone());
+
         if self.disconnected.load(Ordering::Relaxed) {
             warn!(
                 "Not activating disconnected vhost-user device for socket {}",
@@ -626,17 +629,13 @@ impl VhostUserCommon {
         self.vu = None;
     }
 
-    fn add_memory_region_internal(
+    fn add_memory_region_internal(&self, region: &Arc<GuestRegionMmap>) -> Result<()> {
-        &self,
-        guest_memory: &Option<GuestMemoryAtomic<GuestMemoryMmap>>,
-        region: &Arc<GuestRegionMmap>,
-    ) -> Result<()> {
         if let Some(vu) = &self.vu {
             if self.acked_protocol_features & VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS.bits()
                 != 0
             {
                 return vu.lock().unwrap().add_memory_region(region);
-            } else if let Some(guest_memory) = guest_memory {
+            } else if let Some(guest_memory) = &self.guest_memory {
                 return vu
                     .lock()
                     .unwrap()
@@ -649,7 +648,6 @@ impl VhostUserCommon {
 
     pub fn add_memory_region(
         &mut self,
-        guest_memory: &Option<GuestMemoryAtomic<GuestMemoryMmap>>,
         region: &Arc<GuestRegionMmap>,
     ) -> std::result::Result<(), crate::Error> {
         if self.disconnected.load(Ordering::Relaxed) {
@@ -659,7 +657,7 @@ impl VhostUserCommon {
             );
         }
 
-        if let Err(e) = self.add_memory_region_internal(guest_memory, region) {
+        if let Err(e) = self.add_memory_region_internal(region) {
             if e.is_transport_lost() {
                 warn!(
                     "Failed updating memory on vhost-user backend for socket {}: {e:?}; \
@@ -784,12 +782,9 @@ impl VhostUserCommon {
         Ok(snapshot)
     }
 
-    pub fn start_dirty_log(
+    pub fn start_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
-        &mut self,
-        guest_memory: &Option<GuestMemoryAtomic<GuestMemoryMmap>>,
-    ) -> std::result::Result<(), MigratableError> {
         if let Some(vu) = &self.vu {
-            if let Some(guest_memory) = guest_memory {
+            if let Some(guest_memory) = &self.guest_memory {
                 let last_ram_addr = guest_memory.memory().last_addr().raw_value();
                 vu.lock()
                     .unwrap()
@@ -821,12 +816,9 @@ impl VhostUserCommon {
         }
     }
 
-    pub fn dirty_log(
+    pub fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        &mut self,
-        guest_memory: &Option<GuestMemoryAtomic<GuestMemoryMmap>>,
-    ) -> std::result::Result<MemoryRangeTable, MigratableError> {
         if let Some(vu) = &self.vu {
-            if let Some(guest_memory) = guest_memory {
+            if let Some(guest_memory) = &self.guest_memory {
                 let last_ram_addr = guest_memory.memory().last_addr().raw_value();
                 vu.lock().unwrap().dirty_log(last_ram_addr).map_err(|e| {
                     MigratableError::DirtyLog(anyhow!(

virtio-devices/src/vhost_user/net.rs

@@ -18,7 +18,7 @@ use virtio_bindings::virtio_net::{
 };
 use virtio_bindings::virtio_ring::VIRTIO_RING_F_EVENT_IDX;
 use virtio_queue::QueueT;
-use vm_memory::{ByteValued, GuestMemoryAtomic};
+use vm_memory::ByteValued;
 use vm_migration::protocol::MemoryRangeTable;
 use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
 use vmm_sys_util::eventfd::EventFd;
@@ -27,8 +27,8 @@ use crate::seccomp_filters::Thread;
 use crate::vhost_user::vu_common_ctrl::{VhostUserConfig, VhostUserHandle};
 use crate::vhost_user::{DEFAULT_VIRTIO_FEATURES, Error, Result, VhostUserCommon, VhostUserState};
 use crate::{
-    ActivateResult, GuestMemoryMmap, GuestRegionMmap, NetCtrlEpollHandler,
+    ActivateResult, GuestRegionMmap, NetCtrlEpollHandler, VIRTIO_F_ACCESS_PLATFORM, VirtioCommon,
-    VIRTIO_F_ACCESS_PLATFORM, VirtioCommon, VirtioDevice, VirtioDeviceType,
+    VirtioDevice, VirtioDeviceType,
 };
 
 const DEFAULT_QUEUE_NUMBER: usize = 2;
@@ -42,7 +42,6 @@ pub struct Net {
     vu_common: VhostUserCommon,
     id: String,
     config: VirtioNetConfig,
-    guest_memory: Option<GuestMemoryAtomic<GuestMemoryMmap>>,
     seccomp_action: SeccompAction,
     exit_evt: EventFd,
     access_platform_enabled: bool,
@@ -218,7 +217,6 @@ impl Net {
                 ..Default::default()
             },
             config,
-            guest_memory: None,
             seccomp_action,
             exit_evt,
             access_platform_enabled,
@@ -271,7 +269,6 @@ impl VirtioDevice for Net {
         self.vu_common
             .virtio_common
             .activate(&queues, interrupt_cb.clone())?;
-        self.guest_memory = Some(mem.clone());
 
         let num_queues = queues.len();
         let event_idx = self
@@ -370,7 +367,7 @@ impl VirtioDevice for Net {
         &mut self,
         region: &Arc<GuestRegionMmap>,
     ) -> std::result::Result<(), crate::Error> {
-        self.vu_common.add_memory_region(&self.guest_memory, region)
+        self.vu_common.add_memory_region(region)
     }
 }
 
@@ -399,7 +396,7 @@ impl Transportable for Net {}
 
 impl Migratable for Net {
     fn start_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
-        self.vu_common.start_dirty_log(&self.guest_memory)
+        self.vu_common.start_dirty_log()
     }
 
     fn stop_dirty_log(&mut self) -> std::result::Result<(), MigratableError> {
@@ -407,7 +404,7 @@ impl Migratable for Net {
     }
 
     fn dirty_log(&mut self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        self.vu_common.dirty_log(&self.guest_memory)
+        self.vu_common.dirty_log()
     }
 
     fn start_migration(&mut self) -> std::result::Result<(), MigratableError> {