snapshot-editor: expanding CLI to support tsc set/clear commands

CHANGELOG.md

@@ -10,6 +10,9 @@ and this project adheres to
 
 ### Added
 
+- Added snapshot-editor `tsc set/clear` commands to stamp or remove TSC
+  frequencies in vmstate snapshots in place (auto-detects host TSC on x86_64).
+
 ### Changed
 
 ### Deprecated

docs/snapshotting/snapshot-editor.md

@@ -136,3 +136,38 @@ Firecracker snapshot consists of 2 files:
 > ```bash
 > ./snapshot-editor info-vmstate vm-state --vmstate-path ./vmstate_file
 > ```
+
+### `tsc` command
+
+#### `set` subcommand
+
+> Set the saved TSC frequency (in kHz) for every vCPU in a vmstate file.
+>
+> Arguments:
+>
+> - `VMSTATE_PATH` - path to the `vmstate` file
+> - `--tsc-khz` - explicit TSC frequency (required on non-x86_64); on x86_64 it
+>   is auto-detected from the host if omitted
+>
+> Usage:
+>
+> ```bash
+> snapshot-editor tsc set \
+>     --vmstate-path ./vmstate_file \
+>     --tsc-khz 2500000
+> ```
+
+#### `clear` subcommand
+
+> Remove the saved TSC frequency so Firecracker skips scaling on restore.
+>
+> Arguments:
+>
+> - `VMSTATE_PATH` - path to the `vmstate` file
+>
+> Usage:
+>
+> ```bash
+> snapshot-editor tsc clear \
+>     --vmstate-path ./vmstate_file
+> ```

src/snapshot-editor/Cargo.toml

@@ -27,5 +27,8 @@ vmm-sys-util = "0.15.0"
 [target.'cfg(target_arch = "aarch64")'.dependencies]
 clap-num = "1.2.0"
 
+[target.'cfg(target_arch = "x86_64")'.dependencies]
+kvm-ioctls = "0.19.1"
+
 [lints]
 workspace = true

src/snapshot-editor/src/edit_vmstate.rs

@@ -51,9 +51,9 @@ fn edit(
     output_path: &PathBuf,
     f: impl Fn(MicrovmState) -> Result<MicrovmState, EditVmStateError>,
 ) -> Result<(), EditVmStateError> {
-    let snapshot = open_vmstate(vmstate_path)?;
-    let microvm_state = f(snapshot.data)?;
-    save_vmstate(microvm_state, output_path)?;
+    let mut snapshot = open_vmstate(vmstate_path)?;
+    snapshot.data = f(snapshot.data)?;
+    save_vmstate(&snapshot, output_path)?;
     Ok(())
 }
 

src/snapshot-editor/src/main.rs

@@ -7,12 +7,14 @@ mod edit_memory;
 #[cfg(target_arch = "aarch64")]
 mod edit_vmstate;
 mod info;
+mod tsc;
 mod utils;
 
 use edit_memory::{EditMemoryError, EditMemorySubCommand, edit_memory_command};
 #[cfg(target_arch = "aarch64")]
 use edit_vmstate::{EditVmStateError, EditVmStateSubCommand, edit_vmstate_command};
 use info::{InfoVmStateError, InfoVmStateSubCommand, info_vmstate_command};
+use tsc::{TscCommandError, TscSubCommand, tsc_command};
 
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 enum SnapEditorError {
@@ -23,6 +25,8 @@ enum SnapEditorError {
     EditVmState(#[from] EditVmStateError),
     /// Error during getting info from a vmstate file: {0}
     InfoVmState(#[from] InfoVmStateError),
+    /// Error during updating TSC metadata: {0}
+    EditTsc(#[from] TscCommandError),
 }
 
 #[derive(Debug, Parser)]
@@ -41,6 +45,8 @@ enum Command {
     EditVmstate(EditVmStateSubCommand),
     #[command(subcommand)]
     InfoVmstate(InfoVmStateSubCommand),
+    #[command(subcommand)]
+    Tsc(TscSubCommand),
 }
 
 fn main_exec() -> Result<(), SnapEditorError> {
@@ -51,6 +57,7 @@ fn main_exec() -> Result<(), SnapEditorError> {
         #[cfg(target_arch = "aarch64")]
         Command::EditVmstate(command) => edit_vmstate_command(command)?,
         Command::InfoVmstate(command) => info_vmstate_command(command)?,
+        Command::Tsc(command) => tsc_command(command)?,
     }
 
     Ok(())

src/snapshot-editor/src/tsc.rs

@@ -0,0 +1,165 @@
+// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::PathBuf;
+
+use clap::{Args, Subcommand};
+
+use crate::utils::{UtilsError, open_vmstate, save_vmstate};
+
+#[derive(Debug, thiserror::Error)]
+pub enum TscCommandError {
+    #[error("{0}")]
+    Utils(#[from] UtilsError),
+    #[cfg_attr(target_arch = "x86_64", allow(dead_code))]
+    #[error("Missing --tsc-khz value; provide a target frequency in kHz.")]
+    MissingFrequency,
+    #[cfg(target_arch = "x86_64")]
+    #[error("Failed to open /dev/kvm: {0}")]
+    DetectOpenKvm(kvm_ioctls::Error),
+    #[cfg(target_arch = "x86_64")]
+    #[error("Failed to create KVM VM: {0}")]
+    DetectCreateVm(kvm_ioctls::Error),
+    #[cfg(target_arch = "x86_64")]
+    #[error("Failed to create KVM vCPU: {0}")]
+    DetectCreateVcpu(kvm_ioctls::Error),
+    #[cfg(target_arch = "x86_64")]
+    #[error("Failed to query TSC frequency from KVM: {0}")]
+    DetectQueryTsc(kvm_ioctls::Error),
+}
+
+#[derive(Debug, Subcommand)]
+pub enum TscSubCommand {
+    /// Set the saved TSC frequency (in kHz) for every vCPU in the vmstate file.
+    Set(SetTscArgs),
+    /// Remove the saved TSC frequency so Firecracker skips scaling on restore.
+    Clear(ClearTscArgs),
+}
+
+#[derive(Debug, Args)]
+pub struct SetTscArgs {
+    /// Path to the vmstate file to update.
+    #[arg(long)]
+    pub vmstate_path: PathBuf,
+    /// TSC frequency in kHz to embed in the vmstate snapshot.
+    #[arg(long, value_parser = clap::value_parser!(u32))]
+    pub tsc_khz: Option<u32>,
+}
+
+#[derive(Debug, Args)]
+pub struct ClearTscArgs {
+    /// Path to the vmstate file to update.
+    #[arg(long)]
+    pub vmstate_path: PathBuf,
+}
+
+pub fn tsc_command(command: TscSubCommand) -> Result<(), TscCommandError> {
+    match command {
+        TscSubCommand::Set(args) => set_tsc(args),
+        TscSubCommand::Clear(args) => clear_tsc(args),
+    }
+}
+
+fn set_tsc(args: SetTscArgs) -> Result<(), TscCommandError> {
+    #[cfg(target_arch = "x86_64")]
+    let freq = match args.tsc_khz {
+        Some(freq) => freq,
+        None => detect_host_tsc_khz()?,
+    };
+    #[cfg(not(target_arch = "x86_64"))]
+    let freq = args.tsc_khz.ok_or(TscCommandError::MissingFrequency)?;
+
+    let mut snapshot = open_vmstate(&args.vmstate_path)?;
+    for vcpu in &mut snapshot.data.vcpu_states {
+        vcpu.tsc_khz = Some(freq);
+    }
+    save_vmstate(&snapshot, &args.vmstate_path)?;
+    Ok(())
+}
+
+#[cfg(target_arch = "x86_64")]
+fn detect_host_tsc_khz() -> Result<u32, TscCommandError> {
+    use kvm_ioctls::Kvm;
+
+    let kvm = Kvm::new().map_err(TscCommandError::DetectOpenKvm)?;
+    let vm = kvm.create_vm().map_err(TscCommandError::DetectCreateVm)?;
+    let vcpu = vm
+        .create_vcpu(0)
+        .map_err(TscCommandError::DetectCreateVcpu)?;
+    vcpu.get_tsc_khz().map_err(TscCommandError::DetectQueryTsc)
+}
+
+fn clear_tsc(args: ClearTscArgs) -> Result<(), TscCommandError> {
+    let mut snapshot = open_vmstate(&args.vmstate_path)?;
+    for vcpu in &mut snapshot.data.vcpu_states {
+        vcpu.tsc_khz = None;
+    }
+    save_vmstate(&snapshot, &args.vmstate_path)?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    use vmm::persist::MicrovmState;
+    use vmm::snapshot::Snapshot;
+
+    use super::*;
+    use crate::utils::save_vmstate;
+
+    fn temp_vmstate_path() -> PathBuf {
+        let mut path = std::env::temp_dir();
+        let nanos = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos();
+        path.push(format!(
+            "snapshot-editor-tsc-{}-{}.bin",
+            std::process::id(),
+            nanos
+        ));
+        path
+    }
+
+    #[test]
+    fn test_tsc_set_and_clear_in_place() {
+        let vmstate_path = temp_vmstate_path();
+
+        // Start from a valid vmstate snapshot.
+        let snapshot = Snapshot::new(MicrovmState::default());
+        save_vmstate(&snapshot, &vmstate_path).expect("save initial vmstate");
+
+        let set_freq = 123_456u32;
+        set_tsc(SetTscArgs {
+            vmstate_path: vmstate_path.clone(),
+            tsc_khz: Some(set_freq),
+        })
+        .expect("tsc set should succeed");
+
+        let snapshot = open_vmstate(&vmstate_path).expect("vmstate after set");
+        assert!(
+            snapshot
+                .data
+                .vcpu_states
+                .iter()
+                .all(|vcpu| vcpu.tsc_khz == Some(set_freq))
+        );
+
+        clear_tsc(ClearTscArgs {
+            vmstate_path: vmstate_path.clone(),
+        })
+        .expect("tsc clear should succeed");
+
+        let snapshot = open_vmstate(&vmstate_path).expect("vmstate after clear");
+        assert!(
+            snapshot
+                .data
+                .vcpu_states
+                .iter()
+                .all(|vcpu| vcpu.tsc_khz.is_none())
+        );
+
+        let _ = std::fs::remove_file(vmstate_path);
+    }
+}

src/snapshot-editor/src/utils.rs

@@ -27,16 +27,17 @@ pub fn open_vmstate(snapshot_path: &PathBuf) -> Result<Snapshot<MicrovmState>, U
     Snapshot::load(&mut snapshot_reader).map_err(UtilsError::VmStateLoad)
 }
 
-// This method is used only in aarch64 code so far
 #[allow(unused)]
-pub fn save_vmstate(microvm_state: MicrovmState, output_path: &PathBuf) -> Result<(), UtilsError> {
+pub fn save_vmstate(
+    snapshot: &Snapshot<MicrovmState>,
+    output_path: &PathBuf,
+) -> Result<(), UtilsError> {
     let mut output_file = OpenOptions::new()
         .create(true)
         .write(true)
         .truncate(true)
         .open(output_path)
         .map_err(UtilsError::OutputFileOpen)?;
-    let mut snapshot = Snapshot::new(microvm_state);
     snapshot
         .save(&mut output_file)
         .map_err(UtilsError::VmStateSave)?;