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

Code before:
                ))
            }
            _ => {
                METRICS.vcpu.failures.inc();
                error!("Failure during vcpu run: {}", err);
                Err(VcpuError::FaultyKvmExit(format!("{}", err)))
            }
        },
    }
}

impl Drop for Vcpu {
    fn drop(&mut self) {
        Self::TLS_VCPU_PTR.with(|cell| {
            // The reason for not asserting TLS being set here is that
            // it can happen that Vcpu::Drop is called on vcpus which never were
            // put on their threads. This can happen if some error occurs during Vmm
            // setup before `start_threaded` call.
            if let Some(_vcpu_ptr) = cell.get() {
                // During normal runtime there is a strong assumption that vcpus will be
                // put on their own threads, thus TLS will be initialized with the
                // correct pointer.
                // In test we do not put vcpus on separate threads, so TLS will have a value
                // of the last created vcpu.
                #[cfg(not(test))]
                assert!(std::ptr::eq(_vcpu_ptr, self));

                Self::TLS_VCPU_PTR.take();
            }
        })
    }
}

/// List of events that the Vcpu can receive.
#[derive(Debug, Clone)]
pub enum VcpuEvent {
    /// The vCPU thread will end when receiving this message.
    Finish,
    /// Pause the Vcpu.
    Pause,
    /// Event to resume the Vcpu.
    Resume,
    /// Event to save the state of a paused Vcpu.
    SaveState,
    /// Event to dump CPU configuration of a paused Vcpu.
    DumpCpuConfig,
}

/// List of responses that the Vcpu reports.
pub enum VcpuResponse {
    /// Requested action encountered an error.

Review comment (suggestion):
I think this is one of those cases where having a separate function return a result and then the caller decide what to do with it is better than having a assert that depends on whether we're testing or not, as it makes the code actually testable without hacks like this.

All this change would be is unwrap the result of `self.reset_thread_local_data()`

Code after:
    /// Event to dump CPU configuration of a paused Vcpu.
    DumpCpuConfig,
}

/// List of responses that the Vcpu reports.
pub enum VcpuResponse {
    /// Requested action encountered an error.
    Error(VcpuError),
    /// Vcpu is stopped.
    Exited(FcExitCode),
    /// Requested action not allowed.
    NotAllowed(String),
    /// Vcpu is paused.
    Paused,
    /// Vcpu is resumed.
    Resumed,
    /// Vcpu state is saved.
    SavedState(Box<VcpuState>),
    /// Vcpu is in the state where CPU config is dumped.
    DumpedCpuConfig(Box<CpuConfiguration>),
}

impl fmt::Debug for VcpuResponse {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use crate::VcpuResponse::*;
        match self {
            Paused => write!(f, "VcpuResponse::Paused"),
            Resumed => write!(f, "VcpuResponse::Resumed"),
            Exited(code) => write!(f, "VcpuResponse::Exited({:?})", code),
            SavedState(_) => write!(f, "VcpuResponse::SavedState"),
            Error(err) => write!(f, "VcpuResponse::Error({:?})", err),
            NotAllowed(reason) => write!(f, "VcpuResponse::NotAllowed({})", reason),
            DumpedCpuConfig(_) => write!(f, "VcpuResponse::DumpedCpuConfig"),
        }
    }
}

/// Wrapper over Vcpu that hides the underlying interactions with the Vcpu thread.
#[derive(Debug)]
pub struct VcpuHandle {
    event_sender: Sender<VcpuEvent>,
    response_receiver: Receiver<VcpuResponse>,
    // Rust JoinHandles have to be wrapped in Option if you ever plan on 'join()'ing them.
    // We want to be able to join these threads in tests.
    vcpu_thread: Option<thread::JoinHandle<()>>,
}

/// Error type for [`VcpuHandle::send_event`].
#[derive(Debug, derive_more::From, thiserror::Error)]
#[error("Failed to signal vCPU: {0}")]
pub struct VcpuSendEventError(pub vmm_sys_util::errno::Error);