snap/restore: extend the number of saved MSRs

KVM can get/set only a maximum of KVM_MAX_MSR_ENTRIES MSRs at a time
through ioctl.
In some scenarios we need to save more than KVM_MAX_MSR_ENTRIES in the
snapshot when the Guest CPU is emulating a lot of MSRs.

Change the list of supported MSRs to a HashSet. This allows better
flexibility in adding and comparing MSR address lists.
Also save a `Vec<Msrs>` instead of a single `Msrs` struct to allow
saving an unlimited number of MSRs in the snapshot.

Signed-off-by: Alexandru-Cezar Sardan <[email protected]>

Author: Alexandru-Cezar Sardan

src/vmm/src/version_map.rs

@@ -53,6 +53,8 @@ lazy_static! {
 
         // v1.2 state change mappings.
         version_map.new_version().set_type_version(VmInfo::type_id(), 2);
+        #[cfg(target_arch = "x86_64")]
+        version_map.set_type_version(VcpuState::type_id(), 3);
 
         version_map
     };

src/vmm/src/vstate/vcpu/x86_64.rs

@@ -5,6 +5,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the THIRD-PARTY file.
 
+use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::{fmt, result};
 
@@ -14,11 +15,11 @@ use arch::x86_64::regs::{SetupFpuError, SetupRegistersError, SetupSpecialRegiste
 use cpuid::{c3, filter_cpuid, msrs_to_save_by_cpuid, t2, t2s, VmSpec};
 use kvm_bindings::{
     kvm_debugregs, kvm_lapic_state, kvm_mp_state, kvm_regs, kvm_sregs, kvm_vcpu_events, kvm_xcrs,
-    kvm_xsave, CpuId, MsrList, Msrs,
+    kvm_xsave, CpuId, Msrs, KVM_MAX_MSR_ENTRIES,
 };
 use kvm_ioctls::{VcpuExit, VcpuFd};
 use logger::{error, warn, IncMetric, METRICS};
-use versionize::{VersionMap, Versionize, VersionizeResult};
+use versionize::{VersionMap, Versionize, VersionizeError, VersionizeResult};
 use versionize_derive::Versionize;
 use vm_memory::{Address, GuestAddress, GuestMemoryMmap};
 
@@ -216,7 +217,7 @@ pub struct KvmVcpu {
     pub pio_bus: Option<devices::Bus>,
     pub mmio_bus: Option<devices::Bus>,
 
-    msr_list: MsrList,
+    msr_list: HashSet<u32>,
 }
 
 impl KvmVcpu {
@@ -234,7 +235,7 @@ impl KvmVcpu {
             fd: kvm_vcpu,
             pio_bus: None,
             mmio_bus: None,
-            msr_list: vm.supported_msrs().clone(),
+            msr_list: vm.supported_msrs().as_slice().iter().copied().collect(),
         })
     }
 
@@ -294,11 +295,7 @@ impl KvmVcpu {
         // the extra MSRs that we need to save based on a dependency map.
         let extra_msrs =
             msrs_to_save_by_cpuid(&cpuid).map_err(KvmVcpuConfigureError::FilterCpuid)?;
-        for msr in extra_msrs {
-            self.msr_list
-                .push(msr)
-                .map_err(KvmVcpuConfigureError::PushMsrEntries)?;
-        }
+        self.msr_list.extend(extra_msrs);
 
         // TODO: Some MSRs depend on values of other MSRs. This dependency will need to
         // be implemented. For now we define known dependencies statically in the CPU
@@ -310,11 +307,7 @@ impl KvmVcpu {
         // to save at snapshot as well.
         // C3 and T2 currently don't have extra MSRs to save/set
         if vcpu_config.cpu_template == CpuFeaturesTemplate::T2S {
-            for msr in t2s::msr_entries_to_save() {
-                self.msr_list
-                    .push(*msr)
-                    .map_err(KvmVcpuConfigureError::PushMsrEntries)?;
-            }
+            self.msr_list.extend(t2s::msr_entries_to_save());
             t2s::update_msr_entries(&mut msr_boot_entries);
         }
         // By this point we know that at snapshot, the list of MSRs we need to
@@ -367,17 +360,24 @@ impl KvmVcpu {
         // GET_MSRS requires a pre-populated data structure to do something
         // meaningful. For SET_MSRS it will then contain good data.
 
-        // Build the list of MSRs we want to save.
-        let num_msrs = self.msr_list.as_fam_struct_ref().nmsrs as usize;
-        let mut msrs = Msrs::new(num_msrs).map_err(Error::Fam)?;
-        {
-            let indices = self.msr_list.as_slice();
+        // Build the list of MSRs we want to save. Sometimes we need to save
+        // more than KVM_MAX_MSR_ENTRIES in the snapshot, so we use a Vec<Msrs>
+        // to allow an unlimited number.
+        let mut all_msrs: Vec<Msrs> = Vec::new();
+        let msr_list: Vec<&u32> = self.msr_list.iter().collect();
+
+        // KVM only supports getting KVM_MAX_MSR_ENTRIES at a time so chunk
+        // them up into `Msrs` so it's easy to pass to the ioctl.
+        for chunk in msr_list.chunks(KVM_MAX_MSR_ENTRIES) {
+            let mut msrs = Msrs::new(chunk.len()).map_err(Error::Fam)?;
             let msr_entries = msrs.as_mut_slice();
-            assert_eq!(indices.len(), msr_entries.len());
-            for (pos, index) in indices.iter().enumerate() {
-                msr_entries[pos].index = *index;
+            assert_eq!(chunk.len(), msr_entries.len());
+            for (pos, index) in chunk.iter().enumerate() {
+                msr_entries[pos].index = **index;
             }
+            all_msrs.push(msrs);
         }
+
         let mp_state = self.fd.get_mp_state().map_err(Error::VcpuGetMpState)?;
         let regs = self.fd.get_regs().map_err(Error::VcpuGetRegs)?;
         let sregs = self.fd.get_sregs().map_err(Error::VcpuGetSregs)?;
@@ -392,9 +392,12 @@ impl KvmVcpu {
             warn!("TSC freq not available. Snapshot cannot be loaded on a different CPU model.");
             None
         });
-        let nmsrs = self.fd.get_msrs(&mut msrs).map_err(Error::VcpuGetMsrs)?;
-        if nmsrs != num_msrs {
-            return Err(Error::VcpuGetMSRSIncomplete);
+        for msrs in all_msrs.iter_mut() {
+            let expected_nmsrs = msrs.as_slice().len();
+            let nmsrs = self.fd.get_msrs(msrs).map_err(Error::VcpuGetMsrs)?;
+            if nmsrs != expected_nmsrs {
+                return Err(Error::VcpuGetMSRSIncomplete);
+            }
         }
         let vcpu_events = self
             .fd
@@ -406,7 +409,8 @@ impl KvmVcpu {
                 .fd
                 .get_cpuid2(kvm_bindings::KVM_MAX_CPUID_ENTRIES)
                 .map_err(Error::VcpuGetCpuid)?,
-            msrs,
+            saved_msrs: all_msrs,
+            msrs: Msrs::new(0).map_err(Error::Fam)?,
             debug_regs,
             lapic,
             mp_state,
@@ -486,10 +490,11 @@ impl KvmVcpu {
         self.fd
             .set_lapic(&state.lapic)
             .map_err(Error::VcpuSetLapic)?;
-        let nmsrs = self.fd.set_msrs(&state.msrs).map_err(Error::VcpuSetMsrs)?;
-        let num_msrs = state.msrs.as_fam_struct_ref().nmsrs as usize;
-        if nmsrs < num_msrs {
-            return Err(Error::VcpuSetMSRSIncomplete);
+        for msrs in &state.saved_msrs {
+            let nmsrs = self.fd.set_msrs(msrs).map_err(Error::VcpuSetMsrs)?;
+            if nmsrs < msrs.as_fam_struct_ref().nmsrs as usize {
+                return Err(Error::VcpuSetMSRSIncomplete);
+            }
         }
         self.fd
             .set_vcpu_events(&state.vcpu_events)
@@ -535,7 +540,10 @@ impl KvmVcpu {
 // NOTICE: Any changes to this structure require a snapshot version bump.
 pub struct VcpuState {
     pub cpuid: CpuId,
+    #[version(end = 3, default_fn = "default_msrs")]
     msrs: Msrs,
+    #[version(start = 3, de_fn = "de_saved_msrs", ser_fn = "ser_saved_msrs")]
+    saved_msrs: Vec<Msrs>,
     debug_regs: kvm_debugregs,
     lapic: kvm_lapic_state,
     mp_state: kvm_mp_state,
@@ -565,6 +573,43 @@ impl VcpuState {
 
         Ok(())
     }
+
+    fn default_msrs(_source_version: u16) -> Msrs {
+        // Safe to unwrap since Msrs::new() only returns an error if the number
+        // of elements exceeds KVM_MAX_MSR_ENTRIES
+        Msrs::new(0).unwrap()
+    }
+
+    fn de_saved_msrs(&mut self, source_version: u16) -> VersionizeResult<()> {
+        if source_version < 3 {
+            self.saved_msrs.push(self.msrs.clone());
+        }
+        Ok(())
+    }
+
+    fn ser_saved_msrs(&mut self, target_version: u16) -> VersionizeResult<()> {
+        match self.saved_msrs.len() {
+            0 => Err(VersionizeError::Serialize(
+                "Cannot serialize MSRs because the MSR list is empty".to_string(),
+            )),
+            1 => {
+                if target_version < 3 {
+                    self.msrs = self.saved_msrs[0].clone();
+                    Ok(())
+                } else {
+                    Err(VersionizeError::Serialize(format!(
+                        "Cannot serialize MSRs to target version {}",
+                        target_version
+                    )))
+                }
+            }
+            _ => Err(VersionizeError::Serialize(
+                "Cannot serialize MSRs. The uVM state needs to save
+                 more MSRs than the target snapshot version supports."
+                    .to_string(),
+            )),
+        }
+    }
 }
 
 #[cfg(test)]
@@ -585,6 +630,7 @@ mod tests {
             VcpuState {
                 cpuid: CpuId::new(1).unwrap(),
                 msrs: Msrs::new(1).unwrap(),
+                saved_msrs: vec![Msrs::new(1).unwrap()],
                 debug_regs: Default::default(),
                 lapic: Default::default(),
                 mp_state: Default::default(),