diff --git a/Cargo.lock b/Cargo.lock index 04075b601..3f39f89ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8309,6 +8309,7 @@ dependencies = [ "serde_json", "storage-util", "tempfile", + "tikv-jemallocator", "tokio", "toml", "tracing", diff --git a/storage/overlaybd/src/lsmt/file/helper.rs b/storage/overlaybd/src/lsmt/file/helper.rs index e61eabf24..2d275c52b 100644 --- a/storage/overlaybd/src/lsmt/file/helper.rs +++ b/storage/overlaybd/src/lsmt/file/helper.rs @@ -9,7 +9,7 @@ use std::io::{self, ErrorKind}; use std::mem::size_of; use std::os::unix::fs::{FileExt, OpenOptionsExt}; use std::path::{Path, PathBuf}; -use std::sync::{Arc, OnceLock, Weak}; +use std::sync::{Arc, Mutex as StdMutex, OnceLock, Weak}; use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::{Mutex, OwnedMutexGuard}; use uuid::Uuid; @@ -472,50 +472,157 @@ pub(super) fn decode_premerged_index_artifact( ))) } +/// Growth in artifact bytes that amortizes one full-dir prune scan. +const PREMERGED_INDEX_PRUNE_SCAN_FRACTION: u64 = 8; + +/// Per-cache-dir growth accounting and scan serialization. +pub(super) struct PremergedIndexPruneState { + account: StdMutex, +} + +struct PremergedIndexScanAccount { + bytes_since_scan: u64, + scan_in_progress: bool, +} + +pub(super) async fn premerged_index_prune_state(cache_dir: &Path) -> Arc { + static STATES: OnceLock>>> = + OnceLock::new(); + STATES + .get_or_init(|| Mutex::new(HashMap::new())) + .lock() + .await + .entry(cache_dir.to_path_buf()) + .or_insert_with(|| { + Arc::new(PremergedIndexPruneState { + account: StdMutex::new(PremergedIndexScanAccount { + bytes_since_scan: 0, + scan_in_progress: false, + }), + }) + }) + .clone() +} + +impl PremergedIndexPruneState { + /// Charge `written_bytes` and elect the single caller that runs the next scan. + pub(super) fn elect_scan(&self, written_bytes: u64, max_dir_bytes: u64) -> bool { + let threshold = (max_dir_bytes / PREMERGED_INDEX_PRUNE_SCAN_FRACTION).max(1); + let mut account = self.account.lock().unwrap_or_else(|err| err.into_inner()); + account.bytes_since_scan = account.bytes_since_scan.saturating_add(written_bytes); + if account.bytes_since_scan >= threshold && !account.scan_in_progress { + account.bytes_since_scan = 0; + account.scan_in_progress = true; + return true; + } + false + } + + /// End the scan; returns true when growth charged while it ran crosses + /// the trigger and a follow-up scan should run. + pub(super) fn scan_finished(&self, max_dir_bytes: u64) -> bool { + let threshold = (max_dir_bytes / PREMERGED_INDEX_PRUNE_SCAN_FRACTION).max(1); + let mut account = self.account.lock().unwrap_or_else(|err| err.into_inner()); + account.scan_in_progress = false; + if account.bytes_since_scan >= threshold { + account.bytes_since_scan = 0; + account.scan_in_progress = true; + return true; + } + false + } + + /// Release a failed scan; concurrent charges are kept so the next write + /// re-elects. + pub(super) fn scan_aborted(&self) { + let mut account = self.account.lock().unwrap_or_else(|err| err.into_inner()); + account.scan_in_progress = false; + } +} + +/// One blocking task for the whole scan: async `tokio::fs` would cost a +/// blocking-pool round-trip per directory entry. async fn prune_premerged_index_dir(dir: &Path, max_dir_bytes: u64) -> Result<()> { + let dir = dir.to_path_buf(); + tokio::task::spawn_blocking(move || prune_premerged_index_dir_blocking(&dir, max_dir_bytes)) + .await + .context("join premerged index cache prune task")? +} + +struct PremergedArtifactEntry { + path: PathBuf, + len: u64, + modified: SystemTime, +} + +/// Delete oldest premerged index artifacts until `dir` fits `max_dir_bytes`. +fn prune_premerged_index_dir_blocking(dir: &Path, max_dir_bytes: u64) -> Result<()> { let mut entries = Vec::new(); let mut total = 0u64; - let mut reader = tokio::fs::read_dir(dir).await?; - - while let Some(entry) = reader.next_entry().await? { + for entry in std::fs::read_dir(dir) + .with_context(|| format!("read premerged index cache dir {}", dir.display()))? + { + let entry = entry?; let path = entry.path(); if path.extension().and_then(|v| v.to_str()) != Some(PREMERGED_INDEX_EXT) { continue; } - let metadata = entry.metadata().await?; + let metadata = match entry.metadata() { + Ok(metadata) => metadata, + // Entries deleted concurrently with the scan only shrink it. + Err(err) if err.kind() == ErrorKind::NotFound => continue, + Err(err) => { + return Err(err).with_context(|| format!("stat {}", path.display())); + } + }; if !metadata.is_file() { continue; } let len = metadata.len(); let modified = metadata.modified().unwrap_or(UNIX_EPOCH); total = total.saturating_add(len); - entries.push((path, len, modified)); + entries.push(PremergedArtifactEntry { + path, + len, + modified, + }); } if total <= max_dir_bytes { return Ok(()); } - entries.sort_by_key(|(_, _, modified): &(PathBuf, u64, SystemTime)| *modified); - for (path, len, _) in entries { + entries.sort_by_key(|entry| entry.modified); + for entry in entries { if total <= max_dir_bytes { break; } - match tokio::fs::remove_file(&path).await { - Ok(()) => total = total.saturating_sub(len), + match std::fs::remove_file(&entry.path) { + Ok(()) => total = total.saturating_sub(entry.len), + // Already removed by someone else since the scan: those bytes + // left the dir too, so count them as freed. + Err(err) if err.kind() == ErrorKind::NotFound => { + total = total.saturating_sub(entry.len); + } Err(err) => { - tracing::warn!(?err, path = %path.display(), "remove premerged index artifact failed") + tracing::warn!( + ?err, + path = %entry.path.display(), + "remove premerged index artifact failed" + ) } } } Ok(()) } +/// Write the artifact atomically (tmp file + rename) and return its size in +/// bytes so callers can account cache growth. async fn write_premerged_index_artifact( cache_dir: &Path, key: &PremergedIndexCacheKey, index: &ReadOnlyIndex, -) -> Result<()> { +) -> Result { let dir = cache_dir.join(PREMERGED_INDEX_DIR); tokio::fs::create_dir_all(&dir) .await @@ -550,7 +657,7 @@ async fn write_premerged_index_artifact( }); } - Ok(()) + Ok(artifact.len() as u64) } pub(super) async fn try_read_premerged_index_artifact( @@ -582,6 +689,34 @@ pub(super) async fn try_read_premerged_index_artifact( } } +/// Charge `written` artifact bytes for `cache_dir` and run prune scans +/// while the growth trigger keeps electing. +async fn prune_premerged_index_cache(cache_dir: &Path, written: u64, max_dir_bytes: u64) { + let state = premerged_index_prune_state(cache_dir).await; + if !state.elect_scan(written, max_dir_bytes) { + return; + } + let dir = cache_dir.join(PREMERGED_INDEX_DIR); + loop { + match prune_premerged_index_dir(&dir, max_dir_bytes).await { + Ok(()) => { + if !state.scan_finished(max_dir_bytes) { + return; + } + } + Err(err) => { + tracing::warn!( + ?err, + path = %dir.display(), + "prune premerged index cache dir failed" + ); + state.scan_aborted(); + return; + } + } + } +} + pub(super) fn spawn_premerged_index_artifact_write( cache_dir: PathBuf, key: PremergedIndexCacheKey, @@ -599,18 +734,15 @@ pub(super) fn spawn_premerged_index_artifact_write( "write premerged index artifact failed" ); } - let key_digest = key.digest_hex.clone(); + // Release the merged index and the digest lock before the prune tail: + // the index can be hundreds of MB and the prune scan may pin it for + // the whole scan, while the held lock would block the next writer for + // the same digest. + drop(merged); drop(guard); - release_premerged_index_lock(&key_digest, &lock).await; - if write_result.is_ok() { - let dir = cache_dir.join(PREMERGED_INDEX_DIR); - if let Err(err) = prune_premerged_index_dir(&dir, max_dir_bytes).await { - tracing::warn!( - ?err, - path = %dir.display(), - "prune premerged index cache dir failed" - ); - } + release_premerged_index_lock(&key.digest_hex, &lock).await; + if let Ok(written) = write_result { + prune_premerged_index_cache(&cache_dir, written, max_dir_bytes).await; } }); } diff --git a/storage/overlaybd/src/lsmt/file/tests.rs b/storage/overlaybd/src/lsmt/file/tests.rs index 907ccd0f5..ea93c6940 100644 --- a/storage/overlaybd/src/lsmt/file/tests.rs +++ b/storage/overlaybd/src/lsmt/file/tests.rs @@ -2536,6 +2536,60 @@ async fn test_premerged_index_lock_map_drops_idle_entries() { release_premerged_index_lock(&stale_key, &replacement).await; } +#[tokio::test] +async fn test_premerged_index_prune_state_is_per_dir_and_elects_one_scan() { + let temp_dir = TempDir::new().unwrap(); + let cache_a = temp_dir.path().join("cache-a"); + let cache_b = temp_dir.path().join("cache-b"); + + let a = premerged_index_prune_state(&cache_a).await; + assert!(Arc::ptr_eq( + &a, + &premerged_index_prune_state(&cache_a).await + )); + let b = premerged_index_prune_state(&cache_b).await; + assert!(!Arc::ptr_eq(&a, &b)); + + // Budget 64 triggers one scan per 8 new bytes (`max_dir_bytes` / 8). + let budget = 64; + assert!(!a.elect_scan(7, budget)); + // Growth charged to A leaves B below B's own trigger. + assert!(!b.elect_scan(7, budget)); + // The winner's charge resets: the next scan needs a fresh trigger of + // post-election growth. + assert!(a.elect_scan(1, budget)); + + // Crossings while a scan runs lose the election but keep their charge: + // when they cross the trigger, scan_finished chains one follow-up scan + // instead of waiting for new writes. + assert!(!a.elect_scan(8, budget)); + assert!(a.scan_finished(budget)); + assert!(!a.elect_scan(1, budget)); + // Below the trigger: the chain ends and the gate is released. + assert!(!a.scan_finished(budget)); + assert!(a.elect_scan(8, budget)); + // No growth during the scan: no follow-up. + assert!(!a.scan_finished(budget)); + + // A failed scan releases the gate but keeps concurrent charges. + assert!(a.elect_scan(8, budget)); + assert!(!a.elect_scan(7, budget)); + a.scan_aborted(); + assert!(a.elect_scan(1, budget)); + assert!(!a.scan_finished(budget)); + + // A sub-fraction budget floors the trigger at one byte instead of zero + // (which would elect a scan per write). + assert!(a.elect_scan(1, 0)); + assert!(!a.scan_finished(0)); + + // Saturating accounting: a u64::MAX charge on a near-trigger counter + // must cross the trigger, not wrap 7 + MAX back below it. + assert!(!a.elect_scan(7, budget)); + assert!(a.elect_scan(u64::MAX, budget)); + assert!(!a.scan_finished(budget)); +} + fn premerged_artifact_count(cache_dir: &std::path::Path) -> usize { let dir = cache_dir.join(PREMERGED_INDEX_DIR); match std::fs::read_dir(dir) { diff --git a/storage/ublk-daemon/Cargo.toml b/storage/ublk-daemon/Cargo.toml index 69736265b..61b53b860 100644 --- a/storage/ublk-daemon/Cargo.toml +++ b/storage/ublk-daemon/Cargo.toml @@ -24,6 +24,7 @@ tokio = { version = "1.44.2", features = ["full"] } tracing = "0.1.41" tracing-log = "0.2.0" reqwest = { version = "0.13", default-features = false, features = ["rustls", "json"] } +tikv-jemallocator = { version = "0.6", features = ["background_threads"] } [dev-dependencies] tempfile = "3" diff --git a/storage/ublk-daemon/src/main.rs b/storage/ublk-daemon/src/main.rs index 4a9450263..260b70702 100644 --- a/storage/ublk-daemon/src/main.rs +++ b/storage/ublk-daemon/src/main.rs @@ -15,6 +15,10 @@ use uvm_ublk_daemon::{server::UblkDaemonServer, ResizeToolSpec}; mod metrics_server; +// Mirrors src/bin/server.rs. +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[derive(Debug, Parser)] #[command( name = "uvm-ublk-daemon",