Refine

crates/vm-device/src/device/virtio/virtio_balloon_traditional.rs

@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::Mutex;
 
@@ -6,8 +7,9 @@ use vm_core::arch::irq::InterruptController;
 use vm_mm::allocator::MemoryContainer;
 use vm_mm::manager::MemoryAddressSpace;
 use vm_virtio::device::VirtioDevice;
-use vm_virtio::device::VirtqueueHandler;
-use vm_virtio::device::VirtqueueHandlerFn;
+use vm_virtio::device::transport::TransportContext;
+use vm_virtio::device::virtqueue::VirtqueueHandler;
+use vm_virtio::device::virtqueue::VirtqueueHandlerFn;
 use vm_virtio::result::Result;
 use vm_virtio::transport::VirtioDev;
 use vm_virtio::transport::mmio::VirtioMmioTransport;
@@ -20,12 +22,50 @@ use zerocopy::IntoBytes;
 const INFLATEQ_QUEUE_SIZE_MAX: u32 = 512;
 const DEFLATEQ_QUEUE_SIZE_MAX: u32 = 512;
 
-fn inflateq_handler<C, D>() -> VirtqueueHandlerFn<C, D> {
-    Box::new(|_mm, _dev, _desc_ring, _desc_id| todo!())
+fn inflateq_handler<C>() -> VirtqueueHandlerFn<C, VirtioBalloonTranditional<C>>
+where
+    C: MemoryContainer,
+{
+    Box::new(|mm, dev, desc_ring, desc_id| {
+        let desc = desc_ring.get(desc_id);
+        let len = desc.len;
+        assert!(len.is_multiple_of(4));
+
+        let array = desc.addr(mm).unwrap().as_ptr() as *const u32;
+
+        for i in 0..(len / 4) {
+            let pfn = unsafe { *array.add(i as usize) };
+            assert!(dev.device.balloon.insert(pfn));
+            let gpa = (pfn as u64) << 12;
+            let _hva = mm.gpa_to_hva(gpa).unwrap();
+            // TODO: mmap
+        }
+
+        len
+    })
 }
 
-fn deflateq_handler<C, D>() -> VirtqueueHandlerFn<C, D> {
-    Box::new(|_mm, _dev, _desc_ring, _desc_id| todo!())
+fn deflateq_handler<C>() -> VirtqueueHandlerFn<C, VirtioBalloonTranditional<C>>
+where
+    C: MemoryContainer,
+{
+    Box::new(|mm, dev, desc_ring, desc_id| {
+        let desc = desc_ring.get(desc_id);
+        let len = desc.len;
+        assert!(len.is_multiple_of(4));
+
+        let array = desc.addr(mm).unwrap().as_ptr() as *const u32;
+
+        for i in 0..(len / 4) {
+            let pfn = unsafe { *array.add(i as usize) };
+            assert!(dev.device.balloon.remove(&pfn));
+            let gpa = (pfn as u64) << 12;
+            let _hva = mm.gpa_to_hva(gpa).unwrap();
+            // TODO: mmap
+        }
+
+        len
+    })
 }
 
 pub struct VirtioBalloonTranditional<C>
@@ -36,6 +76,7 @@ where
     irq_chip: Arc<dyn InterruptController>,
     mm: Arc<MemoryAddressSpace<C>>,
     cfg: VirtioBalloonTranditionalConfig,
+    balloon: HashSet<u32>,
 }
 
 impl<C> VirtioBalloonTranditional<C>
@@ -52,8 +93,15 @@ where
             irq_chip,
             mm,
             cfg: VirtioBalloonTranditionalConfig::default(),
+            balloon: Default::default(),
         }
     }
+
+    pub fn set_num_pages(&mut self, num_pages: u32) {
+        self.cfg.num_pages = num_pages;
+
+        todo!("notify config generation changes");
+    }
 }
 
 impl<C> VirtioDevice<C> for VirtioBalloonTranditional<C>
@@ -112,15 +160,43 @@ where
         }
     }
 
-    fn read_config(&self, offset: usize, len: usize, buf: &mut [u8]) -> Result<()> {
-        buf.copy_from_slice(&self.cfg.as_bytes()[offset..offset + len]);
+    fn read_config(&self, offset: usize, buf: &mut [u8]) -> Result<()> {
+        buf.copy_from_slice(&self.cfg.as_bytes()[offset..offset + buf.len()]);
         Ok(())
     }
 
-    fn write_config(&mut self, offset: usize, len: usize, buf: &[u8]) -> Result<()> {
-        self.cfg.as_mut_bytes()[offset..len].copy_from_slice(buf);
+    fn write_config(&mut self, offset: usize, buf: &[u8]) -> Result<()> {
+        self.cfg.as_mut_bytes()[offset..offset + buf.len()].copy_from_slice(buf);
         Ok(())
     }
+
+    fn transport_context(&self) -> &dyn TransportContext {
+        todo!()
+    }
+
+    fn transport_context_mit(&mut self) -> &mut dyn TransportContext {
+        todo!()
+    }
+}
+
+pub trait VirtioBalloonApi {
+    fn update_num_pages(&mut self, num_pages: u32);
+}
+
+pub type VirtioBalloonDev<C> = VirtioDev<C, VirtioBalloonTranditional<C>>;
+
+impl<C> VirtioBalloonApi for VirtioBalloonDev<C>
+where
+    C: MemoryContainer,
+{
+    fn update_num_pages(&mut self, num_pages: u32) {
+        if self.device.cfg.num_pages == num_pages {
+            return;
+        }
+
+        self.device.cfg.num_pages = num_pages;
+        self.update_config_generation_and_notify();
+    }
 }
 
 pub type VirtioMmioBalloonDevice<C> = VirtioMmioTransport<C, VirtioBalloonTranditional<C>>;

crates/vm-device/src/device/virtio/virtio_blk.rs

@@ -7,15 +7,16 @@ use vm_mm::allocator::MemoryContainer;
 use vm_mm::manager::MemoryAddressSpace;
 use vm_pci::device::interrupt::legacy::InterruptPin;
 use vm_virtio::device::VirtioDevice;
-use vm_virtio::device::VirtqueueHandler;
-use vm_virtio::device::VirtqueueHandlerFn;
-use vm_virtio::device::blk::config::VirtioBlkConfig;
-use vm_virtio::device::blk::req::VirtioBlkReq;
-use vm_virtio::device::blk::req::VirtioBlkReqType;
-use vm_virtio::device::pci::VirtioPciDevice;
+use vm_virtio::device::transport::TransportContext;
+use vm_virtio::device::virtqueue::VirtqueueHandler;
+use vm_virtio::device::virtqueue::VirtqueueHandlerFn;
 use vm_virtio::result::Result;
 use vm_virtio::transport::VirtioDev;
 use vm_virtio::transport::mmio::VirtioMmioTransport;
+use vm_virtio::transport::pci::VirtioPciDevice;
+use vm_virtio::types::device::blk::config::VirtioBlkConfig;
+use vm_virtio::types::device::blk::req::VirtioBlkReq;
+use vm_virtio::types::device::blk::req::VirtioBlkReqType;
 use vm_virtio::types::device_features::VIRTIO_F_VERSION_1;
 use vm_virtio::types::device_id::DeviceId;
 use zerocopy::IntoBytes;
@@ -128,15 +129,23 @@ where
         })
     }
 
-    fn read_config(&self, offset: usize, len: usize, buf: &mut [u8]) -> Result<()> {
-        buf.copy_from_slice(&self.cfg.as_bytes()[offset..offset + len]);
+    fn read_config(&self, offset: usize, buf: &mut [u8]) -> Result<()> {
+        buf.copy_from_slice(&self.cfg.as_bytes()[offset..offset + buf.len()]);
         Ok(())
     }
 
-    fn write_config(&mut self, offset: usize, len: usize, buf: &[u8]) -> Result<()> {
-        self.cfg.as_mut_bytes()[offset..len].copy_from_slice(buf);
+    fn write_config(&mut self, offset: usize, buf: &[u8]) -> Result<()> {
+        self.cfg.as_mut_bytes()[offset..offset + buf.len()].copy_from_slice(buf);
         Ok(())
     }
+
+    fn transport_context(&self) -> &dyn TransportContext {
+        todo!()
+    }
+
+    fn transport_context_mit(&mut self) -> &mut dyn TransportContext {
+        todo!()
+    }
 }
 
 impl<C> VirtioPciDevice<C> for VirtioBlkDevice<C>

crates/vm-machine/src/device.rs

@@ -1,18 +1,21 @@
 use std::sync::Arc;
+use std::thread::sleep;
+use std::time::Duration;
 
 use vm_core::arch::irq::InterruptController;
 use vm_core::device::device_manager::DeviceManager;
 use vm_core::device::mmio::MmioRange;
 use vm_device::device::Device;
+use vm_device::device::virtio::virtio_balloon_traditional::VirtioBalloonApi;
 use vm_device::device::virtio::virtio_balloon_traditional::VirtioBalloonTranditional;
 use vm_device::device::virtio::virtio_balloon_traditional::VirtioMmioBalloonDevice;
 use vm_device::device::virtio::virtio_blk::VirtioBlkDevice;
 use vm_device::device::virtio::virtio_blk::VirtioMmioBlkDevice;
 use vm_mm::allocator::MemoryContainer;
 use vm_mm::manager::MemoryAddressSpace;
 use vm_pci::root_complex::mmio::PciRootComplexMmio;
-use vm_virtio::device::pci::VirtioPciDevice;
 use vm_virtio::transport::VirtioDev;
+use vm_virtio::transport::pci::VirtioPciDevice;
 
 use crate::error::Error;
 
@@ -88,13 +91,41 @@ impl InitDevice for DeviceManager {
             match device {
                 Device::GicV3 => (), // irq_chip is initialized already
                 Device::VirtioMmioBalloon => {
-                    // TODO: use mmio allocator
+                    let dev = VirtioDev::new(VirtioBalloonTranditional::new(
+                        3,
+                        irq_chip.clone(),
+                        mm.clone(),
+                    ));
+
+                    if false {
+                        std::thread::spawn({
+                            let dev = dev.clone();
+                            move || {
+                                let mut i = 0;
+                                loop {
+                                    i += 1;
+
+                                    {
+                                        sleep(Duration::from_secs(5));
+
+                                        let mut dev = dev.lock().unwrap();
+                                        dev.update_num_pages(i);
+                                    }
+
+                                    {
+                                        sleep(Duration::from_secs(5));
+
+                                        let mut dev = dev.lock().unwrap();
+                                        dev.update_num_pages(0);
+                                    }
+                                }
+                            }
+                        });
+                    }
+
+                    // TODO: use mmio allocator?
                     let virtio_mmio_balloon = VirtioMmioBalloonDevice::new(
-                        VirtioDev::new(VirtioBalloonTranditional::new(
-                            3,
-                            irq_chip.clone(),
-                            mm.clone(),
-                        )),
+                        dev,
                         MmioRange {
                             start: 0x0900_2000,
                             len: 0x1000,

crates/vm-virtio/src/device.rs

@@ -2,89 +2,14 @@ use std::sync::Arc;
 use std::sync::Mutex;
 
 use tokio::sync::Notify;
-use vm_core::arch::irq::InterruptController;
-use vm_mm::allocator::MemoryContainer;
-use vm_mm::manager::MemoryAddressSpace;
 
+use crate::device::transport::TransportContext;
+use crate::device::virtqueue::VirtqueueHandler;
 use crate::result::Result;
 use crate::transport::VirtioDev;
-use crate::types::interrupt_status::InterruptStatus;
-use crate::virtqueue::virtq_desc_table::VirtqDescTableRef;
 
-pub mod blk;
-pub mod pci;
-
-pub type VirtqueueHandlerFn<C, D> = Box<
-    dyn Fn(&MemoryAddressSpace<C>, &mut VirtioDev<C, D>, &VirtqDescTableRef, u16) -> u32
-        + Send
-        + Sync,
->;
-
-pub struct VirtqueueHandler<C, D> {
-    pub queue_sel: usize,
-    pub notifier: Arc<Notify>,
-    pub dev: Arc<Mutex<VirtioDev<C, D>>>,
-    pub mm: Arc<MemoryAddressSpace<C>>,
-    pub irq_chip: Arc<dyn InterruptController>,
-    pub irq_line: u32,
-    pub handle_desc: VirtqueueHandlerFn<C, D>,
-}
-
-impl<C, D> VirtqueueHandler<C, D>
-where
-    C: MemoryContainer,
-    D: VirtioDevice<C>,
-{
-    pub async fn run(self) {
-        let mm = self.mm.as_ref();
-
-        loop {
-            self.notifier.notified().await;
-
-            let mut dev = self.dev.lock().unwrap();
-
-            let mut updated = false;
-            loop {
-                let (desc_table, desc_id) = {
-                    // fetch desc from avail ring
-                    let q = dev.get_virtqueue_mut(self.queue_sel).unwrap();
-                    let avail_ring = q.avail_ring(mm).unwrap();
-
-                    if q.last_available_idx() == avail_ring.idx() {
-                        break;
-                    }
-
-                    let last_available_idx = q.last_available_idx();
-                    let desc_id = avail_ring.ring(last_available_idx);
-                    q.incr_last_available_idx();
-                    (q.desc_table_ref(mm).unwrap(), desc_id)
-                };
-
-                let len = (self.handle_desc)(mm, &mut dev, &desc_table, desc_id);
-
-                {
-                    // update used ring
-                    let q = dev.get_virtqueue_mut(self.queue_sel).unwrap();
-
-                    let mut used_ring = q.used_ring(mm).unwrap();
-                    let used_idx = used_ring.idx();
-                    let used_entry = used_ring.ring(used_idx);
-                    used_entry.id = desc_id as u32;
-                    used_entry.len = len;
-                    used_ring.incr_idx();
-                }
-
-                updated = true;
-            }
-
-            if updated {
-                let mut isr = dev.get_interrupt_status();
-                isr.insert(InterruptStatus::VIRTIO_MMIO_INT_VRING);
-                dev.update_interrupt_status(isr);
-            }
-        }
-    }
-}
+pub mod transport;
+pub mod virtqueue;
 
 pub trait VirtioDevice<C>: Sized + Send + Sync + 'static {
     const NAME: &str;
@@ -115,7 +40,11 @@ pub trait VirtioDevice<C>: Sized + Send + Sync + 'static {
         dev: Arc<Mutex<VirtioDev<C, Self>>>,
     ) -> Option<VirtqueueHandler<C, Self>>;
 
-    fn read_config(&self, offset: usize, len: usize, buf: &mut [u8]) -> Result<()>;
+    fn read_config(&self, offset: usize, buf: &mut [u8]) -> Result<()>;
+
+    fn write_config(&mut self, offset: usize, buf: &[u8]) -> Result<()>;
+
+    fn transport_context(&self) -> &dyn TransportContext;
 
-    fn write_config(&mut self, offset: usize, len: usize, buf: &[u8]) -> Result<()>;
+    fn transport_context_mit(&mut self) -> &mut dyn TransportContext;
 }