feat: 实现 VCpu 操作接口，增强虚拟 CPU 管理功能

Author: ZR233

src/arch/aarch64/cpu.rs

@@ -6,8 +6,9 @@ use arm_vcpu::{Aarch64PerCpu, Aarch64VCpuCreateConfig};
 use axvm_types::addr::*;
 
 use crate::{
+    RunError,
     data::VmDataWeak,
-    vcpu::VCpuCommon,
+    vcpu::{VCpuCommon, VCpuOp},
     vhal::{
         ArchCpuData,
         cpu::{CpuHardId, CpuId, HCpuExclusive},
@@ -125,8 +126,18 @@ impl VCpu {
         .unwrap();
         Ok(VCpu { vcpu, common })
     }
+}
+
+impl VCpuOp for VCpu {
+    fn bind_id(&self) -> CpuId {
+        self.common.bind_id()
+    }
+
+    fn hard_id(&self) -> CpuHardId {
+        self.common.hard_id()
+    }
 
-    pub fn run(&mut self) -> anyhow::Result<()> {
+    fn run(&mut self) -> Result<(), RunError> {
         info!("Starting vCPU {}", self.bind_id());
 
         while self.is_active() {

src/vcpu/mod.rs

@@ -1,10 +1,16 @@
 use crate::{
-    CpuId,
+    CpuId, RunError,
     arch::HCpu,
     data::{VmData, VmDataWeak},
     vhal::cpu::{CpuHardId, HCpuExclusive},
 };
 
+pub trait VCpuOp: Send + 'static{
+    fn bind_id(&self) -> CpuId;
+    fn hard_id(&self) -> CpuHardId;
+    fn run(&mut self) -> Result<(), RunError>;
+}
+
 #[derive(Debug)]
 pub struct VCpuCommon {
     pub(crate) hcpu: HCpuExclusive,

src/vm/data.rs

@@ -11,7 +11,7 @@ use spin::RwLock;
 
 use crate::{
     AxVMConfig, RunError, VmId, VmMachineInitedOps, VmMachineRunningOps, VmMachineUninitOps,
-    arch::{VmMachineInited, VmMachineUninit},
+    arch::{VmMachineInited, VmMachineRunning, VmMachineUninit},
     config::AxVCpuConfig,
     vm::machine::{AtomicState, VMStatus, VmMachineState},
 };
@@ -207,6 +207,23 @@ impl VmDataWeak {
             inner.status.store(VMStatus::Stopped);
         }
     }
+
+    pub(crate) fn with_machine_running<F, R>(&self, f: F) -> Result<R, RunError>
+    where
+        F: FnOnce(&VmMachineRunning) -> R,
+    {
+        let vmdata = self.try_upgrade()?;
+        let status = vmdata.machine.read();
+        let running = match &*status {
+            VmMachineState::Running(running) => running,
+            _ => {
+                return Err(RunError::ExitWithError(anyhow!(
+                    "VM is not in Running state"
+                )));
+            }
+        };
+        Ok(f(running))
+    }
 }
 
 impl Debug for VmDataWeak {

src/vm/machine/running.rs

@@ -8,7 +8,8 @@ use std::{
 use alloc::vec::Vec;
 
 use crate::{
-    RunError, TASK_STACK_SIZE, VmAddrSpace, arch::cpu::VCpu, data::VmDataWeak, vhal::cpu::CpuHardId,
+    RunError, TASK_STACK_SIZE, VmAddrSpace, arch::cpu::VCpu, data::VmDataWeak, vcpu::VCpuOp,
+    vhal::cpu::CpuHardId,
 };
 
 pub struct VmMachineRunningCommon {
@@ -44,7 +45,7 @@ impl VmMachineRunningCommon {
         Ok(cpu)
     }
 
-    pub fn run_cpu(&mut self, mut cpu: VCpu) -> anyhow::Result<()> {
+    pub fn run_cpu<C: VCpuOp>(&mut self, mut cpu: C) -> anyhow::Result<()> {
         let waiter = self.new_waiter();
 
         let bind_id = cpu.bind_id();
@@ -61,7 +62,7 @@ impl VmMachineRunningCommon {
                 let res = cpu.run();
                 if let Err(e) = res {
                     if let Some(vm) = waiter.vm.upgrade() {
-                        vm.set_err(RunError::ExitWithError(e));
+                        vm.set_err(e);
                     }
                 }
                 waiter.running_cpu_count.fetch_sub(1, Ordering::SeqCst);