pci: vfio: Use MemorySlotAllocator for allocating memory slots

Adapt the existing code to replace the closure with the new of the new MemorySlotAllocator. Signed-off-by: Rob Bradford <[email protected]>
UiPath · Nov 24, 2024 · 81f8a27 · 81f8a27
1 parent e3bd5e9
commit 81f8a27
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/pci/src/vfio.rs b/pci/src/vfio.rs
@@ -23,7 +23,7 @@ use vfio_ioctls::{
 use vm_allocator::page_size::{
     align_page_size_down, align_page_size_up, is_4k_aligned, is_4k_multiple, is_page_size_aligned,
 };
-use vm_allocator::{AddressAllocator, SystemAllocator};
+use vm_allocator::{AddressAllocator, MemorySlotAllocator, SystemAllocator};
 use vm_device::dma_mapping::ExternalDmaMapping;
 use vm_device::interrupt::{
     InterruptIndex, InterruptManager, InterruptSourceGroup, MsiIrqGroupConfig,
@@ -1416,7 +1416,7 @@ pub struct VfioPciDevice {
     container: Arc<VfioContainer>,
     common: VfioCommon,
     iommu_attached: bool,
-    memory_slot: Arc<dyn Fn() -> u32 + Send + Sync>,
+    memory_slot_allocator: MemorySlotAllocator,
 }
 
 impl VfioPciDevice {
@@ -1431,7 +1431,7 @@ impl VfioPciDevice {
         legacy_interrupt_group: Option<Arc<dyn InterruptSourceGroup>>,
         iommu_attached: bool,
         bdf: PciBdf,
-        memory_slot: Arc<dyn Fn() -> u32 + Send + Sync>,
+        memory_slot_allocator: MemorySlotAllocator,
         snapshot: Option<Snapshot>,
         x_nv_gpudirect_clique: Option<u8>,
     ) -> Result<Self, VfioPciError> {
@@ -1457,7 +1457,7 @@ impl VfioPciDevice {
             container,
             common,
             iommu_attached,
-            memory_slot,
+            memory_slot_allocator,
         };
 
         Ok(vfio_pci_device)
@@ -1635,7 +1635,7 @@ impl VfioPciDevice {
                     }
 
                     let user_memory_region = UserMemoryRegion {
-                        slot: (self.memory_slot)(),
+                        slot: self.memory_slot_allocator.next_memory_slot(),
                         start: region.start.0 + area.offset,
                         size: area.size,
                         host_addr: host_addr as u64,

diff --git a/pci/src/vfio_user.rs b/pci/src/vfio_user.rs
@@ -13,7 +13,7 @@ use thiserror::Error;
 use vfio_bindings::bindings::vfio::*;
 use vfio_ioctls::VfioIrq;
 use vfio_user::{Client, Error as VfioUserError};
-use vm_allocator::{AddressAllocator, SystemAllocator};
+use vm_allocator::{AddressAllocator, MemorySlotAllocator, SystemAllocator};
 use vm_device::dma_mapping::ExternalDmaMapping;
 use vm_device::interrupt::{InterruptManager, InterruptSourceGroup, MsiIrqGroupConfig};
 use vm_device::{BusDevice, Resource};
@@ -35,7 +35,7 @@ pub struct VfioUserPciDevice {
     vm: Arc<dyn hypervisor::Vm>,
     client: Arc<Mutex<Client>>,
     common: VfioCommon,
-    memory_slot: Arc<dyn Fn() -> u32 + Send + Sync>,
+    memory_slot_allocator: MemorySlotAllocator,
 }
 
 #[derive(Error, Debug)]
@@ -74,7 +74,7 @@ impl VfioUserPciDevice {
         msi_interrupt_manager: Arc<dyn InterruptManager<GroupConfig = MsiIrqGroupConfig>>,
         legacy_interrupt_group: Option<Arc<dyn InterruptSourceGroup>>,
         bdf: PciBdf,
-        memory_slot: Arc<dyn Fn() -> u32 + Send + Sync>,
+        memory_slot_allocator: MemorySlotAllocator,
         snapshot: Option<Snapshot>,
     ) -> Result<Self, VfioUserPciDeviceError> {
         let resettable = client.lock().unwrap().resettable();
@@ -106,7 +106,7 @@ impl VfioUserPciDevice {
             vm: vm.clone(),
             client,
             common,
-            memory_slot,
+            memory_slot_allocator,
         })
     }
 
@@ -178,7 +178,7 @@ impl VfioUserPciDevice {
                     }
 
                     let user_memory_region = UserMemoryRegion {
-                        slot: (self.memory_slot)(),
+                        slot: self.memory_slot_allocator.next_memory_slot(),
                         start: mmio_region.start.0 + s.offset,
                         size: s.size,
                         host_addr: host_addr as u64,

diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
@@ -3432,7 +3432,7 @@ impl DeviceManager {
             legacy_interrupt_group,
             device_cfg.iommu,
             pci_device_bdf,
-            Arc::new(move || memory_manager.lock().unwrap().allocate_memory_slot()),
+            memory_manager.lock().unwrap().memory_slot_allocator(),
             vm_migration::snapshot_from_id(self.snapshot.as_ref(), vfio_name.as_str()),
             device_cfg.x_nv_gpudirect_clique,
         )
@@ -3596,7 +3596,7 @@ impl DeviceManager {
             self.msi_interrupt_manager.clone(),
             legacy_interrupt_group,
             pci_device_bdf,
-            Arc::new(move || memory_manager.lock().unwrap().allocate_memory_slot()),
+            memory_manager.lock().unwrap().memory_slot_allocator(),
             vm_migration::snapshot_from_id(self.snapshot.as_ref(), vfio_user_name.as_str()),
         )
         .map_err(DeviceManagerError::VfioUserCreate)?;