Refactor VM data management and address space handling

- Introduced `VmData` and `VmDataWeak` structures to encapsulate VM state and configuration.
- Removed the obsolete `data2.rs` file and migrated relevant functionality to `data.rs`.
- Enhanced `VmAddrSpace` to manage memory regions more effectively, including methods for adding and initializing memory.
- Updated `FdtBuilder` to streamline CPU and memory setup in the device tree.
- Refactored `VCpuCommon` to utilize the new `HCpu` architecture.
- Improved error handling and state management in VM lifecycle operations.
- Added new traits for VM machine states to facilitate better state transitions and management.

Author: ZR233

src/arch/aarch64/cpu.rs

@@ -6,7 +6,7 @@ use arm_vcpu::{Aarch64PerCpu, Aarch64VCpuCreateConfig};
 use axvm_types::addr::*;
 
 use crate::{
-    data2::VmDataWeak,
+    data::VmDataWeak,
     vcpu::VCpuCommon,
     vhal::{
         ArchCpuData,
@@ -19,6 +19,7 @@ pub struct HCpu {
     pub hard_id: CpuHardId,
     vpercpu: Aarch64PerCpu,
     max_guest_page_table_levels: usize,
+    pub pa_range: core::ops::Range<usize>,
 }
 
 impl HCpu {
@@ -33,12 +34,14 @@ impl HCpu {
             hard_id: CpuHardId::new(hard_id),
             vpercpu,
             max_guest_page_table_levels: 0,
+            pa_range: 0..0,
         }
     }
 
     pub fn init(&mut self) -> anyhow::Result<()> {
         self.vpercpu.hardware_enable();
         self.max_guest_page_table_levels = self.vpercpu.max_guest_page_table_levels();
+        self.pa_range = self.vpercpu.pa_range();
         Ok(())
     }
 
@@ -124,11 +127,15 @@ impl VCpu {
     }
 
     pub fn run(&mut self) -> anyhow::Result<()> {
-        info!("Starting vCPU {}", self.id());
+        info!("Starting vCPU {}", self.bind_id());
 
         while self.is_active() {
             let exit_reason = self.vcpu.run().map_err(|e| anyhow!("{e}"))?;
-            debug!("vCPU {} exited with reason: {:?}", self.id(), exit_reason);
+            debug!(
+                "vCPU {} exited with reason: {:?}",
+                self.bind_id(),
+                exit_reason
+            );
             match exit_reason {
                 arm_vcpu::AxVCpuExitReason::Hypercall { nr, args } => todo!(),
                 arm_vcpu::AxVCpuExitReason::MmioRead {
@@ -149,8 +156,8 @@ impl VCpu {
                 } => todo!(),
                 arm_vcpu::AxVCpuExitReason::CpuDown { _state } => todo!(),
                 arm_vcpu::AxVCpuExitReason::SystemDown => {
-                    info!("vCPU {} requested system shutdown", self.common.bind_id);
-                    self.shutdown()?;
+                    info!("vCPU {} requested system shutdown", self.bind_id());
+                    self.vm()?.stop()?;
                 }
                 arm_vcpu::AxVCpuExitReason::Nothing => {}
                 arm_vcpu::AxVCpuExitReason::SendIPI {

src/arch/aarch64/vm/init.rs

@@ -8,11 +8,11 @@ use std::{
 use arm_vcpu::Aarch64VCpuSetupConfig;
 
 use crate::{
-    GuestPhysAddr, TASK_STACK_SIZE, VmRunCommonData, VmStatusInitOps,
-    arch::{VmStatusRunning, cpu::VCpu},
+    GuestPhysAddr, TASK_STACK_SIZE, VmAddrSpace, VmMachineInitedOps,
+    arch::{VmMachineRunning, cpu::VCpu},
     config::AxVMConfig,
-    data2::VmDataWeak,
-    vm::{MappingFlags, VmId},
+    data::VmDataWeak,
+    vm::VmId,
 };
 
 const VM_ASPACE_BASE: GuestPhysAddr = GuestPhysAddr::from_usize(0);
@@ -23,128 +23,14 @@ const VM_ASPACE_END: GuestPhysAddr =
 pub struct VmMachineInited {
     pub id: VmId,
     pub name: String,
-    // pt_levels: usize,
-    // stop_requested: AtomicBool,
-    pub run_data: VmStatusRunning,
+    pub vcpus: Vec<VCpu>,
+    pub vmspace: VmAddrSpace,
 }
 
-impl VmMachineInited {
-    /// Creates a new VM with the given configuration
-    // pub fn new(config: &AxVMConfig) -> anyhow::Result<Self> {
-    //     let vm = Self {
-    //         id: config.id().into(),
-    //         name: config.name().into(),
-    //         pt_levels: 4,
-    //         stop_requested: AtomicBool::new(false),
-    //         run_data: None,
-    //     };
-    //     Ok(vm)
-    // }
-    
+impl VmMachineInited {}
 
-    /// Initializes the VM, creating vCPUs and setting up memory
-    pub fn init(&mut self, config: AxVMConfig) -> anyhow::Result<()> {
-        debug!("Initializing VM {} ({})", self.id, self.name);
-
-        let vcpus = self.new_vcpus(&config)?;
-
-        let mut run_data = VmStatusRunning::new(
-            VmRunCommonData::new(self.pt_levels, VM_ASPACE_BASE..VM_ASPACE_END)?,
-            vcpus,
-        );
-
-        debug!("Mapping memory regions for VM {} ({})", self.id, self.name);
-        for memory_cfg in &config.memory_regions {
-            use crate::vm::MappingFlags;
-            let m = run_data.data.try_use()?.new_memory(
-                memory_cfg,
-                MappingFlags::READ
-                    | MappingFlags::WRITE
-                    | MappingFlags::EXECUTE
-                    | MappingFlags::USER,
-            );
-            run_data.data.add_memory(m);
-        }
-
-        run_data.data.try_use()?.load_kernel_image(&config)?;
-        run_data.make_dtb(&config)?;
-
-        run_data.data.try_use()?.map_passthrough_regions()?;
-
-        let kernel_entry = run_data.data.try_use()?.kernel_entry();
-        let gpt_root = run_data.data.try_use()?.gpt_root();
-
-        // Setup vCPUs
-        for vcpu in &mut run_data.vcpus {
-            vcpu.vcpu.set_entry(kernel_entry).unwrap();
-            vcpu.vcpu.set_dtb_addr(run_data.dtb_addr).unwrap();
-
-            let setup_config = Aarch64VCpuSetupConfig {
-                passthrough_interrupt: config.interrupt_mode()
-                    == axvmconfig::VMInterruptMode::Passthrough,
-                passthrough_timer: config.interrupt_mode()
-                    == axvmconfig::VMInterruptMode::Passthrough,
-            };
-
-            vcpu.vcpu
-                .setup(setup_config)
-                .map_err(|e| anyhow::anyhow!("Failed to setup vCPU : {e:?}"))?;
-
-            // Set EPT root
-            vcpu.vcpu
-                .set_ept_root(gpt_root)
-                .map_err(|e| anyhow::anyhow!("Failed to set EPT root for vCPU : {e:?}"))?;
-
-            run_data.vcpu_running_count.fetch_add(1, Ordering::SeqCst);
-        }
-
-        self.run_data = Some(run_data);
-
-        Ok(())
-    }
-
-    fn new_vcpus(&mut self, config: &AxVMConfig) -> anyhow::Result<Vec<VCpu>> {
-        // Create vCPUs
-        let mut vcpus = Vec::new();
-
-        let dtb_addr = GuestPhysAddr::from_usize(0);
-
-        match config.cpu_num {
-            crate::config::CpuNumType::Alloc(num) => {
-                for _ in 0..num {
-                    let vcpu = VCpu::new(None, dtb_addr)?;
-                    debug!("Created vCPU with {:?}", vcpu.id);
-                    vcpus.push(vcpu);
-                }
-            }
-            crate::config::CpuNumType::Fixed(ref ids) => {
-                for id in ids {
-                    let vcpu = VCpu::new(Some(*id), dtb_addr)?;
-                    debug!("Created vCPU with {:?}", vcpu.id);
-                    vcpus.push(vcpu);
-                }
-            }
-        }
-
-        let vcpu_count = vcpus.len();
-
-        for vcpu in &vcpus {
-            let max_levels = vcpu.with_hcpu(|cpu| cpu.max_guest_page_table_levels());
-            if max_levels < self.pt_levels {
-                self.pt_levels = max_levels;
-            }
-        }
-
-        debug!(
-            "VM {} ({}) vCPU count: {}, Max Guest Page Table Levels: {}",
-            self.id, self.name, vcpu_count, self.pt_levels
-        );
-        Ok(vcpus)
-    }
-}
-
-impl VmStatusInitOps for VmMachineInited {
-    type Running = VmStatusRunning;
+impl VmMachineInitedOps for VmMachineInited {
+    type Running = VmMachineRunning;
 
     fn id(&self) -> VmId {
         self.id
@@ -155,51 +41,14 @@ impl VmStatusInitOps for VmMachineInited {
     }
 
     fn start(self, vmdata: VmDataWeak) -> Result<Self::Running, (anyhow::Error, Self)> {
-        let mut data = self.run_data.unwrap();
-
-        let mut vcpus = vec![];
-
-        vcpus.append(&mut data.vcpus);
-        let mut vcpu_handles = vec![];
-        let vm_id = self.id;
-
-        for mut vcpu in vcpus.into_iter() {
-            let vcpu_id = vcpu.id;
-            let vcpu_running_count = data.vcpu_running_count.clone();
-            let bind_id = vcpu.binded_cpu_id();
-            let handle = std::thread::Builder::new()
-                .name(format!("{vm_id}-{vcpu_id}"))
-                .stack_size(TASK_STACK_SIZE)
-                .spawn(move || {
-                    assert!(
-                        set_current_affinity(AxCpuMask::one_shot(bind_id.raw())),
-                        "Initialize CPU affinity failed!"
-                    );
-                    match vcpu.run() {
-                        Ok(()) => {
-                            info!("vCPU {} of VM {} exited normally", vcpu_id, vm_id);
-                        }
-                        Err(e) => {
-                            error!(
-                                "vCPU {} of VM {} exited with error: {:?}",
-                                vcpu_id, vm_id, e
-                            );
-                        }
-                    }
-                    vcpu_running_count.fetch_sub(1, Ordering::SeqCst);
-                    vcpu
-                })
-                .unwrap();
-
-            vcpu_handles.push(handle);
-        }
-
+        debug!("Starting VM {} ({})", self.id, self.name);
+        let running = VmMachineRunning::new();
         info!(
             "VM {} ({}) with {} cpus booted successfully.",
             self.id,
             self.name,
-            vcpu_handles.len()
+            self.vcpus.len()
         );
-        Ok(data)
+        Ok(running)
     }
 }