Fix: bound uvm-ublk-daemon memory growth under pause/resume churn

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<PremergedIndexScanAccount>,
+}
+
+struct PremergedIndexScanAccount {
+    bytes_since_scan: u64,
+    scan_in_progress: bool,
+}
+
+pub(super) async fn premerged_index_prune_state(cache_dir: &Path) -> Arc<PremergedIndexPruneState> {
+    static STATES: OnceLock<Mutex<HashMap<PathBuf, Arc<PremergedIndexPruneState>>>> =
+        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<u64> {
     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;
         }
     });
 }

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) {

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"

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",