Repository: firecracker-microvm/firecracker
File: src/vmm/src/vstate/vcpu/x86_64.rs (Rust)

Code before:
        //
        // KVM_GET_LAPIC may change state of LAPIC before returning it.
        //
        // GET_VCPU_EVENTS should probably be last to save. The code looks as
        // it might as well be affected by internal state modifications of the
        // GET ioctls.
        //
        // SREGS saves/restores a pending interrupt, similar to what
        // VCPU_EVENTS also does.

        let mp_state = self
            .fd
            .get_mp_state()
            .map_err(KvmVcpuError::VcpuGetMpState)?;
        let regs = self.fd.get_regs().map_err(KvmVcpuError::VcpuGetRegs)?;
        let sregs = self.fd.get_sregs().map_err(KvmVcpuError::VcpuGetSregs)?;
        let xsave = match self.xsave2_size {
            Some(size) => {
                // SAFETY: The subtraction never underflows because
                // `KVM_CHECK_EXTENSION(KVM_CAP_XSAVE2)` returns at least 4096 bytes which equals
                // the size of `kvm_xsave`.
                // https://docs.kernel.org/virt/kvm/api.html#kvm-get-xsave2
                //
                // Note that the FAM entry type is `__u32` although `xsave2_xize` is in bytes.
                let fam_len = (size - std::mem::size_of::<kvm_xsave>())
                    .div_ceil(std::mem::size_of::<<kvm_xsave2 as FamStruct>::Entry>());
                let mut xsave = Xsave::new(fam_len).map_err(KvmVcpuError::Fam)?;
                // SAFETY: Safe because `xsave` is allocated above with enough size to save XSTATE.
                unsafe { self.fd.get_xsave2(&mut xsave) }.map_err(KvmVcpuError::VcpuGetXsave2)?;
                xsave
            }
            // `KVM_CAP_XSAVE2` not supported
            None => Xsave::from(vec![kvm_xsave2 {
                len: 0,
                xsave: self.fd.get_xsave().map_err(KvmVcpuError::VcpuGetXsave)?,
            }]),
        };
        let xcrs = self.fd.get_xcrs().map_err(KvmVcpuError::VcpuGetXcrs)?;
        let debug_regs = self
            .fd
            .get_debug_regs()
            .map_err(KvmVcpuError::VcpuGetDebugRegs)?;
        let lapic = self.fd.get_lapic().map_err(KvmVcpuError::VcpuGetLapic)?;
        let tsc_khz = self.get_tsc_khz().ok().or_else(|| {
            // v0.25 and newer snapshots without TSC will only work on
            // the same CPU model as the host on which they were taken.
            // TODO: Add negative test for this warning failure.
            warn!("TSC freq not available. Snapshot cannot be loaded on a different CPU model.");
            None
        });
        let cpuid = self.get_cpuid()?;

Review comment (suggestion):
I'm wondering if we can avoid this "magic" or at least hide it inside the `Xsave` struct. It's so confusing at the moment that you have to compute "how many members of size A are exceeding the size B". I had my head spin for a good 5m on this.

Code after:
        msr_index_iter: impl Iterator<Item = u32>,
        chunk_size: usize,
    ) -> Result<Msrs, KvmVcpuError> {
        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)?;
        // GET_MSRS returns a number of successfully set msrs.
        // If number of set msrs is not equal to the length of
        // `msrs`, then the value retuned by GET_MSRS can act
        // as an index to the problematic msr.
        if nmsrs != chunk_size {
            Err(KvmVcpuError::VcpuGetMsr(msrs.as_slice()[nmsrs].index))
        } else {
            Ok(msrs)
        }
    }

    /// Get MSRs for the given MSR index list.
    ///
    /// # Arguments
    ///
    /// * `msr_index_list`: List of MSR indices
    ///
    /// # Errors
    ///
    /// * When `KvmVcpu::get_msr_chunks()` 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)
    }

    /// Save the KVM internal state.