Implement Write/Read Cache (Phase 3)

markbase-core/src/vfs/cache.rs

@@ -0,0 +1,334 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex, RwLock};
+use std::time::{Duration, Instant};
+
+use bytes::Bytes;
+
+const DEFAULT_READ_CACHE_SIZE: usize = 64 * 1024 * 1024; // 64MB
+const DEFAULT_READ_CACHE_BLOCK_SIZE: usize = 64 * 1024; // 64KB blocks
+const DEFAULT_WRITE_CACHE_SIZE: usize = 32 * 1024 * 1024; // 32MB
+const DEFAULT_CACHE_TTL_SECS: u64 = 300; // 5 minutes
+
+#[derive(Clone)]
+pub struct CacheConfig {
+    pub read_cache_size: usize,
+    pub read_cache_block_size: usize,
+    pub write_cache_size: usize,
+    pub cache_ttl_secs: u64,
+}
+
+impl Default for CacheConfig {
+    fn default() -> Self {
+        Self {
+            read_cache_size: DEFAULT_READ_CACHE_SIZE,
+            read_cache_block_size: DEFAULT_READ_CACHE_BLOCK_SIZE,
+            write_cache_size: DEFAULT_WRITE_CACHE_SIZE,
+            cache_ttl_secs: DEFAULT_CACHE_TTL_SECS,
+        }
+    }
+}
+
+pub struct ReadCache {
+    config: CacheConfig,
+    blocks: RwLock<HashMap<u64, CachedBlock>>,
+    total_size: RwLock<usize>,
+    last_access: RwLock<Instant>,
+}
+
+struct CachedBlock {
+    data: Bytes,
+    timestamp: Instant,
+}
+
+impl ReadCache {
+    pub fn new(config: CacheConfig) -> Self {
+        Self {
+            config,
+            blocks: RwLock::new(HashMap::new()),
+            total_size: RwLock::new(0),
+            last_access: RwLock::new(Instant::now()),
+        }
+    }
+
+    pub fn get(&self, block_offset: u64) -> Option<Bytes> {
+        let blocks = self.blocks.read().unwrap();
+        if let Some(block) = blocks.get(&block_offset) {
+            if block.timestamp.elapsed() < Duration::from_secs(self.config.cache_ttl_secs) {
+                *self.last_access.write().unwrap() = Instant::now();
+                return Some(block.data.clone());
+            }
+        }
+        None
+    }
+
+    pub fn put(&self, block_offset: u64, data: Bytes) {
+        let mut blocks = self.blocks.write().unwrap();
+        let mut total_size = self.total_size.write().unwrap();
+
+        // Evict old blocks if cache is full
+        while *total_size + data.len() > self.config.read_cache_size && !blocks.is_empty() {
+            let oldest_key = blocks
+                .iter()
+                .min_by_key(|(_, b)| b.timestamp)
+                .map(|(k, _)| *k);
+            if let Some(key) = oldest_key {
+                if let Some(removed) = blocks.remove(&key) {
+                    *total_size -= removed.data.len();
+                }
+            }
+        }
+
+        // Insert new block
+        if *total_size + data.len() <= self.config.read_cache_size {
+            let data_clone = data.clone();
+            blocks.insert(
+                block_offset,
+                CachedBlock {
+                    data: data_clone,
+                    timestamp: Instant::now(),
+                },
+            );
+            *total_size += data.len();
+        }
+    }
+
+    pub fn invalidate(&self, block_offset: u64) {
+        let mut blocks = self.blocks.write().unwrap();
+        let mut total_size = self.total_size.write().unwrap();
+        if let Some(removed) = blocks.remove(&block_offset) {
+            *total_size -= removed.data.len();
+        }
+    }
+
+    pub fn invalidate_all(&self) {
+        let mut blocks = self.blocks.write().unwrap();
+        blocks.clear();
+        *self.total_size.write().unwrap() = 0;
+    }
+
+    pub fn stats(&self) -> CacheStats {
+        let blocks = self.blocks.read().unwrap();
+        CacheStats {
+            block_count: blocks.len(),
+            total_size: *self.total_size.read().unwrap(),
+            max_size: self.config.read_cache_size,
+        }
+    }
+
+    pub fn block_offset(offset: u64, block_size: usize) -> u64 {
+        offset / block_size as u64
+    }
+
+    pub fn block_range(offset: u64, len: u32, block_size: usize) -> Vec<u64> {
+        let start_block = Self::block_offset(offset, block_size);
+        let end_block = Self::block_offset(offset + len as u64 - 1, block_size);
+        (start_block..=end_block).collect()
+    }
+}
+
+pub struct WriteCache {
+    config: CacheConfig,
+    pending_writes: Mutex<HashMap<u64, Vec<u8>>>,
+    total_size: RwLock<usize>,
+}
+
+impl WriteCache {
+    pub fn new(config: CacheConfig) -> Self {
+        Self {
+            config,
+            pending_writes: Mutex::new(HashMap::new()),
+            total_size: RwLock::new(0),
+        }
+    }
+
+    pub fn put(&self, offset: u64, data: &[u8]) -> bool {
+        let mut pending = self.pending_writes.lock().unwrap();
+        let mut total_size = self.total_size.write().unwrap();
+
+        // Check if we have space
+        if *total_size + data.len() > self.config.write_cache_size {
+            return false;
+        }
+
+        // Merge with existing write at same offset
+        if let Some(existing) = pending.get_mut(&offset) {
+            // Overwrite overlapping range
+            for (i, byte) in data.iter().enumerate() {
+                if i < existing.len() {
+                    existing[i] = *byte;
+                } else {
+                    existing.push(*byte);
+                }
+            }
+        } else {
+            pending.insert(offset, data.to_vec());
+            *total_size += data.len();
+        }
+
+        true
+    }
+
+    pub fn get_pending(&self) -> Vec<(u64, Vec<u8>)> {
+        let pending = self.pending_writes.lock().unwrap();
+        pending.iter().map(|(k, v)| (*k, v.clone())).collect()
+    }
+
+    pub fn clear(&self) {
+        let mut pending = self.pending_writes.lock().unwrap();
+        pending.clear();
+        *self.total_size.write().unwrap() = 0;
+    }
+
+    pub fn stats(&self) -> CacheStats {
+        CacheStats {
+            block_count: self.pending_writes.lock().unwrap().len(),
+            total_size: *self.total_size.read().unwrap(),
+            max_size: self.config.write_cache_size,
+        }
+    }
+}
+
+pub struct CacheStats {
+    pub block_count: usize,
+    pub total_size: usize,
+    pub max_size: usize,
+}
+
+pub struct FileManager {
+    read_cache: Arc<ReadCache>,
+    write_cache: Arc<WriteCache>,
+}
+
+impl FileManager {
+    pub fn new(config: CacheConfig) -> Self {
+        Self {
+            read_cache: Arc::new(ReadCache::new(config.clone())),
+            write_cache: Arc::new(WriteCache::new(config)),
+        }
+    }
+
+    pub fn read_cache(&self) -> Arc<ReadCache> {
+        self.read_cache.clone()
+    }
+
+    pub fn write_cache(&self) -> Arc<WriteCache> {
+        self.write_cache.clone()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_read_cache_basic() {
+        let config = CacheConfig::default();
+        let cache = ReadCache::new(config);
+
+        let block_offset = 0;
+        let data1 = Bytes::from(vec![1, 2, 3, 4, 5]);
+
+        cache.put(block_offset, data1.clone());
+
+        let cached = cache.get(block_offset);
+        assert!(cached.is_some());
+        assert_eq!(cached.unwrap(), data1);
+    }
+
+    #[test]
+    fn test_read_cache_ttl_expiry() {
+        let config = CacheConfig {
+            cache_ttl_secs: 1,
+            ..Default::default()
+        };
+        let cache = ReadCache::new(config);
+
+        let block_offset = 0;
+        let data1 = Bytes::from(vec![1, 2, 3, 4, 5]);
+
+        cache.put(block_offset, data1.clone());
+
+        // Should be available immediately
+        assert!(cache.get(block_offset).is_some());
+
+        // Wait for TTL
+        std::thread::sleep(Duration::from_secs(2));
+
+        // Should be expired
+        assert!(cache.get(block_offset).is_none());
+    }
+
+    #[test]
+    fn test_read_cache_eviction() {
+        let config = CacheConfig {
+            read_cache_size: 100,
+            ..Default::default()
+        };
+        let cache = ReadCache::new(config);
+
+        // Fill cache
+        cache.put(0, Bytes::from(vec![0u8; 50]));
+        cache.put(1, Bytes::from(vec![1u8; 50]));
+
+        // Should evict oldest block
+        cache.put(2, Bytes::from(vec![2u8; 50]));
+
+        // Block 0 should be evicted
+        assert!(cache.get(0).is_none());
+        assert!(cache.get(1).is_some());
+        assert!(cache.get(2).is_some());
+    }
+
+    #[test]
+    fn test_write_cache_basic() {
+        let config = CacheConfig::default();
+        let cache = WriteCache::new(config);
+
+        let data1 = vec![1, 2, 3, 4, 5];
+        cache.put(0, &data1);
+
+        let pending = cache.get_pending();
+        assert_eq!(pending.len(), 1);
+        assert_eq!(pending[0], (0, data1));
+    }
+
+    #[test]
+    fn test_write_cache_merge() {
+        let config = CacheConfig::default();
+        let cache = WriteCache::new(config);
+
+        cache.put(0, &[1, 2, 3]);
+        cache.put(0, &[4, 5, 6, 7, 8]);
+
+        let pending = cache.get_pending();
+        assert_eq!(pending.len(), 1);
+        assert_eq!(pending[0].1, vec![4, 5, 6, 7, 8]);
+    }
+
+    #[test]
+    fn test_write_cache_full() {
+        let config = CacheConfig {
+            write_cache_size: 10,
+            ..Default::default()
+        };
+        let cache = WriteCache::new(config);
+
+        assert!(cache.put(0, &[1, 2, 3, 4, 5]));
+        assert!(cache.put(5, &[6, 7, 8, 9, 10]));
+        assert!(!cache.put(10, &[11, 12])); // Cache full
+    }
+
+    #[test]
+    fn test_block_range() {
+        let block_size = 1024;
+
+        let range = ReadCache::block_range(0, 512, block_size);
+        assert_eq!(range, vec![0]);
+
+        let range = ReadCache::block_range(0, 2048, block_size);
+        assert_eq!(range, vec![0, 1]);
+
+        let range = ReadCache::block_range(1024, 512, block_size);
+        assert_eq!(range, vec![1]);
+    }
+}