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.

Finally, fix saving/restoring the state of ResourceAllocator. We were
actually not resetting it correctly upon snapshot restore. The reason
why this was not a problem is that we don't actually need to perform any
new allocations post restore at the moment. However, like this we are
ready when we need to perform any hot-plugging operations. Also, add a
unit-test to ensure that this logic works correctly.

Signed-off-by: Babis Chalios <bchalios@amazon.es>

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 };
@@ -249,18 +249,19 @@ mod tests {
         let mut writer = AcpiTableWriter {
             mem: vmm.vm.guest_memory(),
         };
+        let mut resource_allocator = vmm.vm.resource_allocator();
 
         // 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 resource_allocator, &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 resource_allocator, &mut sdt)
             .unwrap_err();
         assert!(
             matches!(
@@ -275,27 +276,27 @@ 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 resource_allocator, &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 resource_allocator, &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 resource_allocator, &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 resource_allocator, &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 resource_allocator, &mut sdt)
             .unwrap();
         assert_eq!(addr, SYSTEM_MEM_START + 4115);
     }
@@ -312,11 +313,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

@@ -1,11 +1,14 @@
 // Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::sync::Mutex;
+
 use serde::{Deserialize, Serialize};
 
 use crate::Kvm;
 use crate::arch::aarch64::gic::GicState;
 use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryState};
+use crate::vstate::resources::ResourceAllocator;
 use crate::vstate::vm::{VmCommon, VmError};
 
 /// Structure representing the current architecture's understand of what a "virtual machine" is.
@@ -74,6 +77,7 @@ impl ArchVm {
                 .get_irqchip()
                 .save_device(mpidrs)
                 .map_err(ArchVmError::SaveGic)?,
+            resource_allocator: self.resource_allocator().clone(),
         })
     }
 
@@ -86,6 +90,7 @@ impl ArchVm {
         self.get_irqchip()
             .restore_device(mpidrs, &state.gic)
             .map_err(ArchVmError::RestoreGic)?;
+        self.common.resource_allocator = Mutex::new(state.resource_allocator.clone());
 
         Ok(())
     }
@@ -98,4 +103,6 @@ pub struct VmState {
     pub memory: GuestMemoryState,
     /// GIC state.
     pub gic: GicState,
+    /// resource allocator
+    pub resource_allocator: ResourceAllocator,
 }

src/vmm/src/arch/aarch64/vm.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.resource_allocator(),
         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.resource_allocator(),
         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

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::fmt;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use kvm_bindings::{
     KVM_CLOCK_TSC_STABLE, KVM_IRQCHIP_IOAPIC, KVM_IRQCHIP_PIC_MASTER, KVM_IRQCHIP_PIC_SLAVE,
@@ -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`]
@@ -142,6 +142,7 @@ impl ArchVm {
         self.fd()
             .set_irqchip(&state.ioapic)
             .map_err(ArchVmError::SetIrqChipIoAPIC)?;
+        self.common.resource_allocator = Mutex::new(state.resource_allocator.clone());
         Ok(())
     }
 
@@ -195,7 +196,7 @@ impl ArchVm {
 
         Ok(VmState {
             memory: self.common.guest_memory.describe(),
-            resource_allocator: self.common.resource_allocator.save(),
+            resource_allocator: self.resource_allocator().save(),
             pitstate,
             clock,
             pic_master,
@@ -221,7 +222,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.