refactor(blob): replace global RwLock with per-account Arc<AccountHandle>

This commit is contained in:
rustmailer
2026-06-23 23:13:39 +08:00
parent 495c91b7ae
commit 4e12ae5457
3 changed files with 169 additions and 154 deletions

View File

@@ -1,7 +1,7 @@
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::sync::{Arc, Mutex, RwLock};
use crate::bucket::{self, BucketFile, IndexRecord};
use crate::error::{Error, Result};
@@ -9,26 +9,78 @@ use crate::meta::{AccountMeta, SegmentStats};
use crate::segment::{self, SegmentReader, SegmentWriter};
use crate::types::Codec;
pub struct Account {
pub id: String,
// ── AccountHandle ──────────────────────────────────────────────────────────
pub struct AccountHandle {
id: String,
dir: PathBuf,
meta: AccountMeta,
active_writer: SegmentWriter,
readers: HashMap<u32, SegmentReader>,
write_lock: Mutex<()>,
inner: RwLock<AccountInner>,
pub write_mutex: Mutex<()>,
}
impl Account {
impl AccountHandle {
pub fn id(&self) -> &str {
&self.id
}
pub fn dir(&self) -> &Path {
&self.dir
}
/// Open an existing account.
pub fn open(store_root: &Path, account_id: &str) -> Result<Self> {
pub fn open(store_root: &Path, account_id: &str) -> Result<Arc<Self>> {
let dir = store_root.join("accounts").join(account_id);
if !dir.exists() {
return Err(Error::AccountNotFound(account_id.to_string()));
}
let inner = AccountInner::open(&dir, account_id)?;
Ok(Arc::new(Self {
id: account_id.to_string(),
dir,
inner: RwLock::new(inner),
write_mutex: Mutex::new(()),
}))
}
let meta = AccountMeta::load(&dir)?;
/// Create a new account.
pub fn create(store_root: &Path, account_id: &str) -> Result<Arc<Self>> {
let dir = store_root.join("accounts").join(account_id);
if dir.exists() {
return Err(Error::AccountAlreadyExists(account_id.to_string()));
}
let inner = AccountInner::create(&dir, account_id)?;
Ok(Arc::new(Self {
id: account_id.to_string(),
dir,
inner: RwLock::new(inner),
write_mutex: Mutex::new(()),
}))
}
/// Lock the inner state for reading.
pub fn read(&self) -> std::sync::RwLockReadGuard<'_, AccountInner> {
self.inner.read().unwrap()
}
/// Lock the inner state for writing.
pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, AccountInner> {
self.inner.write().unwrap()
}
}
// ── AccountInner ───────────────────────────────────────────────────────────
pub struct AccountInner {
dir: PathBuf,
pub meta: AccountMeta,
active_writer: SegmentWriter,
readers: HashMap<u32, SegmentReader>,
}
impl AccountInner {
fn open(dir: &Path, _account_id: &str) -> Result<Self> {
let meta = AccountMeta::load(dir)?;
// Open active writer
let seg_path = dir
.join("segments")
.join(segment::segment_filename(meta.active_segment_id));
@@ -39,7 +91,6 @@ impl Account {
SegmentWriter::create(seg_path, meta.active_segment_id)?
};
// Open readers for existing sealed segments
let mut readers = HashMap::new();
for (&seg_id, stats) in &meta.segments {
if stats.sealed {
@@ -53,24 +104,16 @@ impl Account {
}
Ok(Self {
id: account_id.to_string(),
dir,
dir: dir.to_path_buf(),
meta,
active_writer,
readers,
write_lock: Mutex::new(()),
})
}
/// Create a new account.
pub fn create(store_root: &Path, account_id: &str) -> Result<Self> {
let dir = store_root.join("accounts").join(account_id);
if dir.exists() {
return Err(Error::AccountAlreadyExists(account_id.to_string()));
}
fn create(dir: &Path, account_id: &str) -> Result<Self> {
fs::create_dir_all(dir.join("segments"))?;
BucketFile::ensure_dir(&dir)?;
BucketFile::ensure_dir(dir)?;
let meta = AccountMeta::new(account_id.to_string(), 1);
@@ -79,33 +122,21 @@ impl Account {
.join(segment::segment_filename(1));
let active_writer = SegmentWriter::create(seg_path, 1)?;
meta.save(&dir)?;
meta.save(dir)?;
Ok(Self {
id: account_id.to_string(),
dir,
dir: dir.to_path_buf(),
meta,
active_writer,
readers: HashMap::new(),
write_lock: Mutex::new(()),
})
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn meta(&self) -> &AccountMeta {
&self.meta
}
/// Lock for writing. Returns a guard.
pub fn lock_write(&self) -> std::sync::MutexGuard<'_, ()> {
self.write_lock.lock().unwrap()
}
/// Mark the segment as indexed up to the given offset and persist meta.
/// Called after append_index to enable incremental recovery.
pub fn mark_indexed(&mut self, segment_id: u32, indexed_up_to_offset: u64) -> Result<()> {
if let Some(stats) = self.meta.segments.get_mut(&segment_id) {
if indexed_up_to_offset > stats.indexed_up_to_offset {
@@ -115,8 +146,7 @@ impl Account {
self.meta.save(&self.dir)
}
/// Append an entry without fsync. Returns (segment_id, offset, data_size).
/// Caller must hold the write lock and should call `flush_active()` after.
/// Append an entry without fsync.
pub fn append_entry(
&mut self,
key: [u8; 32],
@@ -159,7 +189,7 @@ impl Account {
self.meta.save(&self.dir)
}
/// Write an entry with fsync. Convenience wrapper for single writes.
/// Write an entry with fsync.
pub fn write_entry(
&mut self,
key: [u8; 32],
@@ -181,7 +211,6 @@ impl Account {
.or_insert_with(|| SegmentStats::new(old_id));
old_stats.sealed = true;
// Open reader for the old segment
let seg_path = self
.dir
.join("segments")
@@ -189,7 +218,6 @@ impl Account {
self.readers
.insert(old_id, SegmentReader::open(seg_path, old_id)?);
// Create new segment
let new_id = old_id + 1;
self.meta.active_segment_id = new_id;
let new_path = self
@@ -213,14 +241,14 @@ impl Account {
}
}
/// Append index record to the appropriate bucket file. No fsync.
/// Append index record to the appropriate bucket file.
pub fn append_index(&self, record: &IndexRecord) -> Result<()> {
let bucket_id = bucket::bucket_id(&record.key);
let bf = BucketFile::open(&self.dir, bucket_id);
bf.append(record)
}
/// Return list of sealed segment IDs for GC consideration.
/// Return list of sealed segment IDs.
pub fn sealed_segments(&self) -> Vec<u32> {
self.meta
.segments

View File

@@ -1,24 +1,23 @@
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::RwLock;
use std::sync::{Arc, RwLock};
use crate::account::Account;
use crate::account::AccountHandle;
use crate::bucket::{self, IndexRecord};
use crate::cache::BucketCache;
use crate::compress;
use crate::error::{Error, Result};
use crate::gc::{self, GcStats};
use crate::meta::GlobalMeta;
use crate::recovery;
use crate::segment::{self, SegmentReader};
use crate::segment::SegmentReader;
use crate::types::{Codec, Config, ENTRY_HEADER_SIZE};
pub struct Engine {
root: PathBuf,
config: Config,
cache: BucketCache,
accounts: RwLock<HashMap<String, Account>>,
accounts: RwLock<HashMap<String, Arc<AccountHandle>>>,
}
#[derive(Debug, Clone)]
@@ -31,7 +30,6 @@ pub struct AccountStats {
}
impl Engine {
/// Open or create the store at `path`. Runs recovery on startup.
pub fn open(path: &Path, config: Config) -> Result<Self> {
config.validate()?;
fs::create_dir_all(path)?;
@@ -42,7 +40,6 @@ impl Engine {
let cache = BucketCache::new(config.lru_bucket_count);
// Discover accounts on disk
let accounts_dir = path.join("accounts");
let mut accounts = HashMap::new();
@@ -52,15 +49,13 @@ impl Engine {
if entry.file_type()?.is_dir() {
let account_name = entry.file_name().to_string_lossy().into_owned();
// Clean up temp files from interrupted GC
let _ = recovery::cleanup_temp_files(&entry.path());
let _ = crate::recovery::cleanup_temp_files(&entry.path());
// Run recovery
match recovery::recover_account(&entry.path()) {
match crate::recovery::recover_account(&entry.path()) {
Ok(_meta) => {
match Account::open(path, &account_name) {
Ok(account) => {
accounts.insert(account_name, account);
match AccountHandle::open(path, &account_name) {
Ok(handle) => {
accounts.insert(account_name, handle);
}
Err(e) => {
tracing::warn!(
@@ -83,7 +78,6 @@ impl Engine {
}
}
// Update global account list
global.accounts = accounts.keys().cloned().collect();
global.save(path)?;
@@ -95,15 +89,15 @@ impl Engine {
})
}
// ── Account management ──────────────────────────────────────────────────
// ── Account management ──────────────────────────────────────────────
pub fn create_account(&self, account_id: &str) -> Result<()> {
let mut accounts = self.accounts.write().unwrap();
if accounts.contains_key(account_id) {
return Err(Error::AccountAlreadyExists(account_id.to_string()));
}
let account = Account::create(&self.root, account_id)?;
accounts.insert(account_id.to_string(), account);
let handle = AccountHandle::create(&self.root, account_id)?;
accounts.insert(account_id.to_string(), handle);
let mut global = GlobalMeta::load(&self.root)?;
global.accounts = accounts.keys().cloned().collect();
@@ -114,12 +108,12 @@ impl Engine {
pub fn delete_account(&self, account_id: &str) -> Result<()> {
let mut accounts = self.accounts.write().unwrap();
let account = accounts
let handle = accounts
.remove(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
let account_dir = account.dir().to_path_buf();
drop(account);
let account_dir = handle.dir().to_path_buf();
drop(handle);
fs::remove_dir_all(&account_dir)?;
let mut global = GlobalMeta::load(&self.root)?;
@@ -134,7 +128,7 @@ impl Engine {
accounts.keys().cloned().collect()
}
// ── Read / Write / Delete ───────────────────────────────────────────────
// ── Read / Write / Delete ───────────────────────────────────────────
pub fn write(
&self,
@@ -147,28 +141,28 @@ impl Engine {
return Err(Error::ValueTooLarge { size: value.len() });
}
let mut accounts = self.accounts.write().unwrap();
let account = accounts
.get_mut(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
let handle = {
let accounts = self.accounts.read().unwrap();
accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?
.clone()
};
// Acquire per-account write lock
{
let _lock = account.lock_write();
}
let _write_lock = handle.write_mutex.lock().unwrap();
let mut inner = handle.write();
let (data, actual_codec) =
compress::compress(value, codec, self.config.compress_threshold, self.config.compression_level);
let (segment_id, offset, data_size) =
account.write_entry(key, &data, 0, actual_codec)?;
inner.write_entry(key, &data, 0, actual_codec)?;
let record = IndexRecord::new(key, segment_id, offset, data_size, 0);
account.append_index(&record)?;
inner.append_index(&record)?;
// Update indexed_up_to_offset for incremental recovery
let entry_end = offset + ENTRY_HEADER_SIZE as u64 + data_size as u64;
account.mark_indexed(segment_id, entry_end)?;
inner.mark_indexed(segment_id, entry_end)?;
let bucket_id = bucket::bucket_id(&key);
self.cache.update_record(account_id, bucket_id, record);
@@ -179,37 +173,31 @@ impl Engine {
pub fn read(&self, account_id: &str, key: &[u8; 32]) -> Result<Option<Vec<u8>>> {
let bucket_id = bucket::bucket_id(key);
// Acquire lock only to get bucket records and segment routing info.
let record: Option<IndexRecord> = {
let handle = {
let accounts = self.accounts.read().unwrap();
let account = accounts
accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
let records = self
.cache
.get_or_load(account_id, bucket_id, account.dir())?;
match records.binary_search_by(|r| r.key.cmp(key)) {
Ok(idx) => Some(records[idx].clone()),
Err(_) => None,
}
}; // accounts read lock dropped here — I/O happens outside the lock
let record = match record {
Some(r) => r,
None => return Ok(None),
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?
.clone()
};
if record.is_tombstone() {
return Ok(None);
}
// I/O outside the global lock
let seg_path = self
.root
.join("accounts")
.join(account_id)
.join("segments")
.join(segment::segment_filename(record.segment_id));
let (record, seg_path): (IndexRecord, PathBuf) = {
let inner = handle.read();
let records = self
.cache
.get_or_load(account_id, bucket_id, handle.dir())?;
match records.binary_search_by(|r| r.key.cmp(key)) {
Ok(idx) => {
let r = records[idx].clone();
if r.is_tombstone() {
return Ok(None);
}
let seg_path = inner.segment_path(r.segment_id)?;
(r, seg_path)
}
Err(_) => return Ok(None),
}
};
if !seg_path.exists() {
return Err(Error::SegmentNotFound(record.segment_id));
@@ -224,25 +212,25 @@ impl Engine {
}
pub fn delete(&self, account_id: &str, key: &[u8; 32]) -> Result<()> {
let mut accounts = self.accounts.write().unwrap();
let account = accounts
.get_mut(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
let handle = {
let accounts = self.accounts.read().unwrap();
accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?
.clone()
};
// Acquire per-account write lock
{
let _lock = account.lock_write();
}
let _write_lock = handle.write_mutex.lock().unwrap();
let mut inner = handle.write();
let (segment_id, offset, data_size) =
account.write_entry(*key, &[], 1, Codec::None)?;
inner.write_entry(*key, &[], 1, Codec::None)?;
let record = IndexRecord::new(*key, segment_id, offset, data_size, 1);
account.append_index(&record)?;
inner.append_index(&record)?;
// Update indexed_up_to_offset for incremental recovery
let entry_end = offset + ENTRY_HEADER_SIZE as u64 + data_size as u64;
account.mark_indexed(segment_id, entry_end)?;
inner.mark_indexed(segment_id, entry_end)?;
let bucket_id = bucket::bucket_id(key);
self.cache.update_record(account_id, bucket_id, record);
@@ -250,25 +238,24 @@ impl Engine {
Ok(())
}
// ── Batch write ─────────────────────────────────────────────────────────
// ── Batch write ─────────────────────────────────────────────────────
/// Batch-write multiple entries with a single fsync.
/// Each element is (key, value, codec).
pub fn write_batch(&self, account_id: &str, entries: &[([u8; 32], Vec<u8>, Codec)]) -> Result<()> {
if entries.is_empty() {
return Ok(());
}
let mut accounts = self.accounts.write().unwrap();
let account = accounts
.get_mut(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
let handle = {
let accounts = self.accounts.read().unwrap();
accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?
.clone()
};
{
let _lock = account.lock_write();
}
let _write_lock = handle.write_mutex.lock().unwrap();
let mut inner = handle.write();
// Phase 1: compress and append all entries without fsync
let mut pending: Vec<(IndexRecord, u64)> = Vec::with_capacity(entries.len());
for (key, value, codec) in entries {
if value.len() > crate::types::MAX_VALUE_SIZE {
@@ -278,20 +265,18 @@ impl Engine {
compress::compress(value, *codec, self.config.compress_threshold, self.config.compression_level);
let (segment_id, offset, data_size) =
account.append_entry(*key, &data, 0, actual_codec)?;
inner.append_entry(*key, &data, 0, actual_codec)?;
let entry_end = offset + ENTRY_HEADER_SIZE as u64 + data_size as u64;
let record = IndexRecord::new(*key, segment_id, offset, data_size, 0);
pending.push((record, entry_end));
}
// Phase 2: single fsync
account.flush_active()?;
inner.flush_active()?;
// Phase 3: append indices and update cache
for (record, entry_end) in &pending {
account.append_index(record)?;
account.mark_indexed(record.segment_id, *entry_end)?;
inner.append_index(record)?;
inner.mark_indexed(record.segment_id, *entry_end)?;
let bucket_id = bucket::bucket_id(&record.key);
self.cache.update_record(account_id, bucket_id, record.clone());
@@ -300,21 +285,19 @@ impl Engine {
Ok(())
}
// ── GC ──────────────────────────────────────────────────────────────────
// ── GC ──────────────────────────────────────────────────────────────
pub fn gc(&self, account_id: &str) -> Result<Option<GcStats>> {
// Get account dir while holding read lock, then release before heavy I/O.
let account_dir = {
let accounts = self.accounts.read().unwrap();
let account = accounts
let handle = accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
account.dir().to_path_buf()
}; // read lock released — GC runs without blocking reads
handle.dir().to_path_buf()
};
let result = gc::gc_account(&account_dir, self.config.gc_deleted_ratio)?;
// Invalidate all cached buckets for this account after GC rewrites them
for bid in 0..crate::types::BUCKET_COUNT {
self.cache.invalidate(account_id, bid);
}
@@ -325,10 +308,10 @@ impl Engine {
pub fn compact_buckets(&self, account_id: &str) -> Result<()> {
let account_dir = {
let accounts = self.accounts.read().unwrap();
let account = accounts
let handle = accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
account.dir().to_path_buf()
handle.dir().to_path_buf()
};
gc::compact_buckets(&account_dir)?;
@@ -340,15 +323,19 @@ impl Engine {
Ok(())
}
// ── Stats / Shutdown ────────────────────────────────────────────────────
// ── Stats / Shutdown ────────────────────────────────────────────────
pub fn stats(&self, account_id: &str) -> Result<AccountStats> {
let accounts = self.accounts.read().unwrap();
let account = accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?;
let handle = {
let accounts = self.accounts.read().unwrap();
accounts
.get(account_id)
.ok_or_else(|| Error::AccountNotFound(account_id.to_string()))?
.clone()
};
let meta = account.meta();
let inner = handle.read();
let meta = inner.meta();
let mut total_bytes = 0u64;
let mut deleted_bytes = 0u64;
@@ -357,12 +344,11 @@ impl Engine {
deleted_bytes += seg.deleted_bytes;
}
// Count total live keys from bucket indices
let mut total_keys = 0u64;
for bid in 0..crate::types::BUCKET_COUNT {
if let Ok(records) =
self.cache
.get_or_load(account_id, bid, account.dir())
.get_or_load(account_id, bid, handle.dir())
{
total_keys += records.iter().filter(|r| !r.is_tombstone()).count() as u64;
}
@@ -377,11 +363,11 @@ impl Engine {
})
}
/// Gracefully shut down: fsync all active segments and persist meta.
pub fn shutdown(&self) -> Result<()> {
let mut accounts = self.accounts.write().unwrap();
for (_, account) in accounts.iter_mut() {
account.flush_active()?;
let accounts = self.accounts.read().unwrap();
for (_, handle) in accounts.iter() {
let mut inner = handle.write();
inner.flush_active()?;
}
let global = GlobalMeta::load(&self.root)?;
global.save(&self.root)?;

View File

@@ -11,6 +11,7 @@ pub mod recovery;
pub mod segment;
pub mod types;
pub use account::AccountHandle;
pub use engine::{AccountStats, Engine};
pub use error::{Error, Result};
pub use types::{Codec, Config};