feat(vm): handle shared files properly

Signed-off-by: Anhad Singh <andypythonappdeveloper@gmail.com>

Author: Andy-Python-Programmer

src/aero_kernel/src/fs/block/mod.rs

@@ -32,32 +32,46 @@ use crate::fs::{FileSystem, Result};
 
 use crate::fs::ext2::Ext2;
 use crate::mem::paging::*;
+use crate::mem::AddressSpace;
 use crate::utils::sync::Mutex;
 
 use super::cache::{Cache, CacheArc, CacheItem, Cacheable};
 use super::devfs::{alloc_device_marker, Device};
 use super::inode::INodeInterface;
 
-type PageCacheKey = (usize, usize); // (block device pointer, offset)
-type PageCacheItem = CacheArc<CacheItem<PageCacheKey, CachedPage>>;
+type PageCacheKey = (usize, usize); // (owner ptr, index)
+pub type PageCacheItem = CacheArc<CacheItem<PageCacheKey, CachedPage>>;
 
-struct CachedPage {
-    device: Weak<dyn CachedAccess>,
+struct DirtyMapping {
+    addr_space: AddressSpace,
+    addr: VirtAddr,
+}
+
+pub struct CachedPage {
+    owner: Weak<dyn CachedAccess>,
     offset: usize,
     page: PhysFrame,
     dirty: AtomicBool,
+    dirty_mappings: Mutex<Vec<DirtyMapping>>,
 }
 
 impl CachedPage {
-    fn new(device: Weak<dyn CachedAccess>, offset: usize) -> Self {
-        Self {
-            device,
+    fn new(owner: Weak<dyn CachedAccess>, offset: usize) -> Self {
+        let k = Self {
+            owner,
             offset,
             page: FRAME_ALLOCATOR
                 .allocate_frame()
                 .expect("page_cache: out of memory"),
             dirty: AtomicBool::new(false),
-        }
+            dirty_mappings: Mutex::new(Vec::new()),
+        };
+        // TODO: temporary hack. i mean this is fine but is there a cleaner way to do this. this is
+        // required since when the VM for the process umaps a page that contains a cached page, it
+        // will unmap this page which will decrease the refcnt to 0 and deallocate it.
+        get_vm_frames().unwrap()[k.page.start_address().as_u64() as usize / 4096usize]
+            .inc_ref_count();
+        k
     }
 
     fn data_mut(&self) -> &mut [MaybeUninit<u8>] {
@@ -72,10 +86,14 @@ impl CachedPage {
         unsafe { core::slice::from_raw_parts_mut(data_ptr, Size4KiB::SIZE as usize) }
     }
 
-    fn data_addr(&self) -> PhysAddr {
+    pub fn data_addr(&self) -> PhysAddr {
         self.page.start_address()
     }
 
+    pub fn page(&self) -> PhysFrame {
+        self.page
+    }
+
     fn make_key(device: &Weak<dyn CachedAccess>, offset: usize) -> PageCacheKey {
         (device.as_ptr().addr(), offset)
     }
@@ -85,26 +103,35 @@ impl CachedPage {
         self.dirty.load(Ordering::SeqCst)
     }
 
-    fn mark_dirty(&self) {
+    pub fn mark_dirty(&self) {
+        log::error!("marking dirty --------------------------------------");
         self.dirty.store(true, Ordering::SeqCst);
     }
 
     fn device(&self) -> Arc<dyn CachedAccess> {
-        self.device.upgrade().unwrap()
+        self.owner.upgrade().unwrap()
     }
 
     fn sync(&self) {
         if !self.is_dirty() {
             return;
         }
 
-        // Commit the changes made to the cache to the disk.
-        let disk = self.device();
-
+        // Commit the changes made to the cache to the owner.
+        let owner = self.device();
         let offset_bytes = self.offset * Size4KiB::SIZE as usize;
-        let sector = offset_bytes / disk.block_size();
+        owner.write_direct(offset_bytes, self.page);
+
+        for mut mapping in self.dirty_mappings.lock_irq().drain(..) {
+            let mut offset_table = mapping.addr_space.offset_page_table();
+            offset_table
+                .unmap(Page::<Size4KiB>::containing_address(mapping.addr))
+                .unwrap()
+                .1
+                .flush();
+        }
 
-        disk.write_dma(sector, self.data_addr(), Size4KiB::SIZE as usize);
+        self.dirty.store(false, Ordering::SeqCst);
     }
 }
 
@@ -116,12 +143,12 @@ impl Drop for CachedPage {
 
 impl Cacheable<PageCacheKey> for CachedPage {
     fn cache_key(&self) -> PageCacheKey {
-        Self::make_key(&self.device, self.offset)
+        Self::make_key(&self.owner, self.offset)
     }
 }
 
 lazy_static::lazy_static! {
-    static ref PAGE_CACHE: Arc<Cache<PageCacheKey, CachedPage>> = Cache::new();
+    pub(in crate::fs) static ref PAGE_CACHE: Arc<Cache<PageCacheKey, CachedPage>> = Cache::new();
 }
 
 impl Cache<PageCacheKey, CachedPage> {
@@ -145,16 +172,16 @@ impl Cache<PageCacheKey, CachedPage> {
         let device = device.upgrade().expect("page_cache: device dropped");
 
         let aligned_offset = align_down(offset as u64, Size4KiB::SIZE) as usize;
-        let sector = aligned_offset / device.block_size();
-
         device
-            .read_dma(sector, page.data_addr(), Size4KiB::SIZE as usize)
+            .read_direct(aligned_offset, page.page())
             .expect("page_cache: failed to read block");
 
         PAGE_CACHE.make_item_cached(page)
     }
 }
 
+// TODO: cache hit miss stats
+
 pub struct DirtyRef<T: Sized> {
     cache: PageCacheItem,
     ptr: *mut T,
@@ -202,9 +229,12 @@ pub trait BlockDeviceInterface: Send + Sync {
     fn write_block(&self, sector: usize, buf: &[u8]) -> Option<usize>;
 }
 
-pub trait CachedAccess: BlockDeviceInterface {
+pub trait CachedAccess: Send + Sync {
     fn sref(&self) -> Weak<dyn CachedAccess>;
 
+    fn read_direct(&self, offset: usize, dest: PhysFrame) -> Option<usize>;
+    fn write_direct(&self, offset: usize, src: PhysFrame) -> Option<usize>;
+
     fn read(&self, mut offset: usize, dest: &mut [MaybeUninit<u8>]) -> Option<usize> {
         let mut loc = 0;
 
@@ -236,6 +266,9 @@ pub trait CachedAccess: BlockDeviceInterface {
         let mut loc = 0;
 
         while loc < buffer.len() {
+            // TODO: If it is not found in the page cache, then, when the write perfectly falls on
+            // page size boundaries, the page is not even read from disk, but allocated and
+            // immediately marked dirty.
             let page = PAGE_CACHE.get_page(&self.sref(), offset);
 
             let page_offset = offset % Size4KiB::SIZE as usize;
@@ -318,6 +351,22 @@ impl CachedAccess for BlockDevice {
     fn sref(&self) -> Weak<dyn CachedAccess> {
         self.sref.clone()
     }
+
+    fn read_direct(&self, offset: usize, dest: PhysFrame) -> Option<usize> {
+        self.dev.read_dma(
+            offset / self.dev.block_size(),
+            dest.start_address(),
+            Size4KiB::SIZE as _,
+        )
+    }
+
+    fn write_direct(&self, offset: usize, src: PhysFrame) -> Option<usize> {
+        self.dev.write_dma(
+            offset / self.dev.block_size(),
+            src.start_address(),
+            Size4KiB::SIZE as _,
+        )
+    }
 }
 
 impl INodeInterface for BlockDevice {}

src/aero_kernel/src/fs/devfs.rs

@@ -31,7 +31,7 @@ use crate::mem::paging::*;
 use crate::rendy::RendyInfo;
 
 use super::cache::{DirCacheItem, INodeCacheItem};
-use super::inode::{INodeInterface, PollFlags, PollTable};
+use super::inode::{INodeInterface, MMapPage, PollFlags, PollTable};
 use super::ramfs::RamFs;
 use super::{FileSystem, FileSystemError, Result, MOUNT_MANAGER};
 
@@ -132,6 +132,10 @@ impl INodeInterface for DevINode {
     fn open(&self, handle: Arc<super::file_table::FileHandle>) -> Result<Option<DirCacheItem>> {
         self.0.inode().open(handle)
     }
+
+    fn mmap_v2(&self, offset: usize) -> Result<MMapPage> {
+        self.0.inode().mmap_v2(offset)
+    }
 }
 
 /// Implementation of dev filesystem. (See the module-level documentation for more
@@ -357,6 +361,26 @@ impl INodeInterface for DevFb {
             .expect("/dev/fb: terminal not initialized")
     }
 
+    fn mmap_v2(&self, offset: usize) -> Result<MMapPage> {
+        let rinfo = crate::rendy::get_rendy_info();
+
+        // Make sure we are in bounds.
+        if offset > rinfo.byte_len || offset + Size4KiB::SIZE as usize > rinfo.byte_len {
+            return Err(FileSystemError::NotSupported);
+        }
+
+        let mut lock = crate::rendy::DEBUG_RENDY.get().unwrap().lock_irq();
+        let fb = lock.get_framebuffer();
+
+        let fb_ptr = fb.as_ptr().cast::<u8>();
+        let fb_ptr = unsafe { fb_ptr.add(offset) };
+        let fb_phys_ptr = unsafe { fb_ptr.sub(crate::PHYSICAL_MEMORY_OFFSET.as_u64() as usize) };
+
+        Ok(MMapPage::Direct(PhysFrame::containing_address(unsafe {
+            PhysAddr::new_unchecked(fb_phys_ptr as u64)
+        })))
+    }
+
     fn ioctl(&self, command: usize, arg: usize) -> Result<usize> {
         match command {
             FBIOGET_VSCREENINFO => {

src/aero_kernel/src/fs/ext2/mod.rs

@@ -37,13 +37,13 @@ use crate::socket::SocketAddrRef;
 
 use self::group_desc::GroupDescriptors;
 
-use super::block::{self, BlockDevice, CachedAccess};
+use super::block::{self, BlockDevice, CachedAccess, PAGE_CACHE};
 
 use super::cache::{DirCacheItem, INodeCacheItem};
 use super::path::PathBuf;
 use super::{cache, FileSystemError, Path};
 
-use super::inode::{self, INodeInterface, Metadata, PollFlags, PollTable};
+use super::inode::{self, INodeInterface, MMapPage, Metadata, PollFlags, PollTable};
 use super::FileSystem;
 
 pub struct INode {
@@ -358,6 +358,20 @@ impl INode {
     }
 }
 
+impl CachedAccess for INode {
+    fn sref(&self) -> Weak<dyn CachedAccess> {
+        self.sref.clone()
+    }
+
+    fn read_direct(&self, offset: usize, dest: PhysFrame) -> Option<usize> {
+        INodeInterface::read_at(self, offset, dest.as_slice_mut()).ok()
+    }
+
+    fn write_direct(&self, offset: usize, src: PhysFrame) -> Option<usize> {
+        INodeInterface::write_at(self, offset, src.as_slice_mut()).ok()
+    }
+}
+
 impl INodeInterface for INode {
     fn weak_filesystem(&self) -> Option<Weak<dyn FileSystem>> {
         Some(self.fs.clone())
@@ -587,6 +601,14 @@ impl INodeInterface for INode {
         Ok(private_cp)
     }
 
+    // TODO: cleanup
+    fn mmap_v2(&self, offset: usize) -> super::Result<MMapPage> {
+        Ok(MMapPage::PageCache(PAGE_CACHE.get_page(
+            &(self.sref.clone() as Weak<dyn CachedAccess>),
+            offset,
+        )))
+    }
+
     fn listen(&self, backlog: usize) -> Result<(), SyscallError> {
         if let Some(proxy) = self.proxy.as_ref() {
             return proxy.listen(backlog);

src/aero_kernel/src/fs/inode.rs

@@ -33,6 +33,7 @@ use crate::socket::{SocketAddr, SocketAddrRef};
 use crate::userland::scheduler;
 use crate::utils::sync::{BMutex, Mutex, WaitQueue};
 
+use super::block::PageCacheItem;
 use super::cache::{Cacheable, CachedINode, DirCacheItem, INodeCacheItem};
 use super::devfs::DevINode;
 use super::file_table::FileHandle;
@@ -109,6 +110,11 @@ impl From<PollFlags> for PollEventFlags {
     }
 }
 
+pub enum MMapPage {
+    Direct(PhysFrame),
+    PageCache(PageCacheItem),
+}
+
 /// An inode describes a file. An inode structure holds metadata of the
 /// inode which includes its type, size, the number of links referring to it,
 /// and the list of blocks holding the file's content. For example device files,
@@ -234,6 +240,14 @@ pub trait INodeInterface: Send + Sync {
         Err(FileSystemError::NotSupported)
     }
 
+    fn mmap_v2(&self, _offset: usize) -> Result<MMapPage> {
+        log::error!(
+            "{} does not support mmap_v2!",
+            core::any::type_name_of_val(self),
+        );
+        Err(FileSystemError::NotSupported)
+    }
+
     // Socket operations:
     fn bind(&self, _address: SocketAddrRef, _length: usize) -> Result<()> {
         Err(FileSystemError::NotSocket)

src/aero_kernel/src/fs/ramfs.rs

@@ -33,7 +33,7 @@ use super::cache::{
 };
 use super::devfs::DevINode;
 use super::inode::{
-    DirEntry, FileContents, FileType, INodeInterface, Metadata, PollFlags, PollTable,
+    DirEntry, FileContents, FileType, INodeInterface, MMapPage, Metadata, PollFlags, PollTable,
 };
 use super::{FileSystem, FileSystemError, Result};
 
@@ -384,6 +384,21 @@ impl INodeInterface for LockedRamINode {
         }
     }
 
+    fn mmap_v2(&self, offset: usize) -> Result<MMapPage> {
+        let this = self.0.read();
+
+        match &this.contents {
+            FileContents::Device(dev) => {
+                let device = dev.clone();
+                drop(this);
+
+                device.mmap_v2(offset)
+            }
+
+            _ => todo!(),
+        }
+    }
+
     fn lookup(&self, dir: DirCacheItem, name: &str) -> Result<DirCacheItem> {
         let this = self.0.read();
         let child = this

src/aero_kernel/src/mem/paging/addr.rs

@@ -35,7 +35,7 @@ use bit_field::BitField;
 /// [`TryFrom`](https://doc.rust-lang.org/std/convert/trait.TryFrom.html) trait can be used for performing conversions
 /// between `u64` and `usize`.
 
-#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[repr(transparent)]
 pub struct VirtAddr(u64);