Repository: firecracker-microvm/firecracker
File: src/vmm/src/builder.rs (Rust)

Code before:
    vmm.uffd = uffd;

    #[cfg(target_arch = "x86_64")]
    {
        // Scale TSC to match, extract the TSC freq from the state if specified
        if let Some(state_tsc) = microvm_state.vcpu_states[0].tsc_khz {
            // Scale the TSC frequency for all VCPUs. If a TSC frequency is not specified in the
            // snapshot, by default it uses the host frequency.
            if vcpus[0].kvm_vcpu.is_tsc_scaling_required(state_tsc)? {
                for vcpu in &vcpus {
                    vcpu.kvm_vcpu.set_tsc_khz(state_tsc)?;
                }
            }
        }
    }

    // Need to alloc sys mem for steal time regions before vcpu setup
    #[cfg(target_arch = "aarch64")]
    if let Some(pvtime_ipa) = vcpus[0].kvm_vcpu.pvtime_ipa {
        allocate_pvtime_region(
            &mut vmm,
            vcpus.len(),
            vm_allocator::AllocPolicy::ExactMatch(pvtime_ipa.0),
        )?;
    }

    // Restore vcpus kvm state.
    for (vcpu, state) in vcpus.iter_mut().zip(microvm_state.vcpu_states.iter()) {
        vcpu.kvm_vcpu
            .restore_state(state)
            .map_err(VcpuError::VcpuResponse)
            .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
    }

    #[cfg(target_arch = "aarch64")]
    {
        let mpidrs = construct_kvm_mpidrs(&microvm_state.vcpu_states);
        // Restore kvm vm state.
        vmm.vm.restore_state(&mpidrs, &microvm_state.vm_state)?;
    }

    // Restore kvm vm state.
    #[cfg(target_arch = "x86_64")]
    vmm.vm.restore_state(&microvm_state.vm_state)?;

    // Restore the boot source config paths.
    vm_resources.boot_source.config = microvm_state.vm_info.boot_source;

    // Restore devices states.
    let mmio_ctor_args = MMIODevManagerConstructorArgs {
        mem: vmm.vm.guest_memory(),

Review comment (suggestion):
Ahh, I see now what you meant by restoring the allocator yesterday :) 

On the snapshot restore path, this is not needed - the allocator only allocates address ranges, it has nothing to do with actually allocating the memory. All this call does is "give me an address range that doesnt overlap any address range you've given me previously". The actual memory setup is restored implicitly because we told the guest all the way back before booting where the pvtime memory should be, and it remembers this information until it's shutdown.

And on the restore path, we dont use the allocators to allocate any new resources anymore, so there's no need to restore their old state


Code after:
        .map_err(StartMicrovmError::Internal)?;
    vmm.uffd = uffd;

    #[cfg(target_arch = "x86_64")]
    {
        // Scale TSC to match, extract the TSC freq from the state if specified
        if let Some(state_tsc) = microvm_state.vcpu_states[0].tsc_khz {
            // Scale the TSC frequency for all VCPUs. If a TSC frequency is not specified in the
            // snapshot, by default it uses the host frequency.
            if vcpus[0].kvm_vcpu.is_tsc_scaling_required(state_tsc)? {
                for vcpu in &vcpus {
                    vcpu.kvm_vcpu.set_tsc_khz(state_tsc)?;
                }
            }
        }
    }

    // Restore allocator state
    #[cfg(target_arch = "aarch64")]
    if let Some(pvtime_ipa) = vcpus[0].kvm_vcpu.pvtime_ipa {
        allocate_pvtime_region(
            &mut vmm,
            vcpus.len(),
            vm_allocator::AllocPolicy::ExactMatch(pvtime_ipa.0),
        )?;
    }

    // Restore vcpus kvm state.
    for (vcpu, state) in vcpus.iter_mut().zip(microvm_state.vcpu_states.iter()) {
        vcpu.kvm_vcpu
            .restore_state(state)
            .map_err(VcpuError::VcpuResponse)
            .map_err(BuildMicrovmFromSnapshotError::RestoreVcpus)?;
    }

    #[cfg(target_arch = "aarch64")]
    {
        let mpidrs = construct_kvm_mpidrs(&microvm_state.vcpu_states);
        // Restore kvm vm state.
        vmm.vm.restore_state(&mpidrs, &microvm_state.vm_state)?;
    }

    // Restore kvm vm state.
    #[cfg(target_arch = "x86_64")]
    vmm.vm.restore_state(&microvm_state.vm_state)?;

    // Restore the boot source config paths.
    vm_resources.boot_source.config = microvm_state.vm_info.boot_source;

    // Restore devices states.
    let mmio_ctor_args = MMIODevManagerConstructorArgs {