331 lines
9.9 KiB
Rust

#![allow(clippy::missing_transmute_annotations)]
use core::{
mem::MaybeUninit,
ops::{Deref, DerefMut},
};
use alloc::sync::Arc;
use libk_util::{lru_hash_table::LruCache, sync::spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::error::Error;
use crate::{dma::DmaBuffer, task::sync::AsyncMutex};
use super::BlockDevice;
pub struct CachedSegment {
data: DmaBuffer<[u8]>,
dirty: bool,
}
pub struct UncachedCache {
device: Arc<dyn BlockDevice>,
block_size: usize,
}
pub enum DeviceMapper {
Uncached(UncachedCache),
Cached(BlockCache),
}
pub struct BlockCache {
device: Arc<dyn BlockDevice>,
block_size: usize,
segment_size: usize,
cache: AsyncMutex<LruCache<u64, Arc<IrqSafeRwLock<CachedSegment>>>>,
}
impl DeviceMapper {
pub fn cached_with_capacity(
device: Arc<dyn BlockDevice>,
block_size: usize,
segment_size: usize,
bucket_capacity: usize,
bucket_count: usize,
filesystem: &str,
) -> Result<DeviceMapper, Error> {
BlockCache::with_capacity(
device,
block_size,
segment_size,
bucket_capacity,
bucket_count,
filesystem,
)
.map(DeviceMapper::Cached)
}
pub fn uncached(
device: Arc<dyn BlockDevice>,
block_size: usize,
filesystem: &str,
) -> Result<DeviceMapper, Error> {
if block_size % device.block_size() != 0 {
log::error!(
"Couldn't create block mapper for {filesystem}: \
cache block size is not a multiple of device block size"
);
return Err(Error::InvalidArgument);
}
let uncache = UncachedCache { device, block_size };
Ok(Self::Uncached(uncache))
}
pub fn device(&self) -> &Arc<dyn BlockDevice> {
match self {
Self::Uncached(uncache) => uncache.device(),
Self::Cached(cache) => cache.device(),
}
}
pub async fn try_with<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
&self,
pos: u64,
mapper: F,
) -> Result<T, Error> {
match self {
Self::Uncached(uncache) => uncache.try_with(pos, mapper).await,
Self::Cached(cache) => cache.try_with(pos, mapper).await,
}
}
pub async fn try_with_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
&self,
pos: u64,
size: usize,
mapper: F,
) -> Result<T, Error> {
match self {
Self::Uncached(uncache) => uncache.try_with_mut(pos, size, mapper).await,
Self::Cached(cache) => cache.try_with_mut(pos, size, mapper).await,
}
}
pub async fn flush(&self) -> Result<(), Error> {
match self {
Self::Uncached(_) => Ok(()),
Self::Cached(cache) => {
cache.flush().await;
Ok(())
}
}
}
}
impl UncachedCache {
pub fn device(&self) -> &Arc<dyn BlockDevice> {
&self.device
}
pub async fn try_with<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
&self,
pos: u64,
mapper: F,
) -> Result<T, Error> {
if pos % self.block_size as u64 != 0 {
log::warn!(
"uncached: position {pos} is not a multiple of block size {}",
self.block_size
);
return Err(Error::InvalidArgument);
}
let mut buffer = self.device.allocate_buffer(self.block_size)?;
self.device
.read_aligned(pos, buffer.slice_mut(0..self.block_size))
.await?;
let result = mapper(unsafe { MaybeUninit::slice_assume_init_ref(&buffer[..]) })?;
Ok(result)
}
pub async fn try_with_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
&self,
pos: u64,
size: usize,
mapper: F,
) -> Result<T, Error> {
if pos % self.block_size as u64 != 0 {
log::warn!(
"uncached: position {pos} is not a multiple of block size {}",
self.block_size
);
return Err(Error::InvalidArgument);
}
let mut buffer = self.device.allocate_buffer(self.block_size)?;
// No need to read a block only to then fully rewrite it
if size != self.block_size {
self.device
.read_aligned(pos, buffer.slice_mut(0..self.block_size))
.await?;
}
let mut buffer = unsafe { DmaBuffer::assume_init_slice(buffer) };
let result = mapper(&mut buffer[..])?;
self.device
.write_aligned(pos, buffer.slice(0..self.block_size))
.await?;
Ok(result)
}
}
impl BlockCache {
pub fn with_capacity(
device: Arc<dyn BlockDevice>,
block_size: usize,
segment_size: usize,
bucket_capacity: usize,
bucket_count: usize,
filesystem: &str,
) -> Result<Self, Error> {
if block_size % device.block_size() != 0 {
log::error!(
"Couldn't create block cache for {filesystem}: \
cache block size is not a multiple of device block size"
);
return Err(Error::InvalidArgument);
}
if segment_size % block_size != 0 {
log::error!(
"Couldn't create block cache for {filesystem}: \
segment size is not a multiple of block size"
);
return Err(Error::InvalidArgument);
}
if segment_size < block_size {
log::error!(
"Couldn't create block cache for {filesystem}: \
segment size is smaller than block size"
);
return Err(Error::InvalidArgument);
}
log::info!("New block cache: block: {block_size}B, segment: {segment_size}B, geometry: {bucket_capacity}x{bucket_count}");
log::info!(
"Worst-case memory usage: {}K",
(segment_size * bucket_capacity * bucket_count) / 1024
);
Ok(BlockCache {
device,
block_size,
segment_size,
cache: AsyncMutex::new(LruCache::with_capacity(bucket_capacity, bucket_count)),
})
}
pub fn device(&self) -> &Arc<dyn BlockDevice> {
&self.device
}
async fn evict_block(&self, segment_position: u64, block: Arc<IrqSafeRwLock<CachedSegment>>) {
let read = block.read();
if read.dirty {
assert_eq!(segment_position % self.segment_size as u64, 0);
if let Err(err) = self
.device
.write_aligned(segment_position, read.data.slice(0..self.segment_size))
.await
{
log::error!("Disk error: flushing block {}: {:?}", segment_position, err);
}
}
}
async fn fetch_block(
&self,
segment_position: u64,
) -> Result<Arc<IrqSafeRwLock<CachedSegment>>, Error> {
let mut buffer = self.device.allocate_buffer(self.segment_size)?;
self.device
.read_aligned(segment_position, buffer.slice_mut(0..self.segment_size))
.await?;
let data = unsafe { DmaBuffer::assume_init_slice(buffer) };
Ok(Arc::new(IrqSafeRwLock::new(CachedSegment {
data,
dirty: false,
})))
}
async fn entry(
&self,
segment_position: u64,
) -> Result<Arc<IrqSafeRwLock<CachedSegment>>, Error> {
assert_eq!(segment_position % self.segment_size as u64, 0);
let mut lock = self.cache.lock().await;
let (value, evicted) = lock
.try_get_or_insert_with_async(segment_position, || self.fetch_block(segment_position))
.await?;
if let Some((segment_position, block)) = evicted {
self.evict_block(segment_position, block).await;
}
Ok(value.clone())
}
pub async fn try_with<T, F: FnOnce(&[u8]) -> Result<T, Error>>(
&self,
block_position: u64,
mapper: F,
) -> Result<T, Error> {
if block_position % self.block_size as u64 != 0 {
log::warn!(
"mapper: position {block_position} is not a multiple of block size {}",
self.block_size
);
return Err(Error::InvalidArgument);
}
let segment_position = block_position & !(self.segment_size as u64 - 1);
let segment_offset = (block_position & (self.segment_size as u64 - 1)) as usize;
let block = self.entry(segment_position).await?;
let result = mapper(&block.read()[segment_offset..segment_offset + self.block_size])?;
Ok(result)
}
pub async fn try_with_mut<T, F: FnOnce(&mut [u8]) -> Result<T, Error>>(
&self,
block_position: u64,
_size: usize,
mapper: F,
) -> Result<T, Error> {
if block_position % self.block_size as u64 != 0 {
log::warn!(
"mapper: position {block_position} is not a multiple of block size {}",
self.block_size
);
return Err(Error::InvalidArgument);
}
let segment_position = block_position & !(self.segment_size as u64 - 1);
let segment_offset = (block_position & (self.segment_size as u64 - 1)) as usize;
let block = self.entry(segment_position).await?;
let mut block = block.write();
block.dirty = true;
let result = mapper(&mut block[segment_offset..segment_offset + self.block_size])?;
Ok(result)
}
pub async fn flush(&self) {
for (pos, block) in self.cache.lock().await.flush() {
self.evict_block(pos, block).await;
}
}
}
impl CachedSegment {
pub fn set_dirty(&mut self) {
self.dirty = true;
}
}
impl Deref for CachedSegment {
type Target = DmaBuffer<[u8]>;
fn deref(&self) -> &Self::Target {
&self.data
}
}
impl DerefMut for CachedSegment {
fn deref_mut(&mut self) -> &mut Self::Target {
self.dirty = true;
&mut self.data
}
}