Use DMA framework to create DMAble buffers

Cargo.lock

@@ -7029,6 +7029,7 @@ dependencies = [
  "netvsp",
  "nvme_driver",
  "nvme_resources",
+ "once_cell",
  "openhcl_attestation_protocol",
  "page_pool_alloc",
  "pal",

openhcl/page_pool_alloc/src/lib.rs

@@ -418,7 +418,7 @@ impl Drop for PagePoolHandle {
 /// [`PagePoolAllocatorSpawner::allocator`].
 ///
 /// This struct is considered the "owner" of the pool allowing for save/restore.
-#[derive(Inspect)]
+#[derive(Inspect, Clone)]
 pub struct PagePool {
     #[inspect(flatten)]
     inner: Arc<Mutex<PagePoolInner>>,

openhcl/underhill_core/Cargo.toml

@@ -29,6 +29,7 @@ vmotherboard = { workspace = true, features = [
 
 build_info.workspace = true
 cvm_tracing.workspace = true
+once_cell.workspace = true
 diag_proto.workspace = true
 diag_server.workspace = true
 debug_worker_defs = { workspace = true, optional = true }

openhcl/underhill_core/src/dispatch/mod.rs

@@ -7,6 +7,7 @@ mod pci_shutdown;
 pub mod vtl2_settings_worker;
 
 use self::vtl2_settings_worker::DeviceInterfaces;
+use crate::dma_manager::GlobalDmaManager;
 use crate::emuplat::netvsp::RuntimeSavedState;
 use crate::emuplat::EmuplatServicing;
 use crate::nvme_manager::NvmeManager;
@@ -108,6 +109,7 @@ pub trait LoadedVmNetworkSettings: Inspect {
         partition: Arc<UhPartition>,
         state_units: &StateUnits,
         vmbus_server: &Option<VmbusServerHandle>,
+        dma_manager: GlobalDmaManager,
     ) -> anyhow::Result<RuntimeSavedState>;
 
     /// Callback when network is removed externally.
@@ -179,6 +181,7 @@ pub(crate) struct LoadedVm {
     pub shared_vis_pool: Option<PagePool>,
     pub private_pool: Option<PagePool>,
     pub nvme_keep_alive: bool,
+    pub dma_manager: GlobalDmaManager,
 }
 
 pub struct LoadedVmState<T> {
@@ -736,6 +739,7 @@ impl LoadedVm {
                 self.partition.clone(),
                 &self.state_units,
                 &self.vmbus_server,
+                self.dma_manager.clone(),
             )
             .await?;
 

openhcl/underhill_core/src/dma_manager.rs

@@ -0,0 +1,131 @@
+//use std::sync::{Mutex, Weak};
+
+use memory_range::MemoryRange;
+use std::{collections::HashMap, sync::{Arc, Mutex}};
+use page_pool_alloc::PagePool;
+use user_driver::{memory::MemoryBlock, vfio::VfioDmaBuffer};
+use user_driver::lockmem::LockedMemorySpawner;
+
+#[derive(Clone)]
+pub struct GlobalDmaManager {
+    _physical_ranges: Vec<MemoryRange>,
+    _bounce_buffers_manager: Vec<MemoryRange>,
+    //clients: Mutex<Vec<Weak<DmaClient>>>,
+    //client_thresholds: Mutex<Vec<(Weak<DmaClient>, usize)>>,
+
+    page_pool: Option<PagePool>,
+    clients: HashMap<String, Arc<DmaClient>>,
+}
+
+impl GlobalDmaManager {
+    pub fn new(page_pool: Option<PagePool>) -> Self {
+        GlobalDmaManager {
+            _physical_ranges: Vec::new(),
+            _bounce_buffers_manager: Vec::new(),
+            //clients: Mutex::new(Vec::new()),
+            //client_thresholds: Mutex::new(Vec::new()),
+            page_pool: page_pool.clone(),
+            clients: HashMap::new(),
+        }
+    }
+
+    pub fn create_client(&mut self, pci_id: String) -> DmaClient {
+        let client = DmaClient {
+            dma_manager: Arc::new(Mutex::new(self.clone())),
+            dma_buffer_allocator: None,  // Valid now as `Option<Arc<dyn VfioDmaBuffer>>`
+        };
+
+        let arc_client = Arc::new(client);
+        self.clients.insert(pci_id, arc_client.clone()); // Store the cloned `Arc` in `clients`.
+
+        arc_client.as_ref().clone() // Return the `Arc<DmaClient>`.
+    }
+
+    pub fn get_client(&self, pci_id: &str) -> Option<Arc<DmaClient>> {
+        self.clients.get(pci_id).cloned()
+    }
+
+    pub fn get_dma_buffer_allocator(
+        &mut self,
+        device_name: String,
+    ) -> anyhow::Result<Arc<dyn VfioDmaBuffer>> {
+        self.page_pool
+            .as_ref()
+            .map(|p : &PagePool| -> anyhow::Result<Arc<dyn VfioDmaBuffer>> {
+                p.allocator(device_name)
+                    .map(|alloc| Arc::new(alloc) as Arc<dyn VfioDmaBuffer>)
+            })
+            .unwrap_or(Ok(Arc::new(LockedMemorySpawner)))
+    }
+}
+
+#[derive(Clone)]
+pub struct DmaClient {
+    dma_manager: Arc<Mutex<GlobalDmaManager>>,
+    dma_buffer_allocator: Option<Arc<dyn VfioDmaBuffer>>,
+}
+
+impl user_driver::DmaClient for DmaClient {
+    fn map_dma_ranges(
+        &self,
+        ranges: i32,
+    ) -> anyhow::Result<Vec<i32>> {
+        self.map_dma_ranges(ranges)
+    }
+
+    //fn get_dma_buffer_allocator(
+    //    &mut self,
+    //    device_name: String,
+    //) -> anyhow::Result<Arc<dyn VfioDmaBuffer>> {
+
+    //    let mut manager = self.dma_manager.lock().unwrap();
+
+    //    manager.get_dma_buffer_allocator(device_name).map(|allocator| {
+    //        self.dma_buffer_allocator = Some(allocator.clone());
+    //        allocator
+    //    })
+    //}
+
+    fn allocate_dma_buffer(
+        &self,
+        total_size: usize,
+    ) -> anyhow::Result<MemoryBlock> {
+        if self.dma_buffer_allocator.is_none() {
+            return Err(anyhow::anyhow!("DMA buffer allocator is not set"));
+        }
+
+        let allocator = self.dma_buffer_allocator.as_ref().unwrap();
+
+        allocator.create_dma_buffer(total_size)
+    }
+
+    fn attach_dma_buffer(&self, len: usize, base_pfn: u64) -> anyhow::Result<MemoryBlock>
+    {
+        let allocator = self.dma_buffer_allocator.as_ref().unwrap();
+        allocator.restore_dma_buffer(len, base_pfn)
+    }
+
+}
+
+impl DmaClient {
+    fn map_dma_ranges(
+        &self,
+        _ranges: i32,
+    ) -> anyhow::Result<Vec<i32>>
+    {
+        Ok(Vec::new())
+    }
+
+    pub fn get_dma_buffer_allocator(
+        &mut self,
+        device_name: String,
+    ) -> anyhow::Result<Arc<dyn VfioDmaBuffer>> {
+
+        let mut manager = self.dma_manager.lock().unwrap();
+
+        manager.get_dma_buffer_allocator(device_name).map(|allocator| {
+            self.dma_buffer_allocator = Some(allocator.clone());
+            allocator
+        })
+    }
+}

openhcl/underhill_core/src/emuplat/netvsp.rs

@@ -2,7 +2,9 @@
 // Licensed under the MIT License.
 
 use crate::dispatch::vtl2_settings_worker::wait_for_pci_path;
+use crate::dma_manager::GlobalDmaManager;
 use crate::vpci::HclVpciBusControl;
+use user_driver::DmaClient;
 use anyhow::Context;
 use async_trait::async_trait;
 use futures::lock::Mutex;
@@ -40,7 +42,6 @@ use uevent::UeventListener;
 use user_driver::vfio::vfio_set_device_reset_method;
 use user_driver::vfio::PciDeviceResetMethod;
 use user_driver::vfio::VfioDevice;
-use user_driver::vfio::VfioDmaBuffer;
 use vmcore::vm_task::VmTaskDriverSource;
 use vpci::bus_control::VpciBusControl;
 use vpci::bus_control::VpciBusEvent;
@@ -60,12 +61,12 @@ enum HclNetworkVfManagerMessage {
     PacketCapture(FailableRpc<PacketCaptureParams<Socket>, PacketCaptureParams<Socket>>),
 }
 
-async fn create_mana_device(
+async fn   create_mana_device(
     driver_source: &VmTaskDriverSource,
     pci_id: &str,
     vp_count: u32,
     max_sub_channels: u16,
-    dma_buffer: Arc<dyn VfioDmaBuffer>,
+    dma_client: Arc<dyn DmaClient>,
 ) -> anyhow::Result<ManaDevice<VfioDevice>> {
     // Disable FLR on vfio attach/detach; this allows faster system
     // startup/shutdown with the caveat that the device needs to be properly
@@ -89,7 +90,7 @@ async fn create_mana_device(
             pci_id,
             vp_count,
             max_sub_channels,
-            &dma_buffer,
+            dma_client.clone(),
         )
         .await
         {
@@ -118,9 +119,9 @@ async fn try_create_mana_device(
     pci_id: &str,
     vp_count: u32,
     max_sub_channels: u16,
-    dma_buffer: &Arc<dyn VfioDmaBuffer>,
+    dma_client: Arc<dyn DmaClient>,
 ) -> anyhow::Result<ManaDevice<VfioDevice>> {
-    let device = VfioDevice::new(driver_source, pci_id, dma_buffer.clone())
+    let device = VfioDevice::new(driver_source, pci_id, dma_client)
         .await
         .context("failed to open device")?;
 
@@ -200,9 +201,9 @@ struct HclNetworkVFManagerWorker {
     vtl2_bus_control: HclVpciBusControl,
     vtl2_pci_id: String,
     #[inspect(skip)]
-    dma_buffer: Arc<dyn VfioDmaBuffer>,
-    #[inspect(skip)]
     dma_mode: GuestDmaMode,
+    #[inspect(skip)]
+    dma_manager: GlobalDmaManager,
 }
 
 impl HclNetworkVFManagerWorker {
@@ -217,8 +218,8 @@ impl HclNetworkVFManagerWorker {
         endpoint_controls: Vec<DisconnectableEndpointControl>,
         vp_count: u32,
         max_sub_channels: u16,
-        dma_buffer: Arc<dyn VfioDmaBuffer>,
         dma_mode: GuestDmaMode,
+        dma_manager: GlobalDmaManager,
     ) -> (Self, mesh::Sender<HclNetworkVfManagerMessage>) {
         let (tx_to_worker, worker_rx) = mesh::channel();
         let vtl0_bus_control = if save_state.hidden_vtl0.lock().unwrap_or(false) {
@@ -247,8 +248,8 @@ impl HclNetworkVFManagerWorker {
                 vtl0_bus_control,
                 vtl2_bus_control,
                 vtl2_pci_id,
-                dma_buffer,
                 dma_mode,
+                dma_manager,
             },
             tx_to_worker,
         )
@@ -685,12 +686,13 @@ impl HclNetworkVFManagerWorker {
                     } else {
                         tracing::info!("VTL2 VF arrived");
                     }
+                    let pci_id = self.vtl2_pci_id.clone();
                     let device_bound = match create_mana_device(
                         &self.driver_source,
                         &self.vtl2_pci_id,
                         self.vp_count,
                         self.max_sub_channels,
-                        self.dma_buffer.clone(),
+                        Arc::new(self.dma_manager.create_client(pci_id))
                     )
                     .await
                     {
@@ -865,8 +867,9 @@ impl HclNetworkVFManager {
         vp_count: u32,
         max_sub_channels: u16,
         netvsp_state: &Option<Vec<SavedState>>,
-        dma_buffer: Arc<dyn VfioDmaBuffer>,
         dma_mode: GuestDmaMode,
+        dma_client: Arc<dyn DmaClient>,
+        dma_manager: GlobalDmaManager,
     ) -> anyhow::Result<(
         Self,
         Vec<HclNetworkVFManagerEndpointInfo>,
@@ -877,7 +880,7 @@ impl HclNetworkVFManager {
             &vtl2_pci_id,
             vp_count,
             max_sub_channels,
-            dma_buffer.clone(),
+            dma_client,
         )
         .await?;
         let (mut endpoints, endpoint_controls): (Vec<_>, Vec<_>) = (0..device.num_vports())
@@ -926,8 +929,8 @@ impl HclNetworkVFManager {
             endpoint_controls,
             vp_count,
             max_sub_channels,
-            dma_buffer,
             dma_mode,
+            dma_manager,
         );
 
         // Queue new endpoints.

openhcl/underhill_core/src/lib.rs

@@ -9,6 +9,7 @@
 
 mod diag;
 mod dispatch;
+mod dma_manager;
 mod emuplat;
 mod get_tracing;
 mod inspect_internal;