Skip to content

Commit

Permalink
refactor: simplify get_msr_chunks
Browse files Browse the repository at this point in the history
Move logic to create single msr chunk into separate function.
Also use iterators for msr indexes to remove some unneeded
allocations.

Signed-off-by: Egor Lazarchuk <[email protected]>
  • Loading branch information
ShadowCurse committed Jun 20, 2024
1 parent fc7c896 commit 0524244
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 40 deletions.
3 changes: 1 addition & 2 deletions src/vmm/src/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -770,10 +770,9 @@ pub fn configure_system_for_boot(
use crate::cpu_config::x86_64::cpuid;
let cpuid = cpuid::Cpuid::try_from(vmm.vm.supported_cpuid().clone())
.map_err(GuestConfigError::CpuidFromKvmCpuid)?;
let msr_index_list = cpu_template.get_msr_index_list();
let msrs = vcpus[0]
.kvm_vcpu
.get_msrs(&msr_index_list)
.get_msrs(cpu_template.msr_index_iter())
.map_err(GuestConfigError::VcpuIoctl)?;
CpuConfiguration { cpuid, msrs }
};
Expand Down
9 changes: 3 additions & 6 deletions src/vmm/src/cpu_config/x86_64/custom_cpu_template.rs
Original file line number Diff line number Diff line change
Expand Up @@ -144,12 +144,9 @@ pub struct CustomCpuTemplate {
}

impl CustomCpuTemplate {
/// Get a list of MSR indices that are modified by the CPU template.
pub fn get_msr_index_list(&self) -> Vec<u32> {
self.msr_modifiers
.iter()
.map(|modifier| modifier.addr)
.collect()
/// Get an iterator of MSR indices that are modified by the CPU template.
pub fn msr_index_iter(&self) -> impl ExactSizeIterator<Item = u32> + '_ {
self.msr_modifiers.iter().map(|modifier| modifier.addr)
}

/// Validate the correctness of the template.
Expand Down
99 changes: 67 additions & 32 deletions src/vmm/src/vstate/vcpu/x86_64.rs
Original file line number Diff line number Diff line change
Expand Up @@ -324,42 +324,76 @@ impl KvmVcpu {
///
/// # Arguments
///
/// * `msr_index_list`: List of MSR indices.
/// * `msr_index_iter`: Iterator over MSR indices.
///
/// # Errors
///
/// * When [`kvm_bindings::Msrs::new`] returns errors.
/// * When [`kvm_ioctls::VcpuFd::get_msrs`] returns errors.
/// * When the return value of [`kvm_ioctls::VcpuFd::get_msrs`] (the number of entries that
/// could be gotten) is less than expected.
fn get_msr_chunks(&self, msr_index_list: &[u32]) -> Result<Vec<Msrs>, KvmVcpuError> {
let mut msr_chunks: Vec<Msrs> = Vec::new();

for msr_index_chunk in msr_index_list.chunks(KVM_MAX_MSR_ENTRIES) {
let mut msrs = Msrs::new(msr_index_chunk.len())?;
let msr_entries = msrs.as_mut_slice();
assert_eq!(msr_index_chunk.len(), msr_entries.len());
for (pos, index) in msr_index_chunk.iter().enumerate() {
msr_entries[pos].index = *index;
}

let expected_nmsrs = msrs.as_slice().len();
let nmsrs = self
.fd
.get_msrs(&mut msrs)
.map_err(KvmVcpuError::VcpuGetMsrs)?;
if nmsrs != expected_nmsrs {
return Err(KvmVcpuError::VcpuGetMsr(msr_index_chunk[nmsrs]));
}

msr_chunks.push(msrs);
fn get_msr_chunks(
&self,
mut msr_index_iter: impl ExactSizeIterator<Item = u32>,
) -> Result<Vec<Msrs>, KvmVcpuError> {
let num_chunks = msr_index_iter.len() / KVM_MAX_MSR_ENTRIES;
let excess = msr_index_iter.len() % KVM_MAX_MSR_ENTRIES;

let mut msr_chunks: Vec<Msrs> = Vec::with_capacity(num_chunks);

for _ in 0..num_chunks {
let chunk = self.get_msr_chunk(&mut msr_index_iter, KVM_MAX_MSR_ENTRIES)?;
msr_chunks.push(chunk);
}
if excess != 0 {
let chunk = self.get_msr_chunk(&mut msr_index_iter, excess)?;
msr_chunks.push(chunk);
}

Self::fix_zero_tsc_deadline_msr(&mut msr_chunks);

Ok(msr_chunks)
}

/// Get single MSR chunk for the given MSR index iterator with
/// specified maximum length.
///
/// # Arguments
///
/// * `msr_index_iter`: Iterator over MSR indices.
/// * `max_chunk_len`: Maximum lenght of a chunk to get.
///
/// # Errors
///
/// * When [`kvm_bindings::Msrs::new`] returns errors.
/// * When [`kvm_ioctls::VcpuFd::get_msrs`] returns errors.
/// * When the return value of [`kvm_ioctls::VcpuFd::get_msrs`] (the number of entries that
/// could be gotten) is less than expected.
pub fn get_msr_chunk(
&self,
msr_index_iter: impl ExactSizeIterator<Item = u32>,
max_chunk_len: usize,
) -> Result<Msrs, KvmVcpuError> {
let chunk_size = msr_index_iter.len().min(max_chunk_len);
let chunk_iter = msr_index_iter.take(chunk_size);

let mut msrs = Msrs::new(chunk_size)?;
let msr_entries = msrs.as_mut_slice();
for (pos, msr_index) in chunk_iter.enumerate() {
msr_entries[pos].index = msr_index;
}

let nmsrs = self
.fd
.get_msrs(&mut msrs)
.map_err(KvmVcpuError::VcpuGetMsrs)?;
if nmsrs != chunk_size {
Err(KvmVcpuError::VcpuGetMsr(msrs.as_slice()[nmsrs].index))
} else {
Ok(msrs)
}
}

/// Get MSRs for the given MSR index list.
///
/// # Arguments
Expand All @@ -369,9 +403,12 @@ impl KvmVcpu {
/// # Errors
///
/// * When `KvmVcpu::get_msr_chunks()` returns errors.
pub fn get_msrs(&self, msr_index_list: &[u32]) -> Result<HashMap<u32, u64>, KvmVcpuError> {
pub fn get_msrs(
&self,
msr_index_iter: impl ExactSizeIterator<Item = u32>,
) -> Result<HashMap<u32, u64>, KvmVcpuError> {
let mut msrs: HashMap<u32, u64> = HashMap::new();
self.get_msr_chunks(msr_index_list)?
self.get_msr_chunks(msr_index_iter)?
.iter()
.for_each(|msr_chunk| {
msr_chunk.as_slice().iter().for_each(|msr| {
Expand Down Expand Up @@ -422,8 +459,7 @@ impl KvmVcpu {
None
});
let cpuid = self.get_cpuid()?;
let saved_msrs =
self.get_msr_chunks(&self.msrs_to_save.iter().copied().collect::<Vec<_>>())?;
let saved_msrs = self.get_msr_chunks(self.msrs_to_save.iter().copied())?;
let vcpu_events = self
.fd
.get_vcpu_events()
Expand Down Expand Up @@ -452,7 +488,7 @@ impl KvmVcpu {
let cpuid = cpuid::Cpuid::try_from(self.get_cpuid()?)?;
let kvm = kvm_ioctls::Kvm::new().unwrap();
let msr_index_list = crate::arch::x86_64::msr::get_msrs_to_dump(&kvm)?;
let msrs = self.get_msrs(msr_index_list.as_slice())?;
let msrs = self.get_msrs(msr_index_list.as_slice().iter().copied())?;
Ok(CpuConfiguration { cpuid, msrs })
}

Expand Down Expand Up @@ -687,7 +723,7 @@ mod tests {
let cpuid = Cpuid::try_from(vm.supported_cpuid().clone())
.map_err(GuestConfigError::CpuidFromKvmCpuid)?;
let msrs = vcpu
.get_msrs(&template.get_msr_index_list())
.get_msrs(template.msr_index_iter())
.map_err(GuestConfigError::VcpuIoctl)?;
let base_cpu_config = CpuConfiguration { cpuid, msrs };
let cpu_config = CpuConfiguration::apply_template(base_cpu_config, template)?;
Expand Down Expand Up @@ -956,8 +992,7 @@ mod tests {
// Test `get_msrs()` with the MSR indices that should be serialized into snapshots.
// The MSR indices should be valid and this test should succeed.
let (_, vcpu, _) = setup_vcpu(0x1000);
vcpu.get_msrs(&vcpu.msrs_to_save.iter().copied().collect::<Vec<_>>())
.unwrap();
vcpu.get_msrs(vcpu.msrs_to_save.iter()).unwrap();
}

#[test]
Expand All @@ -968,7 +1003,7 @@ mod tests {

let kvm = kvm_ioctls::Kvm::new().unwrap();
let msrs_to_dump = crate::arch::x86_64::msr::get_msrs_to_dump(&kvm).unwrap();
vcpu.get_msrs(msrs_to_dump.as_slice()).unwrap();
vcpu.get_msrs(msrs_to_dump.as_slice().iter()).unwrap();
}

#[test]
Expand All @@ -978,7 +1013,7 @@ mod tests {
// Currently, MSR indices 2..=4 are not listed as supported MSRs.
let (_, vcpu, _) = setup_vcpu(0x1000);
let msr_index_list: Vec<u32> = vec![2, 3, 4];
match vcpu.get_msrs(&msr_index_list) {
match vcpu.get_msrs(msr_index_list.iter()) {
Err(KvmVcpuError::VcpuGetMsr(_)) => (),
Err(err) => panic!("Unexpected error: {err}"),
Ok(_) => {
Expand Down

0 comments on commit 0524244

Please sign in to comment.