Fix: bound uvm-ublk-daemon memory growth under pause/resume churn - #179
Fix: bound uvm-ublk-daemon memory growth under pause/resume churn#179YXalix wants to merge 2 commits into
Conversation
|
🔍 OpenCodeReview found 4 issue(s) in this PR.
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| const PREMERGED_INDEX_PRUNE_SCAN_FRACTION: u64 = 8; | ||
|
|
||
| /// Artifact bytes written since the last prune scan. | ||
| static PREMERGED_INDEX_BYTES_SINCE_SCAN: AtomicU64 = AtomicU64::new(0); |
There was a problem hiding this comment.
This accounting counter is process-global, but the prune operation is scoped to dir. If two callers use different cache directories (or the same directory with different limits), bytes charged by one are accumulated and can reset the counter for the other; the next scan may therefore be triggered for the wrong directory and the first directory can grow indefinitely without a scan. Keep the byte counter keyed by cache directory/limit (or otherwise scope the election to the cache being pruned).
| /// one caller to run the next prune scan, once growth since the last scan | ||
| /// reaches `max_dir_bytes` / [`PREMERGED_INDEX_PRUNE_SCAN_FRACTION`]. | ||
| fn charge_premerged_index_prune(written_bytes: u64, max_dir_bytes: u64) -> bool { | ||
| let threshold = max_dir_bytes / PREMERGED_INDEX_PRUNE_SCAN_FRACTION; |
There was a problem hiding this comment.
Because max_dir_bytes is an externally constructible public policy field, values below 8 are possible even though hybrid() clamps its result. Integer division makes threshold zero for those values, so every successful write performs an atomic swap and elects a prune scan, defeating the intended amortization. Validate/reject such limits or clamp the threshold to a nonzero, explicitly supported minimum.
|
|
||
| entries.sort_by_key(|(_, _, modified): &(PathBuf, u64, SystemTime)| *modified); | ||
| for (path, len, _) in entries { | ||
| entries.sort_by_key(|entry| entry.modified); |
There was a problem hiding this comment.
An elected writer still awaits a full directory scan that collects every artifact and sorts all entries by modification time. With a large cache directory this is O(n) metadata plus O(n) memory and O(n log n) CPU, and the blocking-pool worker remains occupied for the entire operation, causing latency spikes and potentially delaying unrelated blocking tasks. Consider a bounded/incremental eviction strategy or an independently scheduled maintenance task if large directories are expected.
|
Thanks for the fix and the detailed validation results. While reviewing the changes, I noticed three points that I would like to confirm:
Could you please confirm whether these observations are correct, or whether there is additional context that prevents these cases? |
| .get_or_init(|| Mutex::new(HashMap::new())) | ||
| .lock() | ||
| .await | ||
| .entry(cache_dir.to_path_buf()) |
There was a problem hiding this comment.
The registry key is the raw PathBuf, so aliases such as a relative path versus its absolute form (or paths containing ..) for the same cache directory create independent scan_in_progress flags. Concurrent callers using those spellings can then run overlapping full-directory scans, defeating the stated per-cache-dir serialization and increasing blocking-pool pressure. Normalize the path before using it as the registry key, or enforce a canonical/normalized representation at the API boundary.
| let previous = self | ||
| .bytes_since_scan | ||
| .fetch_add(written_bytes, Ordering::Relaxed); | ||
| let total = previous.wrapping_add(written_bytes); |
There was a problem hiding this comment.
Using fetch_add together with wrapping_add allows the accounting counter to wrap to a value below threshold. After enough accumulated writes (or an unusually large written_bytes), a threshold crossing can be missed and pruning can be deferred until the counter wraps again. Use saturating/checked atomic accounting so any overflow remains at or above the trigger.
| let previous = self | ||
| .bytes_since_scan | ||
| .fetch_add(written_bytes, Ordering::Relaxed); | ||
| let total = previous.wrapping_add(written_bytes); | ||
| total >= threshold && !self.scan_in_progress.swap(true, Ordering::AcqRel) |
There was a problem hiding this comment.
bytes_since_scan is updated before checking scan_in_progress, while scan_finished resets the counter before clearing the flag. A writer can execute fetch_add, then the scanner can reset the counter and clear the flag, and finally that writer's swap(true) can succeed using its pre-reset total. This permits the back-to-back scan that the method explicitly intends to suppress. The election and reset need a single synchronized state transition (for example, one atomic state containing both fields, or a mutex/CAS protocol that revalidates the counter after acquiring the scan flag).
| Ok(()) => total = total.saturating_sub(len), | ||
| match std::fs::remove_file(&entry.path) { | ||
| Ok(()) => total = total.saturating_sub(entry.len), | ||
| Err(err) if err.kind() == ErrorKind::NotFound => {} |
There was a problem hiding this comment.
When another writer or pruning process removes this artifact after the scan, NotFound means its bytes are no longer part of the directory. Leaving total unchanged makes the loop believe the cache is still over budget and can delete additional valid artifacts, causing unnecessary cache eviction. Subtract entry.len in this branch as well.
Suggestion:
| Err(err) if err.kind() == ErrorKind::NotFound => {} | |
| Err(err) if err.kind() == ErrorKind::NotFound => { | |
| total = total.saturating_sub(entry.len); | |
| } |
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_premerged_index_prune_state_is_per_dir_and_elects_one_scan() { |
There was a problem hiding this comment.
The added test covers basic per-directory identity and single election, but not the filesystem race where an enumerated artifact disappears before remove_file, nor the interleaving between elect_scan and scan_finished that can trigger a back-to-back scan. Tests for these cases would help protect the new accounting and concurrency invariants.
| Component::ParentDir => { | ||
| normalized.pop(); | ||
| } |
There was a problem hiding this comment.
Leading parent components in relative paths are discarded. For example, ../cache normalizes to cache, and a/../../cache also collapses too far, so distinct cache directories can share one prune state and suppress each other's scans. Preserve .. whenever there is no normal component available to pop (while handling roots/prefixes separately).
Suggestion:
| Component::ParentDir => { | |
| normalized.pop(); | |
| } | |
| Component::ParentDir => { | |
| let can_pop = matches!(normalized.components().next_back(), Some(Component::Normal(_))); | |
| if can_pop { | |
| normalized.pop(); | |
| } else if !normalized.is_absolute() { | |
| normalized.push(component.as_os_str()); | |
| } | |
| } |
| account.bytes_since_scan = 0; | ||
| account.scan_in_progress = false; |
There was a problem hiding this comment.
This drops all bytes written while the scan was running. A scan can finish enumerating its snapshot, then concurrent writers can add enough artifacts to exceed max_dir_bytes; resetting their charges means no follow-up scan occurs, and the directory can remain over budget indefinitely if the burst stops. Reset the pre-scan charge when electing, retain charges accumulated during the scan, and have completion elect/run a follow-up scan when that retained amount crosses the threshold.
| .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() |
There was a problem hiding this comment.
This process-wide map retains a strong Arc and copied PathBuf for every cache directory ever passed to the public cache-opening API. Callers using transient/per-image directories will grow STATES for the lifetime of the process, even after all writes and scans finish. Store Weak<PremergedIndexPruneState> values and remove/recreate stale entries, as the digest-lock registry already does, to bound registry lifetime.
| 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; | ||
| } |
There was a problem hiding this comment.
A newly created process state starts at zero, so this is no longer guaranteed to inspect an already oversized on-disk cache on the first write. After a restart, cache-limit reduction, or earlier failed cleanup, pruning is deferred until new artifacts total max_dir_bytes / 8 (at least 8 MiB with the production minimum), even though the directory may already exceed max_dir_bytes. There is no startup/maintenance prune elsewhere; initialize a new directory state so its first successful write elects a scan, or otherwise perform an initial size check.
| pub(super) fn scan_aborted(&self) { | ||
| let mut account = self.account.lock().unwrap_or_else(|err| err.into_inner()); | ||
| account.scan_in_progress = false; | ||
| } |
There was a problem hiding this comment.
The charge that elected the failed scan was already reset to zero in elect_scan, but this abort path restores only the gate. Therefore, when a scan fails without at least another threshold of concurrent writes, the next ordinary write does not re-elect as documented; cleanup can be postponed by another full threshold after a transient read_dir, metadata, or join failure. Preserve/restore the triggering threshold on abort (while retaining concurrent charges) so the next write retries.
| let a = premerged_index_prune_state(&cache_a).await; | ||
| assert!(Arc::ptr_eq( | ||
| &a, | ||
| &premerged_index_prune_state(&cache_a).await | ||
| )); |
There was a problem hiding this comment.
This test only invokes elect_scan sequentially, so it does not exercise the concurrency invariant implied by the shared per-directory state: simultaneous writers could both race around scan election or charge accounting while this test still passes. Please add a concurrent test (for example, synchronized Tokio tasks issuing charges against the same Arc) and assert that exactly one caller is elected and pending growth is preserved.
Suggestion:
| let a = premerged_index_prune_state(&cache_a).await; | |
| assert!(Arc::ptr_eq( | |
| &a, | |
| &premerged_index_prune_state(&cache_a).await | |
| )); | |
| let a = premerged_index_prune_state(&cache_a).await; | |
| assert!(Arc::ptr_eq( | |
| &a, | |
| &premerged_index_prune_state(&cache_a).await | |
| )); | |
| // Exercise election and accounting from concurrent callers here. |
| // Budget 64 triggers one scan per 8 new bytes (`max_dir_bytes` / 8). | ||
| let budget = 64; | ||
| assert!(!a.elect_scan(7, budget)); |
There was a problem hiding this comment.
The assertions encode private policy details (max_dir_bytes / 8, counter reset timing, and the exact follow-up-scan state machine) rather than an externally observable pruning contract. This makes the test unnecessarily brittle: a valid change to the amortization threshold or accounting strategy would require rewriting it even if cache pruning behavior remains correct. Consider keeping this as a focused state-machine unit test only if these details are an intentional invariant; otherwise prefer a behavior-oriented test that writes artifacts and verifies the directory is bounded.
Suggestion:
| // Budget 64 triggers one scan per 8 new bytes (`max_dir_bytes` / 8). | |
| let budget = 64; | |
| assert!(!a.elect_scan(7, budget)); | |
| // Verify pruning behavior through artifact writes and the resulting directory size. |
…cache growth Every artifact write scheduled a full tokio::fs prune of the cache dir: one blocking-pool round-trip per directory entry. - Run the scan as one spawn_blocking task. - Trigger a scan per 1/8th of max_dir_bytes of new writes, at most one per dir at a time; growth charged during a scan is retained and chains a follow-up scan when it crosses the trigger. - Drop the merged index before the lock-release/prune tail so queued writes stop pinning it.
glibc malloc leaves RSS plateaued high after pause/resume churn frees multi-MB buffers over and over (issue kvcache-ai#171). - Swap in tikv-jemallocator with background_threads, mirroring src/bin/server.rs.
| static STATES: OnceLock<Mutex<HashMap<PathBuf, Arc<PremergedIndexPruneState>>>> = | ||
| OnceLock::new(); |
There was a problem hiding this comment.
This process-global registry retains a strong Arc for every distinct raw PathBuf indefinitely. A long-running daemon that receives many cache-directory paths can therefore grow this map without bound; equivalent relative/absolute or symlinked paths also get separate states and can scan the same physical directory concurrently. Canonicalize the key and use weak references or an eviction/removal policy if the set of cache directories is not strictly fixed.
| bytes_since_scan: 0, | ||
| scan_in_progress: false, |
There was a problem hiding this comment.
The accounting starts at zero and only grows when this process successfully writes an artifact; it never inspects an existing directory. Because the cache directory persists across restarts and max_dir_bytes can change, a pre-populated or already-oversized cache will not be pruned until another max_dir_bytes / 8 bytes are written. With low write activity, this can leave the cache over its configured limit indefinitely, whereas the previous implementation checked after every successful write. Consider initializing the state with an initial scan/election (or otherwise ensuring the first use prunes an existing over-limit directory).
| pub(super) fn scan_aborted(&self) { | ||
| let mut account = self.account.lock().unwrap_or_else(|err| err.into_inner()); | ||
| account.scan_in_progress = false; | ||
| } |
There was a problem hiding this comment.
On an aborted scan, the charge that elected the scan was already reset in elect_scan, and this method only clears the gate. If the scan fails before any concurrent writes add credit, the next write must accumulate a full new threshold before retrying; repeated transient read/join failures can therefore keep an oversized cache unpruned. Preserve a retry-pending/threshold charge when aborting (or explicitly schedule a bounded retry) instead of discarding the election.
| let a = premerged_index_prune_state(&cache_a).await; | ||
| assert!(Arc::ptr_eq( | ||
| &a, | ||
| &premerged_index_prune_state(&cache_a).await | ||
| )); |
There was a problem hiding this comment.
This only checks sequential lookups. It does not exercise concurrent first-time initialization of the global state map, so a check-then-insert implementation could still create multiple Arcs for the same directory and lose accounting while these assertions pass. Add concurrent lookups (for example, several tasks synchronized by a barrier) and assert all returned Arcs are pointer-identical.
@Hygge-Gezelligheid thanks for the careful review — all three observations are correct. Here is where they stand now:
|
2 participants
What
Bound
uvm-ublk-daemonmemory growth under pause/resume churn, via two commits:8939f2bfix(overlaybd): offload premerged index prune scan and trigger it by cache growth — the prune scan runs as onespawn_blockingtask overstd::fs(one blocking op per scan, not per directory entry), triggers by cache growth (one scan permax_dir_bytes / 8of new writes instead of one queued scan per write), and no longer pins the merged index while queued.355692bfix(ublk-daemon): adopt jemalloc — globaltikv-jemallocatorwith thebackground_threadsfeature (compile-timebackground_thread:true), mirroringsrc/bin/server.rs, so freed pages return to the OS instead of lingering in glibc arenas. No run-timemalloc_confexport is kept; the default 10s dirty/muzzy decay is what the validation numbers below actually ran with.Stress replay of the issue's workload — 600 sandboxes, 100 running concurrently,
continuous pause/resume (16,128 timing samples per run)
8939f2bonly (prune)Why
Issue #171: a 600-VM pause/resume stress cycle grows daemon anonymous memory at 20–40 MB/s up to tens of GB.
memleak-bpfccpoints at the premerged-index write path. Two compounding mechanisms:tokio::fs::ReadDir(one blocking-pool round-trip per entry — the 101kread_dirallocations in the trace; a large cache dir becomes tens of thousands of queued blocking ops). Each queued task also pinned itsArc<ReadOnlyIndex>(can be hundreds of MB) and a manifestVec<(PathBuf, u64, SystemTime)>of the whole directory (the ~27.5 MBgrow_onestack in the trace).On this workload the baseline reproduces the reported behavior (28.39 GB daemon anon memory); the two commits together bring it to 0.72 GB (−97%) — see Validation.
Related issue
Closes #171
Scope and non-goals
Included: only the two commits above —
storage/overlaybd/src/lsmt/file/{helper,tests}.rs(prune path) andstorage/ublk-daemon/{Cargo.toml,src/main.rs}(allocator).Non-goals: no redesign of the cache layout (e.g. incremental artifact registry or generation directories — considered, deferred until scan cost is shown to matter beyond this fix); no changes to the premerged artifact format; no unrelated refactoring.
Design and behavior changes
spawn_blockingtask usingstd::fs; entries deleted concurrently with the scan are skipped instead of aborting the prune.cache_dir, like the other premerged helpers) charges written artifact bytes; once growth since the last scan reachesmax_dir_bytes / 8(floored at one byte, so tiny budgets cannot degenerate into a scan per write), a guarded swap elects exactly one scanner. At most one scan per dir runs at a time, and bytes charged during a running scan are voided when it ends — scans neither overlap nor run back to back, and their frequency is bounded by write rate, not dir size. Bursts coalesce into one scan; idle or slow caches never scan. Failed writes charge nothing (tmp file is removed, no net growth).drop(merged)before the lock-release/prune tail, so queued tasks stop pinning the merged index.src/bin/server.rsexactly —tikv-jemallocatorwith thebackground_threadsfeature, which compiles jemalloc with a build-timebackground_thread:trueso purging runs off the allocation path; dirty/muzzy decay keeps jemalloc's 10s defaults. A run-timemalloc_confexport is deliberately not kept (see Risks and reviewer notes).max_dir_bytes / 8of new bytes per cache dir (≥ 8 MiB given the 64 MiB minimum budget, so it cannot degenerate).Compatibility and operations
PMIDX001v1 format; only eviction scheduling changes.tikv-jemallocatortouvm-ublk-daemon(the same crate and feature the server already uses; bundled jemalloc, no system package needed); slightly larger daemon binary; no new ports or permissions.Validation
make fmtmake clippymake test-unitmake -C services test(required whenservices/changes)maketargetSkipped checks and reasons:
Risks and reviewer notes
fetch_adds its bytes, and the crossing writer wins an in-flight flag swap — exactly one scan fires per threshold crossing per dir. Losing writers keep their bytes accumulated until the elected scan voids them at its end. Relaxed ordering suffices because the counter guards no other memory.dirty_decay_ms:1000,muzzy_decay_ms:1000,background_thread:truevia amalloc_confsymbol, which the prefixed jemalloc build (--with-jemalloc-prefix=_rjem_) silently never reads — the validated 0.72 GB therefore reflects the default decay this PR now ships verbatim. The run-time export was dropped rather than kept as a silent-failure footgun.storage/overlaybd/src/lsmt/file/helper.rs,storage/ublk-daemon/src/main.rs.Checklist