feat: added configurable vcpu features to cpu templates

 Added ability to configure vcpu features from cpu templates.
To make code clearer I moved `vcpu.init` out of `create_vcpus` method.
Because of this I had to update restoration path to init vcpus earlier.
This in turn removed need for `VcupEvent::RestoreState` and
`restore_vcpu_states` method.
Because this is the first FC release with this feature it is
invalid to create snapshots for older versions if any of
additional vcpu features are requested, so there is an
additional check for this.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>

src/vmm/src/builder.rs

@@ -450,7 +450,7 @@ pub fn build_microvm_from_snapshot(
     })?;
 
     // Build Vmm.
-    let (mut vmm, vcpus) = create_vmm_and_vcpus(
+    let (mut vmm, mut vcpus) = create_vmm_and_vcpus(
         instance_info,
         event_manager,
         guest_memory.clone(),
@@ -474,13 +474,34 @@ pub fn build_microvm_from_snapshot(
         }
     }
 
+    // Restore vcpus kvm state.
     #[cfg(target_arch = "aarch64")]
     {
+        for (i, (vcpu, state)) in vcpus
+            .iter_mut()
+            .zip(microvm_state.vcpu_states.iter())
+            .enumerate()
+        {
+            vcpu.kvm_vcpu
+                .restore_state(vmm.vm.fd(), state, i as u8)
+                .map_err(crate::vstate::vcpu::VcpuError::VcpuResponse)
+                .map_err(RestoreVcpusError::RestoreVcpuState)
+                .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
+        }
         let mpidrs = construct_kvm_mpidrs(&microvm_state.vcpu_states);
         // Restore kvm vm state.
         vmm.vm.restore_state(&mpidrs, &microvm_state.vm_state)?;
     }
 
+    #[cfg(target_arch = "x86_64")]
+    for (vcpu, state) in vcpus.iter_mut().zip(microvm_state.vcpu_states.iter()) {
+        vcpu.kvm_vcpu
+            .restore_state(state)
+            .map_err(crate::vstate::vcpu::VcpuError::VcpuResponse)
+            .map_err(RestoreVcpusError::RestoreVcpuState)
+            .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
+    }
+
     // Restore kvm vm state.
     #[cfg(target_arch = "x86_64")]
     vmm.vm.restore_state(&microvm_state.vm_state)?;
@@ -520,9 +541,6 @@ pub fn build_microvm_from_snapshot(
             .clone(),
     )?;
 
-    // Restore vcpus kvm state.
-    vmm.restore_vcpu_states(microvm_state.vcpu_states)?;
-
     let vmm = Arc::new(Mutex::new(vmm));
     event_manager.add_subscriber(vmm.clone());
 
@@ -723,11 +741,7 @@ fn create_vcpus(vm: &Vm, vcpu_count: u8, exit_evt: &EventFd) -> Result<Vec<Vcpu>
     let mut vcpus = Vec::with_capacity(vcpu_count as usize);
     for cpu_idx in 0..vcpu_count {
         let exit_evt = exit_evt.try_clone().map_err(VmmError::EventFd)?;
-
         let vcpu = Vcpu::new(cpu_idx, vm, exit_evt).map_err(VmmError::VcpuCreate)?;
-        #[cfg(target_arch = "aarch64")]
-        vcpu.kvm_vcpu.init(vm.fd()).map_err(VmmError::VcpuInit)?;
-
         vcpus.push(vcpu);
     }
     Ok(vcpus)
@@ -764,6 +778,14 @@ pub fn configure_system_for_boot(
     let cpu_config = {
         use crate::arch::aarch64::regs::Aarch64RegisterVec;
         use crate::arch::aarch64::vcpu::get_registers;
+
+        for vcpu in vcpus.iter_mut() {
+            vcpu.kvm_vcpu
+                .init(vmm.vm.fd(), &cpu_template.vcpu_features)
+                .map_err(VmmError::VcpuInit)
+                .map_err(Internal)?;
+        }
+
         let mut regs = Aarch64RegisterVec::default();
         get_registers(&vcpus[0].kvm_vcpu.fd, &cpu_template.reg_list(), &mut regs)
             .map_err(GuestConfigError)?;

src/vmm/src/lib.rs

@@ -561,38 +561,6 @@ impl Vmm {
         Ok(vcpu_states)
     }
 
-    /// Restores vcpus kvm states.
-    pub fn restore_vcpu_states(
-        &mut self,
-        mut vcpu_states: Vec<VcpuState>,
-    ) -> Result<(), RestoreVcpusError> {
-        if vcpu_states.len() != self.vcpus_handles.len() {
-            return Err(RestoreVcpusError::InvalidInput);
-        }
-        for (handle, state) in self.vcpus_handles.iter().zip(vcpu_states.drain(..)) {
-            handle.send_event(VcpuEvent::RestoreState(Box::new(state)))?;
-        }
-
-        let vcpu_responses = self
-            .vcpus_handles
-            .iter()
-            // `Iterator::collect` can transform a `Vec<Result>` into a `Result<Vec>`.
-            .map(|handle| handle.response_receiver().recv_timeout(RECV_TIMEOUT_SEC))
-            .collect::<Result<Vec<VcpuResponse>, RecvTimeoutError>>()
-            .map_err(|_| RestoreVcpusError::UnexpectedVcpuResponse)?;
-
-        for response in vcpu_responses.into_iter() {
-            match response {
-                VcpuResponse::RestoredState => Ok(()),
-                VcpuResponse::Error(err) => Err(RestoreVcpusError::RestoreVcpuState(err)),
-                VcpuResponse::NotAllowed(reason) => Err(RestoreVcpusError::NotAllowed(reason)),
-                _ => Err(RestoreVcpusError::UnexpectedVcpuResponse),
-            }?;
-        }
-
-        Ok(())
-    }
-
     /// Dumps CPU configuration.
     pub fn dump_cpu_config(&mut self) -> Result<Vec<CpuConfiguration>, DumpCpuConfigError> {
         for handle in self.vcpus_handles.iter() {

src/vmm/src/persist.rs

@@ -370,6 +370,12 @@ pub fn additional_snapshot_version_check(
     {
         return Err(CreateSnapshotError::UnsupportedVersion);
     }
+
+    // We forbid snapshots older then 1.5 if any additional vcpu features are requested
+    #[cfg(target_arch = "aarch64")]
+    if microvm_state.vcpu_states[0].kvi.is_some() && version < FC_V1_5_SNAP_VERSION {
+        return Err(CreateSnapshotError::UnsupportedVersion);
+    }
     Ok(())
 }
 

src/vmm/src/vstate/vcpu/aarch64.rs

@@ -18,6 +18,7 @@ use crate::arch::aarch64::vcpu::{
     get_all_registers, get_all_registers_ids, get_mpidr, get_mpstate, get_registers, set_mpstate,
     set_registers, setup_boot_regs, VcpuError as ArchError,
 };
+use crate::cpu_config::aarch64::custom_cpu_template::VcpuFeatures;
 use crate::cpu_config::templates::CpuConfiguration;
 use crate::vcpu::{VcpuConfig, VcpuError};
 use crate::vstate::vcpu::VcpuEmulation;
@@ -40,6 +41,8 @@ pub enum KvmVcpuError {
     ApplyCpuTemplate(ArchError),
     #[error("Failed to restore the state of the vcpu: {0}")]
     RestoreState(ArchError),
+    #[error("Failed to restore the state of the vcpu: no kvi provided")]
+    MissingKvi,
     #[error("Failed to save the state of the vcpu: {0}")]
     SaveState(ArchError),
 }
@@ -53,6 +56,7 @@ pub struct KvmVcpu {
     pub fd: VcpuFd,
     pub mmio_bus: Option<crate::devices::Bus>,
     mpidr: u64,
+    kvi: Option<kvm_bindings::kvm_vcpu_init>,
 }
 
 impl KvmVcpu {
@@ -73,6 +77,7 @@ impl KvmVcpu {
             fd: kvm_vcpu,
             mmio_bus: None,
             mpidr: 0,
+            kvi: None,
         })
     }
 
@@ -118,20 +123,43 @@ impl KvmVcpu {
     /// # Arguments
     ///
     /// * `vm_fd` - The kvm `VmFd` for this microvm.
-    pub fn init(&self, vm_fd: &VmFd) -> Result<(), KvmVcpuError> {
-        let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
+    pub fn init(
+        &mut self,
+        vm_fd: &VmFd,
+        vcpu_features: &[VcpuFeatures],
+    ) -> Result<(), KvmVcpuError> {
+        let mut kvi = Self::default_kvi(vm_fd, self.index)?;
+
+        for feature in vcpu_features.iter() {
+            kvi.features[feature.index as usize] |= feature.bitmap;
+        }
 
+        self.kvi = if !vcpu_features.is_empty() {
+            Some(kvi)
+        } else {
+            None
+        };
+        self.fd.vcpu_init(&kvi).map_err(KvmVcpuError::Init)
+    }
+
+    pub fn default_kvi(
+        vm_fd: &VmFd,
+        index: u8,
+    ) -> Result<kvm_bindings::kvm_vcpu_init, KvmVcpuError> {
+        let mut kvi: kvm_bindings::kvm_vcpu_init = kvm_bindings::kvm_vcpu_init::default();
         // This reads back the kernel's preferred target type.
         vm_fd
             .get_preferred_target(&mut kvi)
             .map_err(KvmVcpuError::GetPreferredTarget)?;
         // We already checked that the capability is supported.
         kvi.features[0] |= 1 << kvm_bindings::KVM_ARM_VCPU_PSCI_0_2;
+
         // Non-boot cpus are powered off initially.
-        if self.index > 0 {
+        if index > 0 {
             kvi.features[0] |= 1 << kvm_bindings::KVM_ARM_VCPU_POWER_OFF;
         }
-        self.fd.vcpu_init(&kvi).map_err(KvmVcpuError::Init)
+
+        Ok(kvi)
     }
 
     /// Save the KVM internal state.
@@ -142,11 +170,23 @@ impl KvmVcpu {
         };
         get_all_registers(&self.fd, &mut state.regs).map_err(KvmVcpuError::SaveState)?;
         state.mpidr = get_mpidr(&self.fd).map_err(KvmVcpuError::SaveState)?;
+        state.kvi = self.kvi;
         Ok(state)
     }
 
     /// Use provided state to populate KVM internal state.
-    pub fn restore_state(&self, state: &VcpuState) -> Result<(), KvmVcpuError> {
+    pub fn restore_state(
+        &mut self,
+        vm_fd: &VmFd,
+        state: &VcpuState,
+        index: u8,
+    ) -> Result<(), KvmVcpuError> {
+        let kvi = match state.kvi {
+            Some(kvi) => kvi,
+            None => Self::default_kvi(vm_fd, index)?,
+        };
+        self.fd.vcpu_init(&kvi).map_err(KvmVcpuError::Init)?;
+        self.kvi = state.kvi;
         set_registers(&self.fd, &state.regs).map_err(KvmVcpuError::RestoreState)?;
         set_mpstate(&self.fd, state.mp_state).map_err(KvmVcpuError::RestoreState)?;
         Ok(())
@@ -192,13 +232,19 @@ pub struct VcpuState {
     // The VmState will give this away for saving restoring the icc and redistributor
     // registers.
     pub mpidr: u64,
+    #[version(start = 2, default_fn = "default_kvi")]
+    pub kvi: Option<kvm_bindings::kvm_vcpu_init>,
 }
 
 impl VcpuState {
     fn default_old_regs(_: u16) -> Vec<Aarch64RegisterOld> {
         Vec::default()
     }
 
+    fn default_kvi(_: u16) -> Option<kvm_bindings::kvm_vcpu_init> {
+        None
+    }
+
     fn de_regs(&mut self, _source_version: u16) -> VersionizeResult<()> {
         let mut regs = Aarch64RegisterVec::default();
         for reg in self.old_regs.iter() {
@@ -227,31 +273,23 @@ mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
     use std::os::unix::io::AsRawFd;
 
-    use kvm_bindings::KVM_REG_SIZE_U64;
     use utils::vm_memory::GuestMemoryMmap;
 
     use super::*;
-    use crate::arch::aarch64::regs::Aarch64RegisterRef;
     use crate::cpu_config::aarch64::CpuConfiguration;
     use crate::vcpu::VcpuConfig;
     use crate::vstate::vm::tests::setup_vm;
     use crate::vstate::vm::Vm;
 
     fn setup_vcpu(mem_size: usize) -> (Vm, KvmVcpu, GuestMemoryMmap) {
         let (mut vm, vm_mem) = setup_vm(mem_size);
-        let vcpu = KvmVcpu::new(0, &vm).unwrap();
-        vcpu.init(vm.fd()).unwrap();
+        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
+        vcpu.init(vm.fd(), &[]).unwrap();
         vm.setup_irqchip(1).unwrap();
 
         (vm, vcpu, vm_mem)
     }
 
-    fn init_vcpu(vcpu: &VcpuFd, vm: &VmFd) {
-        let mut kvi = kvm_bindings::kvm_vcpu_init::default();
-        vm.get_preferred_target(&mut kvi).unwrap();
-        vcpu.vcpu_init(&kvi).unwrap();
-    }
-
     #[test]
     fn test_create_vcpu() {
         let (vm, _) = setup_vm(0x1000);
@@ -302,9 +340,9 @@ mod tests {
 
     #[test]
     fn test_faulty_init_vcpu() {
-        let (vm, vcpu, _) = setup_vcpu(0x10000);
+        let (vm, mut vcpu, _) = setup_vcpu(0x10000);
         unsafe { libc::close(vm.fd().as_raw_fd()) };
-        let err = vcpu.init(vm.fd());
+        let err = vcpu.init(vm.fd(), &[]);
         assert!(err.is_err());
         assert_eq!(
             err.err().unwrap().to_string(),
@@ -315,7 +353,7 @@ mod tests {
     #[test]
     fn test_vcpu_save_restore_state() {
         let (mut vm, _vm_mem) = setup_vm(0x1000);
-        let vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
         vm.setup_irqchip(1).unwrap();
 
         // Calling KVM_GET_REGLIST before KVM_VCPU_INIT will result in error.
@@ -327,26 +365,14 @@ mod tests {
         ));
 
         // Try to restore the register using a faulty state.
-        let mut regs = Aarch64RegisterVec::default();
-        let mut reg = Aarch64RegisterRef::new(KVM_REG_SIZE_U64, &[0; 8]);
-        reg.id = 0;
-        regs.push(reg);
-        let faulty_vcpu_state = VcpuState {
-            regs,
-            ..Default::default()
-        };
-
-        let res = vcpu.restore_state(&faulty_vcpu_state);
+        let faulty_vcpu_state = VcpuState::default();
+        let res = vcpu.restore_state(vm.fd(), &faulty_vcpu_state, 0);
         assert!(res.is_err());
-        assert!(matches!(
-            res.unwrap_err(),
-            KvmVcpuError::RestoreState(ArchError::SetOneReg(0, _))
-        ));
 
-        init_vcpu(&vcpu.fd, vm.fd());
+        vcpu.init(vm.fd(), &[]).unwrap();
         let state = vcpu.save_state().expect("Cannot save state of vcpu");
         assert!(!state.regs.is_empty());
-        vcpu.restore_state(&state)
+        vcpu.restore_state(vm.fd(), &state, 0)
             .expect("Cannot restore state of vcpu");
     }
 
@@ -367,20 +393,20 @@ mod tests {
     fn test_dump_cpu_config_after_init() {
         // Test `dump_cpu_config()` after `KVM_VCPU_INIT`.
         let (mut vm, _vm_mem) = setup_vm(0x1000);
-        let vcpu = KvmVcpu::new(0, &vm).unwrap();
+        let mut vcpu = KvmVcpu::new(0, &vm).unwrap();
         vm.setup_irqchip(1).unwrap();
-        vcpu.init(vm.fd()).unwrap();
+        vcpu.init(vm.fd(), &[]).unwrap();
 
         assert!(vcpu.dump_cpu_config().is_ok());
     }
 
     #[test]
     fn test_setup_non_boot_vcpu() {
         let (vm, _) = setup_vm(0x1000);
-        let vcpu1 = KvmVcpu::new(0, &vm).unwrap();
-        assert!(vcpu1.init(vm.fd()).is_ok());
-        let vcpu2 = KvmVcpu::new(1, &vm).unwrap();
-        assert!(vcpu2.init(vm.fd()).is_ok());
+        let mut vcpu1 = KvmVcpu::new(0, &vm).unwrap();
+        assert!(vcpu1.init(vm.fd(), &[]).is_ok());
+        let mut vcpu2 = KvmVcpu::new(1, &vm).unwrap();
+        assert!(vcpu2.init(vm.fd(), &[]).is_ok());
     }
 
     #[test]