refactor: x64: use kvm_bindings::Msrs type directly in CpuConfiguration

Similarly to the previous commit, replace Btree type for Msrs handling
with kvm_bindings::Msrs to remove the unnecessary state transitions and
memory allocations. Uniqueness of entries is preserved with
`msrs_insert` function.

Together with previous commit this reduces the binary sizy by ~30KiB and
reduces the maximum RSS during `test_all_vcpus_online` test by 50-80 KiB

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

Author: ShadowCurse

src/cpu-template-helper/Cargo.toml

@@ -15,6 +15,7 @@ tracing = ["log-instrument", "vmm/tracing"]
 [dependencies]
 clap = { version = "4.6.1", features = ["derive", "string"] }
 displaydoc = "0.2.6"
+kvm-bindings = "0.14.0"
 libc = "0.2.186"
 log-instrument = { path = "../log-instrument", optional = true }
 serde = { version = "1.0.228", features = ["derive"] }
@@ -26,6 +27,3 @@ vmm-sys-util = "0.15.0"
 
 [lints]
 workspace = true
-
-[dev-dependencies]
-kvm-bindings = "0.14.0"

src/cpu-template-helper/src/template/dump/x86_64.rs

@@ -1,8 +1,7 @@
 // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use std::collections::BTreeMap;
-
+use kvm_bindings::Msrs;
 use vmm::MSR_RANGE;
 use vmm::arch::x86_64::generated::msr_index::*;
 use vmm::arch::x86_64::msr::MsrRange;
@@ -47,10 +46,11 @@ fn cpuid_to_modifiers(cpuid: &Cpuid) -> Vec<CpuidLeafModifier> {
     result
 }
 
-fn msrs_to_modifier(msrs: &BTreeMap<u32, u64>) -> Vec<RegisterModifier> {
+fn msrs_to_modifier(msrs: &Msrs) -> Vec<RegisterModifier> {
     let mut msrs: Vec<RegisterModifier> = msrs
+        .as_slice()
         .iter()
-        .map(|(index, value)| msr_modifier!(*index, *value))
+        .map(|entry| msr_modifier!(entry.index, entry.data))
         .collect();
 
     msrs.retain(|modifier| !should_exclude_msr(modifier.addr));
@@ -122,8 +122,7 @@ fn should_exclude_msr_amd(index: u32) -> bool {
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeMap;
-
+    use kvm_bindings::kvm_msr_entry;
     use vmm::cpu_config::x86_64::cpuid::IntelCpuid;
 
     use super::*;
@@ -183,24 +182,29 @@ mod tests {
         ]
     }
 
-    fn build_sample_msrs() -> BTreeMap<u32, u64> {
-        let mut map = BTreeMap::from([
+    fn build_sample_msrs() -> Msrs {
+        let entry = |index, data| kvm_msr_entry {
+            index,
+            data,
+            ..Default::default()
+        };
+        let mut entries = vec![
             // should be sorted in the result.
-            (0x1, 0xffff_ffff_ffff_ffff),
-            (0x5, 0xffff_ffff_0000_0000),
-            (0x3, 0x0000_0000_ffff_ffff),
-            (0x2, 0x0000_0000_0000_0000),
-        ]);
+            entry(0x1, 0xffff_ffff_ffff_ffff),
+            entry(0x5, 0xffff_ffff_0000_0000),
+            entry(0x3, 0x0000_0000_ffff_ffff),
+            entry(0x2, 0x0000_0000_0000_0000),
+        ];
         // should be excluded from the result.
         MSR_EXCLUSION_LIST
             .iter()
             .chain(MSR_EXCLUSION_LIST_AMD.iter())
             .for_each(|range| {
                 (range.base..(range.base + range.nmsrs)).for_each(|id| {
-                    map.insert(id, 0);
+                    entries.push(entry(id, 0));
                 })
             });
-        map
+        Msrs::from_entries(&entries).unwrap()
     }
 
     fn build_expected_msr_modifiers() -> Vec<RegisterModifier> {

src/vmm/src/arch/x86_64/msr.rs

@@ -426,21 +426,43 @@ pub fn create_boot_msr_entries() -> Vec<kvm_msr_entry> {
     ]
 }
 
+/// Insert or update an MSR entry in `msrs`, mirroring `BTreeMap::insert` semantics: when an entry
+/// with the given index already exists its `data` is overwritten, otherwise a new entry is pushed.
+///
+/// # Errors
+///
+/// Propagates [`vmm_sys_util::fam::Error`] from [`Msrs::push`] when the FAM struct is at capacity.
+pub fn msrs_insert(msrs: &mut Msrs, index: u32, data: u64) {
+    if let Some(entry) = msrs
+        .as_mut_slice()
+        .iter_mut()
+        .find(|entry| entry.index == index)
+    {
+        entry.data = data;
+    } else {
+        msrs.push(kvm_msr_entry {
+            index,
+            data,
+            ..Default::default()
+        })
+        .expect("MSRS entry count exceeded KVM_MAX_MSR_ENTRIES");
+    }
+}
+
 /// Configure Model Specific Registers (MSRs) required to boot Linux for a given x86_64 vCPU.
 ///
 /// # Arguments
 ///
 /// * `vcpu` - Structure for the VCPU that holds the VCPU's fd.
+/// * `msrs` - The MSRs to write to KVM.
 ///
 /// # Errors
 ///
 /// When:
-/// - Failed to create [`vmm_sys_util::fam::FamStructWrapper`] for MSRs.
 /// - [`kvm_ioctls::ioctls::vcpu::VcpuFd::set_msrs`] errors.
 /// - [`kvm_ioctls::ioctls::vcpu::VcpuFd::set_msrs`] fails to write all given MSRs entries.
-pub fn set_msrs(vcpu: &VcpuFd, msr_entries: &[kvm_msr_entry]) -> Result<(), MsrError> {
-    let msrs = Msrs::from_entries(msr_entries)?;
-    vcpu.set_msrs(&msrs)
+pub fn set_msrs(vcpu: &VcpuFd, msrs: &Msrs) -> Result<(), MsrError> {
+    vcpu.set_msrs(msrs)
         .map_err(MsrError::SetMsrs)
         .and_then(|msrs_written| {
             if msrs_written == msrs.as_fam_struct_ref().nmsrs as usize {
@@ -487,7 +509,8 @@ mod tests {
     fn test_setup_msrs() {
         let vcpu = create_vcpu();
         let msr_boot_entries = create_boot_msr_entries();
-        set_msrs(&vcpu, &msr_boot_entries).unwrap();
+        let msrs = Msrs::from_entries(&msr_boot_entries).unwrap();
+        set_msrs(&vcpu, &msrs).unwrap();
 
         // This test will check against the last MSR entry configured (the tenth one).
         // See create_msr_entries() for details.
@@ -515,12 +538,31 @@ mod tests {
         // Test `set_msrs()` with a valid MSR entry. It should succeed, as IA32_TSC MSR is listed
         // in supported MSRs as of now.
         let vcpu = create_vcpu();
-        let msr_entries = vec![kvm_msr_entry {
+        let msrs = Msrs::from_entries(&[kvm_msr_entry {
             index: MSR_IA32_TSC,
             data: 0,
             ..Default::default()
-        }];
-        set_msrs(&vcpu, &msr_entries).unwrap();
+        }])
+        .unwrap();
+        set_msrs(&vcpu, &msrs).unwrap();
+    }
+
+    #[test]
+    fn test_msrs_insert_overwrites_and_pushes() {
+        let mut msrs = Msrs::new(0).unwrap();
+
+        msrs_insert(&mut msrs, 0x10, 0xaa);
+        msrs_insert(&mut msrs, 0x20, 0xbb);
+        assert_eq!(msrs.as_slice().len(), 2);
+        assert_eq!(msrs.as_slice()[0].data, 0xaa);
+        assert_eq!(msrs.as_slice()[1].data, 0xbb);
+
+        msrs_insert(&mut msrs, 0x10, 0xcc);
+        assert_eq!(msrs.as_slice().len(), 2);
+        assert_eq!(msrs.as_slice()[0].index, 0x10);
+        assert_eq!(msrs.as_slice()[0].data, 0xcc);
+        assert_eq!(msrs.as_slice()[1].index, 0x20);
+        assert_eq!(msrs.as_slice()[1].data, 0xbb);
     }
 
     #[test]
@@ -529,12 +571,13 @@ mod tests {
         // listed in supported MSRs as of now. If hardware vendor adds this MSR index and KVM
         // supports this MSR, we need to change the index as needed.
         let vcpu = create_vcpu();
-        let msr_entries = vec![kvm_msr_entry {
+        let msrs = Msrs::from_entries(&[kvm_msr_entry {
             index: 2,
             ..Default::default()
-        }];
+        }])
+        .unwrap();
         assert_eq!(
-            set_msrs(&vcpu, &msr_entries).unwrap_err(),
+            set_msrs(&vcpu, &msrs).unwrap_err(),
             MsrError::SetMsrsIncomplete
         );
     }

src/vmm/src/arch/x86_64/vcpu.rs

@@ -5,7 +5,6 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the THIRD-PARTY file.
 
-use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::sync::Arc;
 
@@ -20,7 +19,7 @@ use vmm_sys_util::fam::{self, FamStruct};
 use crate::arch::EntryPoint;
 use crate::arch::x86_64::generated::msr_index::{MSR_IA32_TSC, MSR_IA32_TSC_DEADLINE};
 use crate::arch::x86_64::interrupts;
-use crate::arch::x86_64::msr::{MsrError, create_boot_msr_entries};
+use crate::arch::x86_64::msr::{MsrError, create_boot_msr_entries, msrs_insert};
 use crate::arch::x86_64::regs::{SetupFpuError, SetupRegistersError, SetupSpecialRegistersError};
 use crate::cpu_config::x86_64::{CpuConfiguration, cpuid};
 use crate::logger::{IncMetric, METRICS, error, warn};
@@ -222,12 +221,13 @@ impl KvmVcpu {
 
         // Clone MSR entries that are modified by CPU template from `VcpuConfig`.
         let mut msrs = vcpu_config.cpu_config.msrs.clone();
-        self.msrs_to_save.extend(msrs.keys());
+        self.msrs_to_save
+            .extend(msrs.as_slice().iter().map(|entry| entry.index));
 
         // Apply MSR modification to comply the linux boot protocol.
-        create_boot_msr_entries().into_iter().for_each(|entry| {
-            msrs.insert(entry.index, entry.data);
-        });
+        for entry in create_boot_msr_entries() {
+            msrs_insert(&mut msrs, entry.index, entry.data);
+        }
 
         // TODO - Add/amend MSRs for vCPUs based on cpu_config
         // By this point the Guest CPUID is established. Some CPU features require MSRs
@@ -247,16 +247,7 @@ impl KvmVcpu {
         // save is `architectural MSRs` + `MSRs inferred through CPUID` + `other
         // MSRs defined by the template`
 
-        let kvm_msrs = msrs
-            .into_iter()
-            .map(|entry| kvm_bindings::kvm_msr_entry {
-                index: entry.0,
-                data: entry.1,
-                ..Default::default()
-            })
-            .collect::<Vec<_>>();
-
-        crate::arch::x86_64::msr::set_msrs(&self.fd, &kvm_msrs)?;
+        crate::arch::x86_64::msr::set_msrs(&self.fd, &msrs)?;
         crate::arch::x86_64::regs::setup_regs(&self.fd, kernel_entry_point)?;
         crate::arch::x86_64::regs::setup_fpu(&self.fd)?;
         crate::arch::x86_64::regs::setup_sregs(guest_mem, &self.fd, kernel_entry_point.protocol)?;
@@ -534,19 +525,18 @@ impl KvmVcpu {
     /// # Errors
     ///
     /// * When `KvmVcpu::get_msr_chunks()` returns errors.
+    /// * When [`kvm_bindings::Msrs::new`] returns errors.
     pub fn get_msrs(
         &self,
         msr_index_iter: impl ExactSizeIterator<Item = u32>,
-    ) -> Result<BTreeMap<u32, u64>, KvmVcpuError> {
-        let mut msrs = BTreeMap::new();
-        self.get_msr_chunks(msr_index_iter)?
-            .iter()
-            .for_each(|msr_chunk| {
-                msr_chunk.as_slice().iter().for_each(|msr| {
-                    msrs.insert(msr.index, msr.data);
-                });
-            });
-        Ok(msrs)
+    ) -> Result<Msrs, KvmVcpuError> {
+        let mut combined = Msrs::new(0).map_err(KvmVcpuError::Fam)?;
+        for chunk in self.get_msr_chunks(msr_index_iter)? {
+            for entry in chunk.as_slice() {
+                msrs_insert(&mut combined, entry.index, entry.data);
+            }
+        }
+        Ok(combined)
     }
 
     /// Save the KVM internal state.
@@ -1012,7 +1002,7 @@ mod tests {
             smt: false,
             cpu_config: CpuConfiguration {
                 cpuid: Cpuid::try_from(vm.kvm().supported_cpuid.clone()).unwrap(),
-                msrs: BTreeMap::new(),
+                msrs: Msrs::new(0).unwrap(),
             },
         };
         vcpu.configure(
@@ -1080,7 +1070,7 @@ mod tests {
             smt: false,
             cpu_config: CpuConfiguration {
                 cpuid: Cpuid::try_from(vm.kvm().supported_cpuid.clone()).unwrap(),
-                msrs: BTreeMap::new(),
+                msrs: Msrs::new(0).unwrap(),
             },
         };