drivers: smp: Replace use of .expect() with an Error enum

Use this new error in the primary_init() method via the main module's
Error enum

Temporarily use a .expect() in the _secondary_init as it has not been
split as per primary_init().

See: #100

Signed-off-by: Rob Bradford <rbradford@rivosinc.com>

src/main.rs

@@ -369,6 +369,8 @@ enum Error {
     CpuTopologyGeneration(drivers::cpu::Error),
     /// Creating hypervisor map failed
     CreateHypervisorMap(hyp_map::Error),
+    /// Creating (per CPU) SMP state
+    CreateSmpState(smp::Error),
     /// Problem creating derived device tree
     FdtCreation(DeviceTreeError),
     /// Problem parsing device tree
@@ -383,6 +385,8 @@ enum Error {
     RequiredDeviceProbe(RequiredDeviceProbe),
     /// Setup of user-mode failed
     SetupUserMode(umode::Error),
+    /// Problem running secondary CPUs
+    StartSecondaryCpus(smp::Error),
     /// User-mode NOP failed
     UserModeNop(umode::Error),
 }
@@ -392,16 +396,18 @@ impl Display for Error {
         use Error::*;
         match self {
             BuildMemoryMap(e) => write!(f, "Failed to build memory map: {:?}", e),
-            CreateHypervisorMap(e) => write!(f, "Cannot create Hypervisor map: {:?}", e),
             CpuMissingFeature(feature) => write!(f, "Missing required CPU feature: {:?}", feature),
             CpuTopologyGeneration(e) => write!(f, "Failed to generate CPU topology: {}", e),
+            CreateHypervisorMap(e) => write!(f, "Cannot create Hypervisor map: {:?}", e),
+            CreateSmpState(e) => write!(f, "Error during (per CPU) SMP setup: {:?}", e),
             FdtCreation(e) => write!(f, "Failed to construct device-tree: {}", e),
             FdtParsing(e) => write!(f, "Failed to read FDT: {}", e),
             HeapOutOfSpace => write!(f, "Not enough free memory for hypervisor heap"),
             KernelMissing => write!(f, "No host kernel image"),
             LoadUserMode(e) => write!(f, "Cannot load user-mode ELF binary: {:?}", e),
             RequiredDeviceProbe(e) => write!(f, "Failed to probe required device: {}", e),
             SetupUserMode(e) => write!(f, "Failed to setup user-mode: {:?}", e),
+            StartSecondaryCpus(e) => write!(f, "Error running secondary CPUs: {:?}", e),
             UserModeNop(e) => write!(f, "Failed to execute a NOP in user-mode: {:?}", e),
         }
     }
@@ -586,7 +592,7 @@ fn primary_init(hart_id: u64, fdt_addr: u64) -> Result<CpuParams, Error> {
     HypMap::init(mem_map, &umode_elf).map_err(Error::CreateHypervisorMap)?;
 
     // Set up per-CPU memory and prepare the structures for secondary CPUs boot.
-    PerCpu::init(hart_id, &mut hyp_mem);
+    PerCpu::init(hart_id, &mut hyp_mem).map_err(Error::CreateSmpState)?;
 
     // Find and initialize the IOMMU.
     match Iommu::probe_from(PcieRoot::get(), &mut || {
@@ -621,7 +627,7 @@ fn primary_init(hart_id: u64, fdt_addr: u64) -> Result<CpuParams, Error> {
     // Lock down the boot time allocator before allowing the host VM to be entered.
     HYPERVISOR_ALLOCATOR.get().unwrap().seal();
 
-    smp::start_secondary_cpus();
+    smp::start_secondary_cpus().map_err(Error::StartSecondaryCpus)?;
 
     HOST_VM.call_once(|| host);
 
@@ -690,7 +696,7 @@ extern "C" fn _secondary_main() {}
 #[no_mangle]
 extern "C" fn _secondary_init(hart_id: u64) -> CpuParams {
     setup_csrs();
-    PerCpu::setup_this_cpu(hart_id);
+    PerCpu::setup_this_cpu(hart_id).expect("Failed to create SMP state");
 
     test_declare_pass!("secondary init", hart_id);
 

src/smp.rs

@@ -4,6 +4,7 @@
 
 use core::arch::asm;
 use core::cell::{RefCell, RefMut};
+use core::fmt;
 use drivers::{imsic::Imsic, CpuId, CpuInfo};
 use page_tracking::HypPageAlloc;
 use riscv_pages::{
@@ -41,14 +42,34 @@ pub struct PerCpu {
 /// The base address of the per-CPU memory region.
 static PER_CPU_BASE: Once<SupervisorPageAddr> = Once::new();
 
+/// Error for SMP handling
+#[derive(Debug)]
+pub enum Error {
+    MapHartIdToCpuId(u64),
+    StartHart(sbi_rs::Error, u64),
+}
+
+impl fmt::Display for Error {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use Error::*;
+
+        match self {
+            MapHartIdToCpuId(hart_id) => write!(f, "Error mapping HART id ({}) to state", hart_id),
+            Self::StartHart(e, hart_id) => {
+                write!(f, "Error starting HART id ({}): {:?}", hart_id, e)
+            }
+        }
+    }
+}
+
 impl PerCpu {
     /// Initializes the `PerCpu` structures for each CPU, taking memory from `mem_map`. This (the
     /// boot CPU's) per-CPU area is initialized and loaded into TP as well.
-    pub fn init(boot_hart_id: u64, hyp_mem: &mut HypPageAlloc) {
+    pub fn init(boot_hart_id: u64, hyp_mem: &mut HypPageAlloc) -> Result<(), Error> {
         let cpu_info = CpuInfo::get();
         let boot_cpu = cpu_info
             .hart_id_to_cpu(boot_hart_id as u32)
-            .expect("Cannot find boot CPU ID");
+            .ok_or(Error::MapHartIdToCpuId(boot_hart_id))?;
 
         // Allocate memory for per-CPU structures.
         let pcpu_size = core::mem::size_of::<PerCpu>() as u64 * cpu_info.num_cpus() as u64;
@@ -63,7 +84,7 @@ impl PerCpu {
             let cpu_id = CpuId::new(i);
             // Boot CPU is special. Doesn't use the PCPU_BASE area as stack.
             let stack_pages = if cpu_id == boot_cpu {
-                Self::boot_cpu_stack()
+                Self::boot_cpu_stack()?
             } else {
                 // Unwrap okay. We allocated the area for `num_cpus() - 1` pages and we skip the
                 // bootstrap CPU above.
@@ -101,27 +122,30 @@ impl PerCpu {
         }
 
         // Initialize TP register and set this CPU online to be consistent with secondary CPUs.
-        Self::setup_this_cpu(boot_hart_id);
+        Self::setup_this_cpu(boot_hart_id)?;
         let me = Self::this_cpu();
         me.set_online();
+
+        Ok(())
     }
 
     /// Initializes the TP pointer to point to PerCpu data.
-    pub fn setup_this_cpu(hart_id: u64) {
+    pub fn setup_this_cpu(hart_id: u64) -> Result<(), Error> {
         let cpu_info = CpuInfo::get();
         let cpu_id = cpu_info
             .hart_id_to_cpu(hart_id as u32)
-            .expect("Cannot find CPU ID");
+            .ok_or(Error::MapHartIdToCpuId(hart_id))?;
 
         // Load TP with the address of our PerCpu struct.
         let tp = Self::ptr_for_cpu(cpu_id) as u64;
         unsafe {
             // Safe since we're the only users of TP.
             asm!("mv tp, {rs}", rs = in(reg) tp)
         };
+        Ok(())
     }
 
-    fn boot_cpu_stack() -> SequentialPages<InternalDirty> {
+    fn boot_cpu_stack() -> Result<SequentialPages<InternalDirty>, Error> {
         // Get the pages of the current stack as created by the linker.
         // Safe because `_stack_start` is created by the linker.
         let stack_startaddr = unsafe {
@@ -136,8 +160,10 @@ impl PerCpu {
         // Safe because the pages in this range are in the `HypervisorImage` memory region and are only
         // used for the boot cpu stack.
         unsafe {
-            SequentialPages::from_page_range(stack_startaddr, stack_endaddr, PageSize::Size4k)
-                .unwrap()
+            Ok(
+                SequentialPages::from_page_range(stack_startaddr, stack_endaddr, PageSize::Size4k)
+                    .unwrap(),
+            )
         }
     }
 
@@ -226,7 +252,7 @@ pub fn send_ipi(cpu: CpuId) {
 
 /// Boots secondary CPUs, using the HSM SBI call. Upon return, all secondary CPUs will have
 /// entered secondary_init().
-pub fn start_secondary_cpus() {
+pub fn start_secondary_cpus() -> Result<(), Error> {
     let cpu_info = CpuInfo::get();
     let boot_cpu = PerCpu::this_cpu().cpu_id();
     for i in 0..cpu_info.num_cpus() {
@@ -240,20 +266,19 @@ pub fn start_secondary_cpus() {
         let pcpu = unsafe { PerCpu::ptr_for_cpu(cpu_id).as_ref().unwrap() };
         let sp = pcpu.stack_top_addr();
 
+        let hart_id = cpu_info.cpu_to_hart_id(cpu_id).unwrap() as u64;
         // Safety: _secondary_start is guaranteed by the linker to be the code to start secondary
         // CPUs. pcpu will only be shared with one cpu.
         unsafe {
-            state::hart_start(
-                cpu_info.cpu_to_hart_id(cpu_id).unwrap() as u64,
-                (_secondary_start as *const fn()) as u64,
-                sp.bits(),
-            )
-            .expect("Failed to start CPU {i}");
+            state::hart_start(hart_id, (_secondary_start as *const fn()) as u64, sp.bits())
+                .map_err(|e| Error::StartHart(e, hart_id))?;
         }
 
         // Synchronize with the CPU coming online. TODO: Timeout?
         pcpu.online.wait();
     }
 
     println!("Brought online {} CPU(s)", cpu_info.num_cpus());
+
+    Ok(())
 }