yggdrasil/kernel/driver/fs/ext2/src/inode/cache.rs

use core::ops::{AsyncFnOnce, Deref, DerefMut};

use alloc::sync::Arc;
use bytemuck::Zeroable;
use libk::{
    error::Error,
    task::sync::AsyncMutex,
    time::real_time,
    vfs::{Metadata, NodeRef},
};
use libk_util::{lru_hash_table::LruCache, sync::spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::io::{FileMode, FileType};

use crate::{data::InodeMode, dir::DirectoryNode, file::RegularNode, Ext2Fs, Inode};

pub struct InodeHolder {
    inode: Inode,
    dirty: bool,
}

pub struct InodeCache {
    fs: Arc<Ext2Fs>,
    synchronous: bool,
    cache: AsyncMutex<LruCache<u32, Arc<IrqSafeRwLock<InodeHolder>>>>,
}

pub struct InodeAccess {
    inode_cache: Arc<InodeCache>,
    ino: u32,
}

impl InodeAccess {
    pub fn new(inode_cache: Arc<InodeCache>, ino: u32) -> Self {
        Self { inode_cache, ino }
    }

    pub fn ino(&self) -> u32 {
        self.ino
    }

    pub fn cache(&self) -> &Arc<InodeCache> {
        &self.inode_cache
    }

    pub async fn map<T, F: FnOnce(&Inode) -> Result<T, Error>>(
        &self,
        mapper: F,
    ) -> Result<T, Error> {
        let inode = self.inode_cache.entry(self.ino).await?;
        let lock = inode.read();
        mapper(&lock.inode)
    }

    pub async fn map_mut<T, F: FnOnce(&mut Inode) -> Result<T, Error>>(
        &self,
        mapper: F,
    ) -> Result<T, Error> {
        let inode = self.inode_cache.entry(self.ino).await?;
        let mut lock = inode.write();
        let result = mapper(&mut lock.inode);
        self.inode_cache.put(self.ino, &mut lock).await?;
        result
    }

    pub async fn amap<T, F: AsyncFnOnce(&Inode) -> Result<T, Error>>(
        &self,
        mapper: F,
    ) -> Result<T, Error> {
        let inode = self.inode_cache.entry(self.ino).await?;
        let lock = inode.read();
        mapper(&lock.inode).await
    }

    pub async fn amap_mut<T, F: AsyncFnOnce(&mut Inode) -> Result<T, Error>>(
        &self,
        mapper: F,
    ) -> Result<T, Error> {
        let inode = self.inode_cache.entry(self.ino).await?;
        let mut lock = inode.write();
        let result = mapper(&mut lock.inode).await;
        self.inode_cache.put(self.ino, &mut lock).await?;
        result
    }

    pub async fn map_blocks<T, F: Fn(&Inode, usize, &[u8]) -> Result<Option<T>, Error>>(
        &self,
        mapper: F,
    ) -> Result<Option<T>, Error> {
        self.amap(async |inode| {
            let block_count = inode.blocks(&self.inode_cache.fs);

            for index in 0..block_count {
                let result = self
                    .inode_cache
                    .fs
                    .with_inode_block(inode, index as u32, |block| mapper(inode, index, block))
                    .await?;

                if let Some(result) = result {
                    return Ok(Some(result));
                }
            }

            Ok(None)
        })
        .await
    }

    pub async fn map_blocks_mut<T, F: Fn(&Inode, usize, &mut [u8]) -> Result<Option<T>, Error>>(
        &self,
        mapper: F,
    ) -> Result<Option<T>, Error> {
        self.amap(async |inode| {
            let block_count = inode.blocks(&self.inode_cache.fs);

            for index in 0..block_count {
                let result = self
                    .inode_cache
                    .fs
                    .with_inode_block_mut(inode, index as u32, 0, |block| {
                        mapper(inode, index, block)
                    })
                    .await?;

                if let Some(result) = result {
                    return Ok(Some(result));
                }
            }

            Ok(None)
        })
        .await
    }

    pub async fn metadata(&self) -> Result<Metadata, Error> {
        self.map(|inode| Ok(inode.metadata(&self.inode_cache.fs, self.ino)))
            .await
    }

    pub async fn size(&self) -> Result<u64, Error> {
        self.map(|inode| Ok(inode.size(&self.inode_cache.fs))).await
    }

    pub async fn update_metadata(&self, metadata: &Metadata) -> Result<(), Error> {
        let uid = metadata
            .uid
            .bits()
            .try_into()
            .map_err(|_| Error::InvalidArgument)?;
        let gid = metadata
            .gid
            .bits()
            .try_into()
            .map_err(|_| Error::InvalidArgument)?;

        self.map_mut(|inode| {
            inode.mtime = metadata.mtime as _;
            inode.atime = metadata.mtime as _;
            inode.ctime = metadata.ctime as _;

            inode.mode.update_permissions(metadata.mode);
            inode.uid = uid;
            inode.gid = gid;

            Ok(())
        })
        .await
    }

    async fn populate_inode(
        fs: &Arc<Ext2Fs>,
        ty: FileType,
        mode: FileMode,
        parent_ino: Option<u32>,
        ino: u32,
    ) -> Result<NodeRef, Error> {
        log::info!("ext2: allocated inode #{ino}");

        let now = real_time().seconds as u32;

        let mut imode = InodeMode::default_for_type(ty);
        imode.update_permissions(mode);

        fs.write_inode(
            ino,
            &Inode {
                ctime: now,
                mtime: now,
                atime: now,
                mode: imode,
                ..Inode::zeroed()
            },
        )
        .await?;

        let this = InodeAccess::new(fs.inode_cache.get().clone(), ino);
        let fs = fs.clone();
        let node = match ty {
            FileType::Directory => DirectoryNode::create(fs, this, parent_ino).await?,
            FileType::File => RegularNode::new(fs, this),
            FileType::Symlink => todo!(),
            _ => return Err(Error::NotImplemented),
        };

        Ok(node)
    }

    pub async fn allocate(
        fs: &Arc<Ext2Fs>,
        ty: FileType,
        mode: FileMode,
        parent_ino: Option<u32>,
    ) -> Result<NodeRef, Error> {
        if parent_ino.is_none() && ty != FileType::Directory {
            log::warn!("ext2: cannot allocate non-directory inode without a parent");
            return Err(Error::InvalidOperation);
        }
        if !matches!(ty, FileType::Symlink | FileType::Directory | FileType::File) {
            log::warn!("ext2: cannot allocate inode of type {ty:?}");
            return Err(Error::InvalidOperation);
        }

        let ino = fs.allocate_inode(ty == FileType::Directory).await?;

        match Self::populate_inode(fs, ty, mode, parent_ino, ino).await {
            Ok(node) => Ok(node),
            Err(error) => {
                log::warn!("ext2: couldn't set up inode #{ino}: {error:?}");
                // TODO free the inode and flush it from the cache
                Err(error)
            }
        }
    }
}

impl InodeCache {
    pub fn with_capacity(fs: Arc<Ext2Fs>, bucket_capacity: usize, synchronous: bool) -> Self {
        Self {
            fs,
            synchronous,
            cache: AsyncMutex::new(LruCache::with_capacity(bucket_capacity, 4)),
        }
    }

    async fn evict_inode(
        &self,
        ino: u32,
        inode: Arc<IrqSafeRwLock<InodeHolder>>,
    ) -> Result<(), Error> {
        let inode = inode.read();
        if inode.dirty {
            assert!(!self.synchronous);
            log::debug!("Flush dirty inode {ino}");
            self.fs.write_inode(ino, &inode.inode).await?;
        }
        Ok(())
    }

    async fn fetch_inode(&self, ino: u32) -> Result<Arc<IrqSafeRwLock<InodeHolder>>, Error> {
        let inode = self.fs.read_inode(ino).await?;
        Ok(Arc::new(IrqSafeRwLock::new(InodeHolder {
            inode,
            dirty: false,
        })))
    }

    async fn entry(&self, ino: u32) -> Result<Arc<IrqSafeRwLock<InodeHolder>>, Error> {
        if ino < 1 || ino > self.fs.total_inodes {
            return Err(Error::InvalidFile);
        }

        let mut lock = self.cache.lock().await;
        let (value, evicted) = lock
            .try_get_or_insert_with_async(ino, || self.fetch_inode(ino))
            .await?;
        let value = value.clone();

        if let Some((ino, holder)) = evicted {
            if let Err(error) = self.evict_inode(ino, holder).await {
                log::error!("ext2: inode flush error: ino={ino}, error={error:?}");
            }
        }

        Ok(value)
    }

    async fn put(&self, ino: u32, holder: &mut InodeHolder) -> Result<(), Error> {
        if self.synchronous {
            // Immediately write-back
            self.fs.write_inode(ino, &holder.inode).await
        } else {
            // Mark dirty
            holder.dirty = true;
            Ok(())
        }
    }

    pub async fn flush(&self) -> Result<(), Error> {
        let mut last_error = None;

        let mut lock = self.cache.lock().await;
        while let Some((ino, inode)) = lock.pop_entry() {
            if let Err(error) = self.evict_inode(ino, inode).await {
                log::error!("ext2: flush inode #{ino} error: {error:?}");
                last_error = Some(error);
            }
        }

        match last_error {
            None => Ok(()),
            Some(error) => Err(error),
        }
    }
}

impl Deref for InodeHolder {
    type Target = Inode;

    fn deref(&self) -> &Self::Target {
        &self.inode
    }
}

impl DerefMut for InodeHolder {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.dirty = true;
        &mut self.inode
    }
}