Lazily acquire I/O memory

kernel/comps/pci/src/capability/msix.rs

@@ -2,14 +2,14 @@
 
 //! MSI-X capability support.
 
-use alloc::{sync::Arc, vec::Vec};
+use alloc::vec::Vec;
 
-use ostd::{irq::IrqLine, mm::VmIoOnce};
+use ostd::{io::IoMem, irq::IrqLine, mm::VmIoOnce};
 
 use crate::{
     PciDeviceLocation,
     arch::{MSIX_DEFAULT_MSG_ADDR, construct_remappable_msix_address},
-    cfg_space::{Bar, Command, MemoryBar},
+    cfg_space::{BarAccess, Command},
     common_device::PciCommonDevice,
 };
 
@@ -21,9 +21,9 @@ pub struct CapabilityMsixData {
     table_size: u16,
     /// MSI-X table entry content:
     /// | Vector Control: u32 | Msg Data: u32 | Msg Upper Addr: u32 | Msg Addr: u32 |
-    table_bar: Arc<MemoryBar>,
+    table_bar: IoMem,
     /// Pending bits table.
-    pending_table_bar: Arc<MemoryBar>,
+    pending_table_bar: IoMem,
     table_offset: usize,
     pending_table_offset: usize,
     irqs: Vec<Option<IrqLine>>,
@@ -54,28 +54,30 @@ impl CapabilityMsixData {
         let table_bar;
         let pba_bar;
 
-        let bar_manager = dev.bar_manager();
+        let bar_manager = dev.bar_manager_mut();
         match bar_manager
-            .bar((pba_info & 0b111) as u8)
+            .bar_mut((pba_info & 0b111) as u8)
-            .cloned()
             .expect("MSIX cfg:pba BAR is none")
+            .acquire()
+            .expect("MSIX cfg:pba BAR is unavailable")
         {
-            Bar::Memory(memory) => {
+            BarAccess::Memory(io_mem) => {
-                pba_bar = memory;
+                pba_bar = io_mem.clone();
             }
-            Bar::Io(_) => {
+            BarAccess::Io => {
                 panic!("MSIX cfg:pba BAR is IO type")
             }
         };
         match bar_manager
-            .bar((table_info & 0b111) as u8)
+            .bar_mut((table_info & 0b111) as u8)
-            .cloned()
             .expect("MSIX cfg:table BAR is none")
+            .acquire()
+            .expect("MSIX cfg:table BAR is unavailable")
         {
-            Bar::Memory(memory) => {
+            BarAccess::Memory(io_mem) => {
-                table_bar = memory;
+                table_bar = io_mem.clone();
             }
-            Bar::Io(_) => {
+            BarAccess::Io => {
                 panic!("MSIX cfg:table BAR is IO type")
             }
         }
@@ -90,15 +92,12 @@ impl CapabilityMsixData {
         let message_upper_address = 0u32;
         for i in 0..table_size {
             table_bar
-                .io_mem()
                 .write_once((16 * i) as usize + table_offset, &message_address)
                 .unwrap();
             table_bar
-                .io_mem()
                 .write_once((16 * i + 4) as usize + table_offset, &message_upper_address)
                 .unwrap();
             table_bar
-                .io_mem()
                 .write_once((16 * i + 12) as usize + table_offset, &1_u32)
                 .unwrap();
         }
@@ -145,16 +144,13 @@ impl CapabilityMsixData {
             let address = construct_remappable_msix_address(remapping_index as u32);
 
             self.table_bar
-                .io_mem()
                 .write_once((16 * index) as usize + self.table_offset, &address)
                 .unwrap();
             self.table_bar
-                .io_mem()
                 .write_once((16 * index + 8) as usize + self.table_offset, &0)
                 .unwrap();
         } else {
             self.table_bar
-                .io_mem()
                 .write_once(
                     (16 * index + 8) as usize + self.table_offset,
                     &(irq.num() as u32),
@@ -165,7 +161,6 @@ impl CapabilityMsixData {
         let _old_irq = self.irqs[index as usize].replace(irq);
         // Enable this MSI-X vector.
         self.table_bar
-            .io_mem()
             .write_once((16 * index + 12) as usize + self.table_offset, &0_u32)
             .unwrap();
     }

kernel/comps/pci/src/cfg_space.rs

@@ -4,8 +4,6 @@
 //!
 //! Reference: <https://wiki.osdev.org/PCI>
 
-use alloc::sync::Arc;
-
 use bitflags::bitflags;
 use ostd::{
     Error, Result,
@@ -13,7 +11,6 @@ use ostd::{
     io::IoMem,
     mm::{PodOnce, VmIoOnce},
 };
-use spin::Once;
 
 use super::PciDeviceLocation;
 
@@ -238,15 +235,37 @@ bitflags! {
 }
 
 /// BAR space in PCI common config space.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub enum Bar {
     /// Memory BAR
-    Memory(Arc<MemoryBar>),
+    Memory(MemoryBar),
     /// I/O BAR
-    Io(Arc<IoBar>),
+    Io(IoBar),
 }
 
 impl Bar {
+    /// Acquires access to the BAR.
+    pub fn acquire(&mut self) -> Result<BarAccess> {
+        match self {
+            Self::Memory(mem_bar) => mem_bar.acquire().cloned().map(BarAccess::Memory),
+            // TODO: Implement the `acquire` operation based on `IoPortAllocator`.
+            Self::Io(_) => Ok(BarAccess::Io),
+        }
+    }
+
+    /// Returns access to the BAR.
+    ///
+    /// This method will return `None` if the access is not [`acquire`]d before.
+    ///
+    /// [`acquire`]: Self::acquire
+    pub fn access(&self) -> Option<BarAccess> {
+        match self {
+            Self::Memory(mem_bar) => mem_bar.access().cloned().map(BarAccess::Memory),
+            // TODO: Implement the `access` operation based on `IoPortAllocator`.
+            Self::Io(_) => Some(BarAccess::Io),
+        }
+    }
+
     pub(super) fn new(location: PciDeviceLocation, index: u8) -> Result<Self> {
         if index >= 6 {
             return Err(Error::InvalidArgs);
@@ -258,27 +277,40 @@ impl Bar {
         // Check the "Space Indicator" bit.
         let result = if raw & 1 == 0 {
             // Memory BAR
-            Self::Memory(Arc::new(MemoryBar::new(&location, index, raw)?))
+            Self::Memory(MemoryBar::new(&location, index, raw)?)
         } else {
             // I/O port BAR
-            Self::Io(Arc::new(IoBar::new(&location, index, raw)?))
+            Self::Io(IoBar::new(&location, index, raw)?)
         };
         Ok(result)
     }
+}
 
+/// Access to memory or I/O port BAR.
+#[derive(Debug, Clone)]
+pub enum BarAccess {
+    /// Memory BAR
+    Memory(IoMem),
+    /// I/O BAR
+    Io,
+}
+
+impl BarAccess {
     /// Reads a value of a specified type at a specified offset.
     pub fn read_once<T: PodOnce + PortRead>(&self, offset: usize) -> Result<T> {
         match self {
-            Bar::Memory(mem_bar) => mem_bar.io_mem().read_once(offset),
+            Self::Memory(io_mem) => io_mem.read_once(offset),
-            Bar::Io(io_bar) => io_bar.read(offset as u32),
+            // TODO: Implement the `read` operation based on `IoPortAllocator`.
+            Self::Io => Err(Error::IoError),
         }
     }
 
     /// Writes a value of a specified type at a specified offset.
     pub fn write_once<T: PodOnce + PortWrite>(&self, offset: usize, value: T) -> Result<()> {
         match self {
-            Bar::Memory(mem_bar) => mem_bar.io_mem().write_once(offset, &value),
+            Self::Memory(io_mem) => io_mem.write_once(offset, &value),
-            Bar::Io(io_bar) => io_bar.write(offset as u32, value),
+            // TODO: Implement the `write` operation based on `IoPortAllocator`.
+            Self::Io => Err(Error::IoError),
         }
     }
 }
@@ -290,7 +322,7 @@ pub struct MemoryBar {
     size: u64,
     prefetchable: bool,
     address_length: AddrLen,
-    io_memory: Once<IoMem>,
+    io_memory: Option<IoMem>,
 }
 
 impl MemoryBar {
@@ -315,13 +347,26 @@ impl MemoryBar {
         self.size
     }
 
-    /// Grants I/O memory access
+    /// Acquires access to the memory BAR.
-    pub fn io_mem(&self) -> &IoMem {
+    pub fn acquire(&mut self) -> Result<&IoMem> {
-        self.io_memory.call_once(|| {
+        if self.io_memory.is_none() {
             // Use the `Uncacheable` cache policy for PCI device BARs by default.
             // Device-specific drivers may remap with different cache policies if needed.
-            IoMem::acquire((self.base as usize)..((self.base + self.size) as usize)).unwrap()
+            self.io_memory = Some(IoMem::acquire(
-        })
+                (self.base as usize)..((self.base + self.size) as usize),
+            )?);
+        }
+
+        Ok(self.io_memory.as_ref().unwrap())
+    }
+
+    /// Returns access to the memory BAR.
+    ///
+    /// This method will return `None` if the access is not [`acquire`]d before.
+    ///
+    /// [`acquire`]: Self::acquire
+    pub fn access(&self) -> Option<&IoMem> {
+        self.io_memory.as_ref()
     }
 
     /// Creates a memory BAR structure.
@@ -417,7 +462,7 @@ impl MemoryBar {
             size,
             prefetchable,
             address_length,
-            io_memory: Once::new(),
+            io_memory: None,
         })
     }
 }
@@ -449,36 +494,6 @@ impl IoBar {
         self.size
     }
 
-    /// Reads from port
-    pub fn read<T: PortRead>(&self, offset: u32) -> Result<T> {
-        // Check alignment
-        if !(self.base + offset).is_multiple_of(size_of::<T>() as u32) {
-            return Err(Error::InvalidArgs);
-        }
-        // Check overflow
-        if self.size < size_of::<T>() as u32 || offset > self.size - size_of::<T>() as u32 {
-            return Err(Error::InvalidArgs);
-        }
-
-        // TODO: Implement the read operation based on `IoPortAllocator`.
-        Err(Error::IoError)
-    }
-
-    /// Writes to port
-    pub fn write<T: PortWrite>(&self, offset: u32, _value: T) -> Result<()> {
-        // Check alignment
-        if !(self.base + offset).is_multiple_of(size_of::<T>() as u32) {
-            return Err(Error::InvalidArgs);
-        }
-        // Check overflow
-        if size_of::<T>() as u32 > self.size || offset > self.size - size_of::<T>() as u32 {
-            return Err(Error::InvalidArgs);
-        }
-
-        // TODO: Implement the write operation based on `IoPortAllocator`.
-        Err(Error::IoError)
-    }
-
     /// Creates an I/O port BAR structure.
     fn new(location: &PciDeviceLocation, index: u8, raw: u32) -> Result<Self> {
         debug_assert_eq!(raw & 1, 1);

kernel/comps/pci/src/common_device.rs

@@ -37,16 +37,26 @@ impl PciCommonDevice {
         &self.location
     }
 
-    /// Returns the PCI Base Address Register (BAR) manager.
+    /// Returns a reference to the PCI Base Address Register (BAR) manager.
     pub fn bar_manager(&self) -> &BarManager {
         &self.bar_manager
     }
 
+    /// Returns a mutable reference to the PCI Base Address Register (BAR) manager.
+    pub fn bar_manager_mut(&mut self) -> &mut BarManager {
+        &mut self.bar_manager
+    }
+
     /// Returns the PCI capabilities.
     pub fn capabilities(&self) -> &Vec<Capability> {
         &self.capabilities
     }
 
+    /// Returns the PCI capabilities and a mutable reference to the BAR manager.
+    pub fn capabilities_and_bar_manager_mut(&mut self) -> (&Vec<Capability>, &mut BarManager) {
+        (&self.capabilities, &mut self.bar_manager)
+    }
+
     /// Returns the PCI device type.
     pub fn device_type(&self) -> PciDeviceType {
         self.header_type.device_type()
@@ -195,6 +205,11 @@ impl BarManager {
         self.bars[idx as usize].as_ref()
     }
 
+    /// Gains mutable access to the BAR space and returns `None` if that BAR is absent.
+    pub fn bar_mut(&mut self, idx: u8) -> Option<&mut Bar> {
+        self.bars[idx as usize].as_mut()
+    }
+
     /// Parses the BAR space by PCI device location.
     fn new(device_type: PciDeviceType, location: PciDeviceLocation) -> Self {
         let mut bars = [None, None, None, None, None, None];

kernel/comps/virtio/src/transport/mmio/device.rs

@@ -205,7 +205,7 @@ impl VirtioTransport for VirtioMmioTransport {
         Some(self.common_device.io_mem().slice(0x100..0x200))
     }
 
-    fn device_config_bar(&self) -> Option<(aster_pci::cfg_space::Bar, usize)> {
+    fn device_config_bar(&self) -> Option<(aster_pci::cfg_space::BarAccess, usize)> {
         None
     }
 

kernel/comps/virtio/src/transport/mod.rs

@@ -3,7 +3,7 @@
 use alloc::{boxed::Box, sync::Arc};
 use core::fmt::Debug;
 
-use aster_pci::cfg_space::Bar;
+use aster_pci::cfg_space::BarAccess;
 use aster_util::safe_ptr::SafePtr;
 use ostd::{
     Pod,
@@ -61,7 +61,7 @@ pub trait VirtioTransport: Sync + Send + Debug {
     fn device_config_mem(&self) -> Option<IoMem>;
 
     /// Get access to the device config BAR space.
-    fn device_config_bar(&self) -> Option<(Bar, usize)>;
+    fn device_config_bar(&self) -> Option<(BarAccess, usize)>;
 
     // ====================Virtqueue related APIs====================
 
@@ -113,13 +113,13 @@ pub trait VirtioTransport: Sync + Send + Debug {
 #[derive(Debug)]
 pub struct ConfigManager<T: Pod> {
     modern_space: Option<SafePtr<T, IoMem>>,
-    legacy_space: Option<(Bar, usize)>,
+    legacy_space: Option<(BarAccess, usize)>,
 }
 
 impl<T: Pod> ConfigManager<T> {
     pub(super) fn new(
         modern_space: Option<SafePtr<T, IoMem>>,
-        legacy_space: Option<(Bar, usize)>,
+        legacy_space: Option<(BarAccess, usize)>,
     ) -> Self {
         Self {
             modern_space,

kernel/comps/virtio/src/transport/pci/capability.rs

@@ -1,13 +1,10 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use alloc::sync::Arc;
-
 use aster_pci::{
-    capability::vendor::CapabilityVndrData,
+    capability::vendor::CapabilityVndrData, cfg_space::BarAccess, common_device::BarManager,
-    cfg_space::{Bar, MemoryBar},
-    common_device::BarManager,
 };
 use log::warn;
+use ostd::io::IoMem;
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
 #[repr(u8)]
@@ -26,11 +23,11 @@ pub struct VirtioPciCapabilityData {
     offset: u32,
     length: u32,
     option: Option<u32>,
-    memory_bar: Option<Arc<MemoryBar>>,
+    memory_bar: Option<IoMem>,
 }
 
 impl VirtioPciCapabilityData {
-    pub fn memory_bar(&self) -> Option<&Arc<MemoryBar>> {
+    pub fn memory_bar(&self) -> Option<&IoMem> {
         self.memory_bar.as_ref()
     }
 
@@ -50,7 +47,7 @@ impl VirtioPciCapabilityData {
         self.option
     }
 
-    pub(super) fn new(bar_manager: &BarManager, vendor_cap: CapabilityVndrData) -> Self {
+    pub(super) fn new(bar_manager: &mut BarManager, vendor_cap: CapabilityVndrData) -> Self {
         let cfg_type = vendor_cap.read8(3).unwrap();
         let cfg_type = match cfg_type {
             1 => VirtioPciCpabilityType::CommonCfg,
@@ -60,33 +57,40 @@ impl VirtioPciCapabilityData {
             5 => VirtioPciCpabilityType::PciCfg,
             _ => panic!("Unsupported virtio capability type:{:?}", cfg_type),
         };
-        let bar = vendor_cap.read8(4).unwrap();
+
-        let capability_length = vendor_cap.read8(2).unwrap();
         let offset = vendor_cap.read32(8).unwrap();
         let length = vendor_cap.read32(12).unwrap();
+
+        let capability_length = vendor_cap.read8(2).unwrap();
         let option = if capability_length > 0x10 {
             Some(vendor_cap.read32(16).unwrap())
         } else {
             None
         };
 
-        let mut memory_bar = None;
+        let bar = vendor_cap.read8(4).unwrap();
-        if let Some(bar) = bar_manager.bar(bar) {
+        let memory_bar = if let Some(bar) = bar_manager.bar_mut(bar) {
-            match bar {
+            match bar.acquire() {
-                Bar::Memory(memory) => {
+                Ok(BarAccess::Memory(io_mem)) => Some(io_mem),
-                    memory_bar = Some(memory);
+                Ok(BarAccess::Io) => {
+                    warn!("I/O BAR is not supported");
+                    None
                 }
-                Bar::Io(_) => {
+                Err(err) => {
-                    warn!("`Bar::Io` is not supported")
+                    warn!("BAR is not available: {:?}", err);
+                    None
                 }
             }
+        } else {
+            None
         };
+
         Self {
             cfg_type,
             offset,
             length,
             option,
-            memory_bar: memory_bar.cloned(),
+            memory_bar,
         }
     }
 }

kernel/comps/virtio/src/transport/pci/common_cfg.rs

@@ -31,9 +31,6 @@ pub struct VirtioPciCommonCfg {
 impl VirtioPciCommonCfg {
     pub(super) fn new(cap: &VirtioPciCapabilityData) -> SafePtr<Self, IoMem> {
         debug_assert!(cap.typ() == VirtioPciCpabilityType::CommonCfg);
-        SafePtr::new(
+        SafePtr::new(cap.memory_bar().unwrap().clone(), cap.offset() as usize)
-            cap.memory_bar().unwrap().io_mem().clone(),
-            cap.offset() as usize,
-        )
     }
 }

kernel/comps/virtio/src/transport/pci/device.rs

@@ -4,7 +4,7 @@ use alloc::{boxed::Box, sync::Arc};
 use core::fmt::Debug;
 
 use aster_pci::{
-    PciDeviceId, bus::PciDevice, capability::CapabilityData, cfg_space::Bar,
+    PciDeviceId, bus::PciDevice, capability::CapabilityData, cfg_space::BarAccess,
     common_device::PciCommonDevice,
 };
 use aster_util::{field_ptr, safe_ptr::SafePtr};
@@ -134,13 +134,12 @@ impl VirtioTransport for VirtioPciModernTransport {
             .device_cfg
             .memory_bar()
             .unwrap()
-            .io_mem()
             .slice(offset..offset + length);
 
         Some(io_mem)
     }
 
-    fn device_config_bar(&self) -> Option<(Bar, usize)> {
+    fn device_config_bar(&self) -> Option<(BarAccess, usize)> {
         None
     }
 
@@ -266,7 +265,7 @@ impl VirtioTransport for VirtioPciModernTransport {
 impl VirtioPciModernTransport {
     #[expect(clippy::result_large_err)]
     pub(super) fn new(
-        common_device: PciCommonDevice,
+        mut common_device: PciCommonDevice,
     ) -> Result<Self, (BusProbeError, PciCommonDevice)> {
         let device_id = common_device.device_id().device_id;
         let device_type_value = if device_id <= 0x1040 {
@@ -289,10 +288,11 @@ impl VirtioPciModernTransport {
         let mut notify = None;
         let mut common_cfg = None;
         let mut device_cfg = None;
-        for cap in common_device.capabilities().iter() {
+        let (caps, bar_manager) = common_device.capabilities_and_bar_manager_mut();
+        for cap in caps {
             match cap.capability_data() {
                 CapabilityData::Vndr(vendor) => {
-                    let data = VirtioPciCapabilityData::new(common_device.bar_manager(), *vendor);
+                    let data = VirtioPciCapabilityData::new(bar_manager, *vendor);
                     match data.typ() {
                         VirtioPciCpabilityType::CommonCfg => {
                             common_cfg = Some(VirtioPciCommonCfg::new(&data));
@@ -301,7 +301,7 @@ impl VirtioPciModernTransport {
                             notify = Some(VirtioPciNotify {
                                 offset_multiplier: data.option_value().unwrap(),
                                 offset: data.offset(),
-                                io_memory: data.memory_bar().unwrap().io_mem().clone(),
+                                io_memory: data.memory_bar().unwrap().clone(),
                             });
                         }
                         VirtioPciCpabilityType::IsrCfg => {}

kernel/comps/virtio/src/transport/pci/legacy.rs

@@ -3,7 +3,7 @@
 use alloc::{boxed::Box, sync::Arc};
 use core::fmt::Debug;
 
-use aster_pci::{capability::CapabilityData, cfg_space::Bar, common_device::PciCommonDevice};
+use aster_pci::{capability::CapabilityData, cfg_space::BarAccess, common_device::PciCommonDevice};
 use aster_util::safe_ptr::SafePtr;
 use log::{info, warn};
 use ostd::{
@@ -63,7 +63,7 @@ const DEVICE_CONFIG_OFFSET_WITH_MSIX: usize = 0x18;
 pub struct VirtioPciLegacyTransport {
     device_type: VirtioDeviceType,
     common_device: PciCommonDevice,
-    config_bar: Bar,
+    config_bar: BarAccess,
     num_queues: u16,
     msix_manager: VirtioMsixManager,
 }
@@ -73,7 +73,7 @@ impl VirtioPciLegacyTransport {
 
     #[expect(clippy::result_large_err)]
     pub(super) fn new(
-        common_device: PciCommonDevice,
+        mut common_device: PciCommonDevice,
     ) -> Result<Self, (BusProbeError, PciCommonDevice)> {
         let device_type = match common_device.device_id().device_id {
             0x1000 => VirtioDeviceType::Network,
@@ -93,7 +93,12 @@ impl VirtioPciLegacyTransport {
         };
         info!("[Virtio]: Found device:{:?}", device_type);
 
-        let config_bar = common_device.bar_manager().bar(0).cloned().unwrap();
+        let config_bar = common_device
+            .bar_manager_mut()
+            .bar_mut(0)
+            .unwrap()
+            .acquire()
+            .unwrap();
 
         let mut num_queues = 0u16;
         while num_queues < u16::MAX {
@@ -202,7 +207,7 @@ impl VirtioTransport for VirtioPciLegacyTransport {
         None
     }
 
-    fn device_config_bar(&self) -> Option<(Bar, usize)> {
+    fn device_config_bar(&self) -> Option<(BarAccess, usize)> {
         let bar = self.config_bar.clone();
         let base = if self.msix_manager.is_enabled() {
             DEVICE_CONFIG_OFFSET_WITH_MSIX