refactor: simplify ResourceAllocator internals

Instead of storing internal allocators of ResourceAllocator within an
Arc<Mutex<>> container, just store `ResourceAllocator` itself in an
`Arc<Mutex<>>`.

Apart from that, we get rid of the `ResourceAllocatorState` state
object, and just clone `ResourceAllocator` itself when we want to
save/restore.

Also, make the creation of `ResourceAllocato` infallible, since we know
that the ranges we are using are correct.

Signed-off-by: Babis Chalios <[email protected]>

Author: bchalios

Cargo.lock

@@ -18,7 +18,7 @@ libc = "0.2.172"
 log = "0.4.27"
 serde = { version = "1.0.219", features = ["derive"] }
 thiserror = "2.0.12"
-vm-allocator = "0.1.2"
+vm-allocator = "0.1.3"
 vm-device = { path = "../vm-device" }
 vm-memory = { version = "0.16.1", features = [
   "backend-mmap",

src/pci/Cargo.toml

@@ -52,7 +52,7 @@ userfaultfd = "0.8.1"
 utils = { path = "../utils" }
 uuid = "1.16.0"
 vhost = { version = "0.14.0", features = ["vhost-user-frontend"] }
-vm-allocator = { version = "0.1.2", features = ["serde"] }
+vm-allocator = { version = "0.1.3", features = ["serde"] }
 vm-device = { path = "../vm-device" }
 vm-memory = { version = "0.16.2", features = [
   "backend-mmap",

src/vmm/Cargo.toml

@@ -54,7 +54,7 @@ impl AcpiTableWriter<'_> {
     /// buffer. It returns the address in which it wrote the table.
     fn write_acpi_table<S>(
         &mut self,
-        resource_allocator: &ResourceAllocator,
+        resource_allocator: &mut ResourceAllocator,
         table: &mut S,
     ) -> Result<u64, AcpiError>
     where
@@ -83,7 +83,7 @@ impl AcpiTableWriter<'_> {
     fn build_dsdt(
         &mut self,
         device_manager: &mut DeviceManager,
-        resource_allocator: &ResourceAllocator,
+        resource_allocator: &mut ResourceAllocator,
     ) -> Result<u64, AcpiError> {
         let mut dsdt_data = Vec::new();
 
@@ -111,7 +111,7 @@ impl AcpiTableWriter<'_> {
     /// This includes a pointer with the location of the DSDT in guest memory
     fn build_fadt(
         &mut self,
-        resource_allocator: &ResourceAllocator,
+        resource_allocator: &mut ResourceAllocator,
         dsdt_addr: u64,
     ) -> Result<u64, AcpiError> {
         let mut fadt = Fadt::new(OEM_ID, *b"FCVMFADT", OEM_REVISION);
@@ -129,7 +129,7 @@ impl AcpiTableWriter<'_> {
     /// This includes information about the interrupt controllers supported in the platform
     fn build_madt(
         &mut self,
-        resource_allocator: &ResourceAllocator,
+        resource_allocator: &mut ResourceAllocator,
         nr_vcpus: u8,
     ) -> Result<u64, AcpiError> {
         let mut madt = Madt::new(
@@ -147,7 +147,7 @@ impl AcpiTableWriter<'_> {
     /// Currently, we pass to the guest just FADT and MADT tables.
     fn build_xsdt(
         &mut self,
-        resource_allocator: &ResourceAllocator,
+        resource_allocator: &mut ResourceAllocator,
         fadt_addr: u64,
         madt_addr: u64,
         mcfg_addr: u64,
@@ -164,7 +164,7 @@ impl AcpiTableWriter<'_> {
     /// Build the MCFG table for the guest.
     fn build_mcfg(
         &mut self,
-        resource_allocator: &ResourceAllocator,
+        resource_allocator: &mut ResourceAllocator,
         pci_mmio_config_addr: u64,
     ) -> Result<u64, AcpiError> {
         let mut mcfg = Mcfg::new(OEM_ID, *b"FCMVMCFG", OEM_REVISION, pci_mmio_config_addr);
@@ -197,7 +197,7 @@ impl AcpiTableWriter<'_> {
 pub(crate) fn create_acpi_tables(
     mem: &GuestMemoryMmap,
     device_manager: &mut DeviceManager,
-    resource_allocator: &ResourceAllocator,
+    resource_allocator: &mut ResourceAllocator,
     vcpus: &[Vcpu],
 ) -> Result<(), AcpiError> {
     let mut writer = AcpiTableWriter { mem };
@@ -253,14 +253,20 @@ mod tests {
         // This should succeed
         let mut sdt = MockSdt(vec![0; 4096]);
         let addr = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START);
 
         // Let's try to write two 4K pages plus one byte
         let mut sdt = MockSdt(vec![0; usize::try_from(SYSTEM_MEM_SIZE + 1).unwrap()]);
         let err = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap_err();
         assert!(
             matches!(
@@ -275,27 +281,42 @@ mod tests {
         // succeed.
         let mut sdt = MockSdt(vec![0; 5]);
         let addr = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START + 4096);
         let mut sdt = MockSdt(vec![0; 2]);
         let addr = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START + 4101);
         let mut sdt = MockSdt(vec![0; 4]);
         let addr = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START + 4103);
         let mut sdt = MockSdt(vec![0; 8]);
         let addr = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START + 4107);
         let mut sdt = MockSdt(vec![0; 16]);
         let addr = writer
-            .write_acpi_table(&vmm.vm.common.resource_allocator, &mut sdt)
+            .write_acpi_table(
+                &mut vmm.vm.common.resource_allocator.lock().unwrap(),
+                &mut sdt,
+            )
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START + 4115);
     }
@@ -312,11 +333,11 @@ mod tests {
         let mut writer = AcpiTableWriter {
             mem: vm.guest_memory(),
         };
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
         let mut sdt = MockSdt(vec![0; usize::try_from(SYSTEM_MEM_SIZE).unwrap()]);
         let err = writer
-            .write_acpi_table(&resource_allocator, &mut sdt)
+            .write_acpi_table(&mut resource_allocator, &mut sdt)
             .unwrap_err();
         assert!(
             matches!(

src/vmm/src/acpi/mod.rs

@@ -217,7 +217,7 @@ pub fn configure_system_for_boot(
     // Note that this puts the mptable at the last 1k of Linux's 640k base RAM
     mptable::setup_mptable(
         vm.guest_memory(),
-        &vm.common.resource_allocator,
+        &mut vm.common.resource_allocator.lock().expect("Poisoned lock"),
         vcpu_config.vcpu_count,
     )
     .map_err(ConfigurationError::MpTableSetup)?;
@@ -241,7 +241,7 @@ pub fn configure_system_for_boot(
     create_acpi_tables(
         vm.guest_memory(),
         device_manager,
-        &vm.common.resource_allocator,
+        &mut vm.common.resource_allocator.lock().expect("Poisoned lock"),
         vcpus,
     )?;
     Ok(())
@@ -607,8 +607,8 @@ mod tests {
     fn test_system_configuration() {
         let no_vcpus = 4;
         let gm = single_region_mem(0x10000);
-        let resource_allocator = ResourceAllocator::new().unwrap();
-        let err = mptable::setup_mptable(&gm, &resource_allocator, 1);
+        let mut resource_allocator = ResourceAllocator::new();
+        let err = mptable::setup_mptable(&gm, &mut resource_allocator, 1);
         assert!(matches!(
             err.unwrap_err(),
             mptable::MptableError::NotEnoughMemory
@@ -617,24 +617,24 @@ mod tests {
         // Now assigning some memory that falls before the 32bit memory hole.
         let mem_size = mib_to_bytes(128);
         let gm = arch_mem(mem_size);
-        let resource_allocator = ResourceAllocator::new().unwrap();
-        mptable::setup_mptable(&gm, &resource_allocator, no_vcpus).unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
+        mptable::setup_mptable(&gm, &mut resource_allocator, no_vcpus).unwrap();
         configure_64bit_boot(&gm, GuestAddress(0), 0, &None).unwrap();
         configure_pvh(&gm, GuestAddress(0), &None).unwrap();
 
         // Now assigning some memory that is equal to the start of the 32bit memory hole.
         let mem_size = mib_to_bytes(3328);
         let gm = arch_mem(mem_size);
-        let resource_allocator = ResourceAllocator::new().unwrap();
-        mptable::setup_mptable(&gm, &resource_allocator, no_vcpus).unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
+        mptable::setup_mptable(&gm, &mut resource_allocator, no_vcpus).unwrap();
         configure_64bit_boot(&gm, GuestAddress(0), 0, &None).unwrap();
         configure_pvh(&gm, GuestAddress(0), &None).unwrap();
 
         // Now assigning some memory that falls after the 32bit memory hole.
         let mem_size = mib_to_bytes(3330);
         let gm = arch_mem(mem_size);
-        let resource_allocator = ResourceAllocator::new().unwrap();
-        mptable::setup_mptable(&gm, &resource_allocator, no_vcpus).unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
+        mptable::setup_mptable(&gm, &mut resource_allocator, no_vcpus).unwrap();
         configure_64bit_boot(&gm, GuestAddress(0), 0, &None).unwrap();
         configure_pvh(&gm, GuestAddress(0), &None).unwrap();
     }

src/vmm/src/arch/x86_64/mod.rs

@@ -116,7 +116,7 @@ fn compute_mp_size(num_cpus: u8) -> usize {
 /// Performs setup of the MP table for the given `num_cpus`.
 pub fn setup_mptable(
     mem: &GuestMemoryMmap,
-    resource_allocator: &ResourceAllocator,
+    resource_allocator: &mut ResourceAllocator,
     num_cpus: u8,
 ) -> Result<(), MptableError> {
     if num_cpus > MAX_SUPPORTED_CPUS {
@@ -334,27 +334,27 @@ mod tests {
     fn bounds_check() {
         let num_cpus = 4;
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(num_cpus));
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
-        setup_mptable(&mem, &resource_allocator, num_cpus).unwrap();
+        setup_mptable(&mem, &mut resource_allocator, num_cpus).unwrap();
     }
 
     #[test]
     fn bounds_check_fails() {
         let num_cpus = 4;
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(num_cpus) - 1);
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
-        setup_mptable(&mem, &resource_allocator, num_cpus).unwrap_err();
+        setup_mptable(&mem, &mut resource_allocator, num_cpus).unwrap_err();
     }
 
     #[test]
     fn mpf_intel_checksum() {
         let num_cpus = 1;
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(num_cpus));
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
-        setup_mptable(&mem, &resource_allocator, num_cpus).unwrap();
+        setup_mptable(&mem, &mut resource_allocator, num_cpus).unwrap();
 
         let mpf_intel: mpspec::mpf_intel = mem.read_obj(GuestAddress(SYSTEM_MEM_START)).unwrap();
 
@@ -365,9 +365,9 @@ mod tests {
     fn mpc_table_checksum() {
         let num_cpus = 4;
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(num_cpus));
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
-        setup_mptable(&mem, &resource_allocator, num_cpus).unwrap();
+        setup_mptable(&mem, &mut resource_allocator, num_cpus).unwrap();
 
         let mpf_intel: mpspec::mpf_intel = mem.read_obj(GuestAddress(SYSTEM_MEM_START)).unwrap();
         let mpc_offset = GuestAddress(u64::from(mpf_intel.physptr));
@@ -388,9 +388,9 @@ mod tests {
     fn mpc_entry_count() {
         let num_cpus = 1;
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(num_cpus));
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
-        setup_mptable(&mem, &resource_allocator, num_cpus).unwrap();
+        setup_mptable(&mem, &mut resource_allocator, num_cpus).unwrap();
 
         let mpf_intel: mpspec::mpf_intel = mem.read_obj(GuestAddress(SYSTEM_MEM_START)).unwrap();
         let mpc_offset = GuestAddress(u64::from(mpf_intel.physptr));
@@ -419,8 +419,9 @@ mod tests {
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(MAX_SUPPORTED_CPUS));
 
         for i in 0..MAX_SUPPORTED_CPUS {
-            let resource_allocator = ResourceAllocator::new().unwrap();
-            setup_mptable(&mem, &resource_allocator, i).unwrap();
+            let mut resource_allocator = ResourceAllocator::new();
+
+            setup_mptable(&mem, &mut resource_allocator, i).unwrap();
 
             let mpf_intel: mpspec::mpf_intel =
                 mem.read_obj(GuestAddress(SYSTEM_MEM_START)).unwrap();
@@ -450,9 +451,9 @@ mod tests {
     fn cpu_entry_count_max() {
         let cpus = MAX_SUPPORTED_CPUS + 1;
         let mem = single_region_mem_at(SYSTEM_MEM_START, compute_mp_size(cpus));
-        let resource_allocator = ResourceAllocator::new().unwrap();
+        let mut resource_allocator = ResourceAllocator::new();
 
-        let result = setup_mptable(&mem, &resource_allocator, cpus).unwrap_err();
+        let result = setup_mptable(&mem, &mut resource_allocator, cpus).unwrap_err();
         assert_eq!(result, MptableError::TooManyCpus);
     }
 }

src/vmm/src/arch/x86_64/mptable.rs

@@ -15,7 +15,7 @@ use crate::arch::x86_64::msr::MsrError;
 use crate::snapshot::Persist;
 use crate::utils::u64_to_usize;
 use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryState};
-use crate::vstate::resources::ResourceAllocatorState;
+use crate::vstate::resources::ResourceAllocator;
 use crate::vstate::vm::{VmCommon, VmError};
 
 /// Error type for [`Vm::restore_state`]
@@ -195,7 +195,12 @@ impl ArchVm {
 
         Ok(VmState {
             memory: self.common.guest_memory.describe(),
-            resource_allocator: self.common.resource_allocator.save(),
+            resource_allocator: self
+                .common
+                .resource_allocator
+                .lock()
+                .expect("Poisoned lock")
+                .save(),
             pitstate,
             clock,
             pic_master,
@@ -221,7 +226,7 @@ pub struct VmState {
     /// guest memory state
     pub memory: GuestMemoryState,
     /// resource allocator
-    pub resource_allocator: ResourceAllocatorState,
+    pub resource_allocator: ResourceAllocator,
     pitstate: kvm_pit_state2,
     clock: kvm_clock_data,
     // TODO: rename this field to adopt inclusive language once Linux updates it, too.

src/vmm/src/arch/x86_64/vm.rs

@@ -254,7 +254,10 @@ pub fn build_microvm_for_boot(
 
     #[cfg(target_arch = "aarch64")]
     if vcpus[0].kvm_vcpu.supports_pvtime() {
-        setup_pvtime(&vm.common.resource_allocator, &mut vcpus)?;
+        setup_pvtime(
+            &mut vm.common.resource_allocator.lock().expect("Poisoned lock"),
+            &mut vcpus,
+        )?;
     } else {
         log::warn!("Vcpus do not support pvtime, steal time will not be reported to guest");
     }
@@ -515,7 +518,7 @@ const STEALTIME_STRUCT_MEM_SIZE: u64 = 64;
 /// Helper method to allocate steal time region
 #[cfg(target_arch = "aarch64")]
 fn allocate_pvtime_region(
-    resource_allocator: &ResourceAllocator,
+    resource_allocator: &mut ResourceAllocator,
     vcpu_count: usize,
     policy: vm_allocator::AllocPolicy,
 ) -> Result<GuestAddress, StartMicrovmError> {
@@ -529,7 +532,7 @@ fn allocate_pvtime_region(
 /// Sets up pvtime for all vcpus
 #[cfg(target_arch = "aarch64")]
 fn setup_pvtime(
-    resource_allocator: &ResourceAllocator,
+    resource_allocator: &mut ResourceAllocator,
     vcpus: &mut [Vcpu],
 ) -> Result<(), StartMicrovmError> {
     // Alloc sys mem for steal time region