From 578780ed0c03a0009874d4a286710d79f324c253 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Fri, 4 Nov 2022 10:27:22 +0100
Subject: [PATCH] vmm: cpu: Split vCPU creation

Split the vCPU creation into two distincts parts. On the one hand we
create the actual Vcpu object with the creation of the hypervisor::Vcpu.
And on the other hand, we configure the existing Vcpu, setting registers
to proper values (such as setting the entry point).

This will allow for further work to move the creation earlier in the
boot, so that the hypervisor::Vcpu will be already created when the
DeviceManager gets created.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
Signed-off-by: Michael Zhao <michael.zhao@arm.com>
---
 vmm/src/cpu.rs | 56 +++++++++++++++++++++++++++++++++++---------------
 vmm/src/vm.rs  | 15 ++++++++------
 2 files changed, 48 insertions(+), 23 deletions(-)

diff --git a/vmm/src/cpu.rs b/vmm/src/cpu.rs
index 681256db0..4e693f98d 100644
--- a/vmm/src/cpu.rs
+++ b/vmm/src/cpu.rs
@@ -720,15 +720,28 @@ impl CpuManager {
         })))
     }
 
-    fn create_vcpu(
-        &mut self,
-        cpu_id: u8,
+    fn create_vcpu(&mut self, cpu_id: u8) -> Result<Arc<Mutex<Vcpu>>> {
+        info!("Creating vCPU: cpu_id = {}", cpu_id);
+
+        let vcpu = Arc::new(Mutex::new(Vcpu::new(
+            cpu_id,
+            &self.vm,
+            Some(self.vm_ops.clone()),
+        )?));
+
+        // Adding vCPU to the CpuManager's vCPU list.
+        self.vcpus.push(vcpu.clone());
+
+        Ok(vcpu)
+    }
+
+    pub fn configure_vcpu(
+        &self,
+        vcpu: Arc<Mutex<Vcpu>>,
         entry_point: Option<EntryPoint>,
         snapshot: Option<Snapshot>,
     ) -> Result<()> {
-        info!("Creating vCPU: cpu_id = {}", cpu_id);
-
-        let mut vcpu = Vcpu::new(cpu_id, &self.vm, Some(self.vm_ops.clone()))?;
+        let mut vcpu = vcpu.lock().unwrap();
 
         if let Some(snapshot) = snapshot {
             // AArch64 vCPUs should be initialized after created.
@@ -751,15 +764,12 @@ impl CpuManager {
                 .expect("Failed to configure vCPU");
         }
 
-        // Adding vCPU to the CpuManager's vCPU list.
-        let vcpu = Arc::new(Mutex::new(vcpu));
-        self.vcpus.push(vcpu);
-
         Ok(())
     }
 
     /// Only create new vCPUs if there aren't any inactive ones to reuse
-    fn create_vcpus(&mut self, desired_vcpus: u8, entry_point: Option<EntryPoint>) -> Result<()> {
+    fn create_vcpus(&mut self, desired_vcpus: u8) -> Result<Vec<Arc<Mutex<Vcpu>>>> {
+        let mut vcpus: Vec<Arc<Mutex<Vcpu>>> = vec![];
         info!(
             "Request to create new vCPUs: desired = {}, max = {}, allocated = {}, present = {}",
             desired_vcpus,
@@ -774,10 +784,10 @@ impl CpuManager {
 
         // Only create vCPUs in excess of all the allocated vCPUs.
         for cpu_id in self.vcpus.len() as u8..desired_vcpus {
-            self.create_vcpu(cpu_id, entry_point, None)?;
+            vcpus.push(self.create_vcpu(cpu_id)?);
         }
 
-        Ok(())
+        Ok(vcpus)
     }
 
     #[cfg(target_arch = "aarch64")]
@@ -799,6 +809,10 @@ impl CpuManager {
         Ok(true)
     }
 
+    pub fn vcpus(&self) -> Vec<Arc<Mutex<Vcpu>>> {
+        self.vcpus.clone()
+    }
+
     fn start_vcpu(
         &mut self,
         vcpu: Arc<Mutex<Vcpu>>,
@@ -1102,10 +1116,10 @@ impl CpuManager {
         Ok(())
     }
 
-    pub fn create_boot_vcpus(&mut self, entry_point: Option<EntryPoint>) -> Result<()> {
+    pub fn create_boot_vcpus(&mut self) -> Result<Vec<Arc<Mutex<Vcpu>>>> {
         trace_scoped!("create_boot_vcpus");
 
-        self.create_vcpus(self.boot_vcpus(), entry_point)
+        self.create_vcpus(self.boot_vcpus())
     }
 
     // Starts all the vCPUs that the VM is booting with. Blocks until all vCPUs are running.
@@ -1133,7 +1147,10 @@ impl CpuManager {
 
         match desired_vcpus.cmp(&self.present_vcpus()) {
             cmp::Ordering::Greater => {
-                self.create_vcpus(desired_vcpus, None)?;
+                let vcpus = self.create_vcpus(desired_vcpus)?;
+                for vcpu in vcpus {
+                    self.configure_vcpu(vcpu, None, None)?
+                }
                 self.activate_vcpus(desired_vcpus, true, None)?;
                 Ok(true)
             }
@@ -2066,8 +2083,13 @@ impl Snapshottable for CpuManager {
     fn restore(&mut self, snapshot: Snapshot) -> std::result::Result<(), MigratableError> {
         for (cpu_id, snapshot) in snapshot.snapshots.iter() {
             info!("Restoring VCPU {}", cpu_id);
-            self.create_vcpu(cpu_id.parse::<u8>().unwrap(), None, Some(*snapshot.clone()))
+            let vcpu = self
+                .create_vcpu(cpu_id.parse::<u8>().unwrap())
                 .map_err(|e| MigratableError::Restore(anyhow!("Could not create vCPU {:?}", e)))?;
+            self.configure_vcpu(vcpu, None, Some(*snapshot.clone()))
+                .map_err(|e| {
+                    MigratableError::Restore(anyhow!("Could not configure vCPU {:?}", e))
+                })?
         }
 
         Ok(())
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index c2ac1f8ae..7d7273c51 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -2087,12 +2087,15 @@ impl Vm {
             self.init_tdx()?;
         }
 
-        // Create and configure vcpus
-        self.cpu_manager
-            .lock()
-            .unwrap()
-            .create_boot_vcpus(entry_point)
-            .map_err(Error::CpuManager)?;
+        // Configure the vcpus that have been created
+        let vcpus = self.cpu_manager.lock().unwrap().vcpus();
+        for vcpu in vcpus {
+            self.cpu_manager
+                .lock()
+                .unwrap()
+                .configure_vcpu(vcpu, entry_point, None)
+                .map_err(Error::CpuManager)?;
+        }
 
         #[cfg(feature = "tdx")]
         let sections = if tdx_enabled {