From 11e40750fe493d043163b3ccefec1ced4e8de39d Mon Sep 17 00:00:00 2001 From: rustmailer Date: Tue, 23 Jun 2026 23:23:03 +0800 Subject: [PATCH] feat(blob): add FilePool read cache and fix BucketCache TOCTOU --- crates/blob/src/account.rs | 14 +++ crates/blob/src/cache.rs | 168 +++++++++++++++++++++-------------- crates/blob/src/engine.rs | 5 +- crates/blob/src/file_pool.rs | 53 +++++++++++ crates/blob/src/lib.rs | 1 + crates/blob/src/segment.rs | 89 +++++++++++++++++++ 6 files changed, 259 insertions(+), 71 deletions(-) create mode 100644 crates/blob/src/file_pool.rs diff --git a/crates/blob/src/account.rs b/crates/blob/src/account.rs index 2af2717..b95794d 100644 --- a/crates/blob/src/account.rs +++ b/crates/blob/src/account.rs @@ -5,6 +5,7 @@ use std::sync::{Arc, Mutex, RwLock}; use crate::bucket::{self, BucketFile, IndexRecord}; use crate::error::{Error, Result}; +use crate::file_pool::FilePool; use crate::meta::{AccountMeta, SegmentStats}; use crate::segment::{self, SegmentReader, SegmentWriter}; use crate::types::Codec; @@ -16,6 +17,7 @@ pub struct AccountHandle { dir: PathBuf, inner: RwLock, pub(crate) write_mutex: Mutex<()>, + file_pool: FilePool, } impl AccountHandle { @@ -39,6 +41,7 @@ impl AccountHandle { dir, inner: RwLock::new(inner), write_mutex: Mutex::new(()), + file_pool: FilePool::new(8), })) } @@ -54,6 +57,7 @@ impl AccountHandle { dir, inner: RwLock::new(inner), write_mutex: Mutex::new(()), + file_pool: FilePool::new(8), })) } @@ -66,6 +70,16 @@ impl AccountHandle { pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, AccountInner> { self.inner.write().unwrap() } + + /// Get a cached file handle for a segment. + pub fn get_segment_file(&self, seg_id: u32, path: &Path) -> Result>> { + self.file_pool.get(seg_id, path) + } + + /// Invalidate cached file handles for a segment (after GC). + pub fn invalidate_file_cache(&self, seg_id: u32) { + self.file_pool.invalidate(seg_id); + } } // ── AccountInner ─────────────────────────────────────────────────────────── diff --git a/crates/blob/src/cache.rs b/crates/blob/src/cache.rs index d33cd55..b25830d 100644 --- a/crates/blob/src/cache.rs +++ b/crates/blob/src/cache.rs @@ -1,17 +1,22 @@ use std::collections::HashMap; +use std::path::Path; use std::sync::Mutex; use crate::bucket::{BucketFile, BucketIndex, IndexRecord}; use crate::error::Result; -/// Cache key: (account_name, bucket_id) type CacheKey = (String, u16); -/// LRU bucket cache. Thread-safe. +/// Thread-safe LRU bucket cache with single-lock interior. +/// Eliminates the TOCTOU race in the old two-Mutex design. pub struct BucketCache { + inner: Mutex, +} + +struct CacheInner { max_entries: usize, - entries: Mutex>, - index: Mutex>, + entries: Vec, + index: HashMap, } struct CacheEntry { @@ -22,27 +27,28 @@ struct CacheEntry { impl BucketCache { pub fn new(max_entries: usize) -> Self { Self { - max_entries: max_entries.max(1), - entries: Mutex::new(Vec::new()), - index: Mutex::new(HashMap::new()), + inner: Mutex::new(CacheInner { + max_entries: max_entries.max(1), + entries: Vec::new(), + index: HashMap::new(), + }), } } - /// Get or load a bucket index. Returns the sorted, deduplicated records for the bucket. + /// Get or load a bucket index. Eliminates TOCTOU via double-checked locking. pub fn get_or_load( &self, account: &str, bucket_id: u16, - account_dir: &std::path::Path, + account_dir: &Path, ) -> Result> { let key: CacheKey = (account.to_string(), bucket_id); // Check cache { - let index = self.index.lock().unwrap(); - if let Some(&pos) = index.get(&key) { - let entries = self.entries.lock().unwrap(); - return Ok(entries[pos].index.records.clone()); + let inner = self.inner.lock().unwrap(); + if let Some(&pos) = inner.index.get(&key) { + return Ok(inner.entries[pos].index.records.clone()); } } @@ -51,84 +57,74 @@ impl BucketCache { let bucket_index = bucket_file.load_index()?; let records = bucket_index.records.clone(); - // Insert into cache - self.insert(key, bucket_index); + // Insert with double-check (another thread might have beaten us) + { + let mut inner = self.inner.lock().unwrap(); + if let Some(&pos) = inner.index.get(&key) { + return Ok(inner.entries[pos].index.records.clone()); + } + // Evict if full + if inner.entries.len() >= inner.max_entries { + if let Some(evicted) = inner.entries.pop() { + inner.index.remove(&evicted.key); + } + } + // Insert at front + inner.entries.insert(0, CacheEntry { + key: key.clone(), + index: bucket_index, + }); + // Rebuild index + inner.index.clear(); + for i in 0..inner.entries.len() { + let key = inner.entries[i].key.clone(); + inner.index.insert(key, i); + } + } Ok(records) } - fn insert(&self, key: CacheKey, index: BucketIndex) { - let mut idx_map = self.index.lock().unwrap(); - let mut entries = self.entries.lock().unwrap(); - - // If already exists, update and move to front - if let Some(&pos) = idx_map.get(&key) { - entries[pos].index = index; - let entry = entries.remove(pos); - entries.insert(0, entry); - // Rebuild index - idx_map.clear(); - for (i, e) in entries.iter().enumerate() { - idx_map.insert(e.key.clone(), i); - } - return; - } - - // Evict if full - if entries.len() >= self.max_entries { - if let Some(evicted) = entries.pop() { - idx_map.remove(&evicted.key); - } - } - - // Insert at front (most recently used) - entries.insert(0, CacheEntry { key: key.clone(), index }); - // Rebuild index (positions shifted) - idx_map.clear(); - for (i, e) in entries.iter().enumerate() { - idx_map.insert(e.key.clone(), i); - } - } - - /// Insert or update a single record in a cached bucket. If bucket not cached, no-op. + /// Insert or update a single record in a cached bucket. pub fn update_record(&self, account: &str, bucket_id: u16, record: IndexRecord) { let key: CacheKey = (account.to_string(), bucket_id); - let mut idx_map = self.index.lock().unwrap(); + let mut inner = self.inner.lock().unwrap(); - if let Some(&pos) = idx_map.get(&key) { - let mut entries = self.entries.lock().unwrap(); - entries[pos].index.insert(record); + if let Some(&pos) = inner.index.get(&key) { + inner.entries[pos].index.insert(record); // Move to front - let entry = entries.remove(pos); - entries.insert(0, entry); + let entry = inner.entries.remove(pos); + inner.entries.insert(0, entry); // Rebuild index - idx_map.clear(); - for (i, e) in entries.iter().enumerate() { - idx_map.insert(e.key.clone(), i); + inner.index.clear(); + for i in 0..inner.entries.len() { + let entry_key = inner.entries[i].key.clone(); + inner.index.insert(entry_key, i); } } } - /// Invalidate a cached bucket. + /// Invalidate a cached bucket (after GC rewrites bucket files). pub fn invalidate(&self, account: &str, bucket_id: u16) { let key: CacheKey = (account.to_string(), bucket_id); - let mut idx_map = self.index.lock().unwrap(); - if let Some(&pos) = idx_map.get(&key) { - let mut entries = self.entries.lock().unwrap(); - entries.remove(pos); - idx_map.clear(); - for (i, e) in entries.iter().enumerate() { - idx_map.insert(e.key.clone(), i); + let mut inner = self.inner.lock().unwrap(); + + if let Some(&pos) = inner.index.get(&key) { + inner.entries.remove(pos); + inner.index.clear(); + for i in 0..inner.entries.len() { + let entry_key = inner.entries[i].key.clone(); + inner.index.insert(entry_key, i); } } } pub fn len(&self) -> usize { - self.entries.lock().unwrap().len() + self.inner.lock().unwrap().entries.len() } pub fn is_empty(&self) -> bool { - self.entries.lock().unwrap().is_empty() + self.inner.lock().unwrap().entries.is_empty() } } @@ -163,7 +159,6 @@ mod tests { let cache = BucketCache::new(10); let _ = cache.get_or_load("test", 0, dir.path()).unwrap(); - // Second call should hit cache let records = cache .get_or_load("test", 0, dir.path()) .unwrap(); @@ -186,4 +181,39 @@ mod tests { assert!(cache.len() <= 2); } + + #[test] + fn test_concurrent_get_or_load_no_deadlock() { + use std::sync::Arc; + use std::thread; + + let dir = TempDir::new().unwrap(); + crate::bucket::BucketFile::ensure_dir(dir.path()).unwrap(); + let bf = BucketFile::open(dir.path(), 0); + for i in 0..10u8 { + bf.append(&IndexRecord::new([i; 32], 1, i as u64 * 100, 50, 0)) + .unwrap(); + } + + let cache = Arc::new(BucketCache::new(10)); + let dir_path = dir.path().to_path_buf(); + + let mut handles = vec![]; + for _ in 0..4 { + let cache = cache.clone(); + let dir_path = dir_path.clone(); + handles.push(thread::spawn(move || { + for _ in 0..100 { + let records = cache + .get_or_load("test", 0, &dir_path) + .unwrap(); + assert_eq!(records.len(), 10); + } + })); + } + + for h in handles { + h.join().unwrap(); + } + } } diff --git a/crates/blob/src/engine.rs b/crates/blob/src/engine.rs index eaefb7c..18d13b8 100644 --- a/crates/blob/src/engine.rs +++ b/crates/blob/src/engine.rs @@ -203,8 +203,9 @@ impl Engine { return Err(Error::SegmentNotFound(record.segment_id)); } - let reader = SegmentReader::open(seg_path, record.segment_id)?; - let (entry, _) = reader.read_entry_at(record.offset)?; + let reader = SegmentReader::open(seg_path.clone(), record.segment_id)?; + let file = handle.get_segment_file(record.segment_id, &seg_path)?; + let (entry, _) = reader.read_entry_at_file(record.offset, &file)?; let value = compress::decompress(&entry.data, entry.codec, entry.raw_size as usize)?; diff --git a/crates/blob/src/file_pool.rs b/crates/blob/src/file_pool.rs new file mode 100644 index 0000000..229533b --- /dev/null +++ b/crates/blob/src/file_pool.rs @@ -0,0 +1,53 @@ +use std::collections::VecDeque; +use std::fs::File; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use crate::error::Result; + +/// Simple LRU pool of open file handles, keyed by segment_id. +/// Uses Arc> to allow safe concurrent reads from the same segment. +pub struct FilePool { + max_entries: usize, + entries: Mutex>)>>, +} + +impl FilePool { + pub fn new(max_entries: usize) -> Self { + Self { + max_entries: max_entries.max(1), + entries: Mutex::new(VecDeque::new()), + } + } + + /// Get an open File for the given segment. Reuses cached handle if available. + pub fn get(&self, seg_id: u32, path: &Path) -> Result>> { + let mut entries = self.entries.lock().unwrap(); + + // Check for existing entry + for (i, (id, _)) in entries.iter().enumerate() { + if *id == seg_id { + let (_, file) = entries.remove(i).unwrap(); + entries.push_front((seg_id, file.clone())); + return Ok(file); + } + } + + // Open new file + let file = Arc::new(Mutex::new(File::open(path)?)); + + // Evict oldest if full + if entries.len() >= self.max_entries { + entries.pop_back(); + } + + entries.push_front((seg_id, file.clone())); + Ok(file) + } + + /// Remove a cached file handle (e.g. after GC rewrites a segment). + pub fn invalidate(&self, seg_id: u32) { + let mut entries = self.entries.lock().unwrap(); + entries.retain(|(id, _)| *id != seg_id); + } +} diff --git a/crates/blob/src/lib.rs b/crates/blob/src/lib.rs index 11bb3b9..46feff4 100644 --- a/crates/blob/src/lib.rs +++ b/crates/blob/src/lib.rs @@ -5,6 +5,7 @@ pub mod checksum; pub mod compress; pub mod engine; pub mod error; +pub mod file_pool; pub mod gc; pub mod meta; pub mod recovery; diff --git a/crates/blob/src/segment.rs b/crates/blob/src/segment.rs index 1e3085f..0d7b162 100644 --- a/crates/blob/src/segment.rs +++ b/crates/blob/src/segment.rs @@ -1,6 +1,7 @@ use std::fs::{self, File, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; +use std::sync::Mutex; use crate::checksum; use crate::error::{Error, Result}; @@ -269,6 +270,94 @@ impl SegmentReader { )) } + /// Read a single entry at the given offset using a pre-opened File (via Mutex). + /// This avoids the per-read File::open cost for hot segments. + pub fn read_entry_at_file(&self, offset: u64, file: &Mutex) -> Result<(Entry, u64)> { + use std::io::{Read, Seek, SeekFrom}; + + let mut file = file.lock().unwrap(); + + file.seek(SeekFrom::Start(offset))?; + + // Read magic + let mut magic_buf = [0u8; 4]; + file.read_exact(&mut magic_buf)?; + let magic = u32::from_le_bytes(magic_buf); + if magic != ENTRY_MAGIC { + return Err(Error::CorruptEntry { + path: self.path.clone(), + offset, + reason: format!("bad magic: 0x{:08X}", magic), + }); + } + + // Read CRC32 + let mut crc_buf = [0u8; 4]; + file.read_exact(&mut crc_buf)?; + let stored_crc = u32::from_le_bytes(crc_buf); + + // Read flags, codec + let mut flags_buf = [0u8; 1]; + file.read_exact(&mut flags_buf)?; + let flags = flags_buf[0]; + + let mut codec_buf = [0u8; 1]; + file.read_exact(&mut codec_buf)?; + let codec = Codec::from_u8(codec_buf[0]).ok_or_else(|| Error::CorruptEntry { + path: self.path.clone(), + offset, + reason: format!("unknown codec: {}", codec_buf[0]), + })?; + + // Read key, raw_size, data_size + let mut key = [0u8; 32]; + file.read_exact(&mut key)?; + + let mut raw_size_buf = [0u8; 4]; + file.read_exact(&mut raw_size_buf)?; + let raw_size = u32::from_le_bytes(raw_size_buf); + + let mut data_size_buf = [0u8; 4]; + file.read_exact(&mut data_size_buf)?; + let data_size = u32::from_le_bytes(data_size_buf); + + // Read data + let mut data = vec![0u8; data_size as usize]; + file.read_exact(&mut data)?; + + // Verify CRC32 + let computed_crc = { + let mut hasher = crate::checksum::CrcWriter::new(); + hasher.update(&[flags]); + hasher.update(&[codec as u8]); + hasher.update(&key); + hasher.update(&raw_size.to_le_bytes()); + hasher.update(&data_size.to_le_bytes()); + hasher.update(&data); + hasher.finalize() + }; + + if stored_crc != computed_crc { + return Err(Error::CrcMismatch { + path: self.path.clone(), + offset, + }); + } + + let next_offset = offset + ENTRY_HEADER_SIZE as u64 + data_size as u64; + + Ok(( + Entry { + flags, + codec, + key, + raw_size, + data, + }, + next_offset, + )) + } + /// Read data portion of an entry (for pread-style reads when you already know offset + data_size). pub fn read_data(&self, offset: u64, data_size: u32) -> Result> { let mut file = File::open(&self.path)?;