From 65dc1c83a9af36d6a0a87863ee8421b16bf7b437 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Tue, 31 May 2022 13:40:14 +0200
Subject: [PATCH] vmm: cpu: Save and restore CPU states during snapshot/restore

Based on recent KVM host patches (merged in Linux 5.16), it's forbidden
to call into KVM_SET_CPUID2 after the first successful KVM_RUN returned.
That means saving CPU states during the pause sequence, and restoring
these states during the resume sequence will not work with the current
design starting with kernel version 5.16.

In order to solve this problem, let's simply move the save/restore logic
to the snapshot/restore sequences rather than the pause/resume ones.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 vmm/src/cpu.rs | 39 +++++++++++++++++----------------------
 1 file changed, 17 insertions(+), 22 deletions(-)

diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index 3bd33a443..c728e61fc 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -375,43 +375,38 @@ impl Vcpu {
 }
 
 const VCPU_SNAPSHOT_ID: &str = "vcpu";
-impl Pausable for Vcpu {
-    fn pause(&mut self) -> std::result::Result<(), MigratableError> {
-        self.saved_state =
-            Some(self.vcpu.state().map_err(|e| {
-                MigratableError::Pause(anyhow!("Could not get vCPU state {:?}", e))
-            })?);
-
-        Ok(())
-    }
-
-    fn resume(&mut self) -> std::result::Result<(), MigratableError> {
-        if let Some(vcpu_state) = &self.saved_state {
-            self.vcpu.set_state(vcpu_state).map_err(|e| {
-                MigratableError::Pause(anyhow!("Could not set the vCPU state {:?}", e))
-            })?;
-        }
-
-        Ok(())
-    }
-}
+impl Pausable for Vcpu {}
 impl Snapshottable for Vcpu {
     fn id(&self) -> String {
         VCPU_SNAPSHOT_ID.to_string()
     }
 
     fn snapshot(&mut self) -> std::result::Result<Snapshot, MigratableError> {
+        let saved_state = self
+            .vcpu
+            .state()
+            .map_err(|e| MigratableError::Pause(anyhow!("Could not get vCPU state {:?}", e)))?;
+
         let mut vcpu_snapshot = Snapshot::new(&format!("{}", self.id));
         vcpu_snapshot.add_data_section(SnapshotDataSection::new_from_state(
             VCPU_SNAPSHOT_ID,
-            &self.saved_state,
+            &saved_state,
         )?);
 
+        self.saved_state = Some(saved_state);
+
         Ok(vcpu_snapshot)
     }
 
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
-        self.saved_state = Some(snapshot.to_state(VCPU_SNAPSHOT_ID)?);
+        let saved_state: CpuState = snapshot.to_state(VCPU_SNAPSHOT_ID)?;
+
+        self.vcpu
+            .set_state(&saved_state)
+            .map_err(|e| MigratableError::Pause(anyhow!("Could not set the vCPU state {:?}", e)))?;
+
+        self.saved_state = Some(saved_state);
+
         Ok(())
     }
 }