[PR #179] [MERGED] Fix: bound uvm-ublk-daemon memory growth under pause/resume churn #232

New issue

Closed

opened 2026-10-29 14:21:17 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:21:17 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/AgentENV/pull/179
Author: @YXalix
Created: 8/17/2026
Status: ✅ Merged
Merged: 8/18/2026
Merged by: @yingdi-shan

Base: main ← Head: fix_leak

📝 Commits (2)

6a7800c fix(overlaybd): offload premerged index prune scan and trigger it by cache growth
c4c6dfd fix(ublk-daemon): adopt jemalloc

📊 Changes

5 files changed (+216 additions, -24 deletions)

View changed files

📝 Cargo.lock (+1 -0)
📝 storage/overlaybd/src/lsmt/file/helper.rs (+156 -24)
📝 storage/overlaybd/src/lsmt/file/tests.rs (+54 -0)
📝 storage/ublk-daemon/Cargo.toml (+1 -0)
📝 storage/ublk-daemon/src/main.rs (+4 -0)

📄 Description

What

Bound uvm-ublk-daemon memory growth under pause/resume churn, via two commits:

8939f2b fix(overlaybd): offload premerged index prune scan and trigger it by cache growth — the prune scan runs as one spawn_blocking task over std::fs (one blocking op per scan, not per directory entry), triggers by cache growth (one scan per max_dir_bytes / 8 of new writes instead of one queued scan per write), and no longer pins the merged index while queued.
355692b fix(ublk-daemon): adopt jemalloc — global tikv-jemallocator with the background_threads feature (compile-time background_thread:true), mirroring src/bin/server.rs, so freed pages return to the OS instead of lingering in glibc arenas. No run-time malloc_conf export 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)

Build	daemon anon_mmap	pause P50/P95	resume P50/P95	total P50/avg
baseline (no fixes)	28.39 GB	1.123 / 1.723 s	0.812 / 1.144 s	8.043 / 13.406 s
`8939f2b` only (prune)	1.54 GB	0.322 / 0.606 s	0.193 / 0.388 s	4.172 / 10.754 s
both commits	0.72 GB	0.338 / 0.669 s	0.190 / 0.443 s	4.303 / 10.826 s

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-bpfcc points at the premerged-index write path. Two compounding mechanisms:

Prune amplification — every artifact write spawned a task whose tail ran a full cache-dir prune via tokio::fs::ReadDir (one blocking-pool round-trip per entry — the 101k read_dir allocations in the trace; a large cache dir becomes tens of thousands of queued blocking ops). Each queued task also pinned its Arc<ReadOnlyIndex> (can be hundreds of MB) and a manifest Vec<(PathBuf, u64, SystemTime)> of the whole directory (the ~27.5 MB grow_one stack in the trace).
Allocator retention — glibc's dynamic mmap threshold climbs to 32 MB under multi-MB alloc/free churn, after which large buffers come off arena heaps and freed arena pages are never returned to the OS. Part of the reported "leak" is memory that was freed but retained.

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.

Closes #171

Scope and non-goals

Included: only the two commits above — storage/overlaybd/src/lsmt/file/{helper,tests}.rs (prune path) and storage/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

Scan offload: the scan/remove loop moved into a single spawn_blocking task using std::fs; entries deleted concurrently with the scan are skipped instead of aborting the prune.
Growth-based trigger: per-cache-dir state (keyed by cache_dir, like the other premerged helpers) charges written artifact bytes; once growth since the last scan reaches max_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).
Early index release: drop(merged) before the lock-release/prune tail, so queued tasks stop pinning the merged index.
Allocator: the daemon now mirrors src/bin/server.rs exactly — tikv-jemallocator with the background_threads feature, which compiles jemalloc with a build-time background_thread:true so purging runs off the allocation path; dirty/muzzy decay keeps jemalloc's 10s defaults. A run-time malloc_conf export is deliberately not kept (see Risks and reviewer notes).
Behavior change: scan frequency moves from once per write to once per max_dir_bytes / 8 of new bytes per cache dir (≥ 8 MiB given the 64 MiB minimum budget, so it cannot degenerate).

Compatibility and operations

Public API or generated protocol: none — internal overlaybd helper and daemon allocator only.
Configuration or defaults: none — the 1/8 scan fraction is a code constant; jemalloc defaults are unchanged.
Snapshot manifest, artifact layout, or storage format: none — premerged artifacts keep the existing PMIDX001 v1 format; only eviction scheduling changes.
Upgrade and rollback: the two fixes are independent; server and daemon binaries can roll forward/back independently with no on-disk coordination.
Host requirements, permissions, ports, or dependencies: adds vendored tikv-jemallocator to uvm-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 fmt
make clippy
make test-unit
Relevant Rust integration tests
make -C services test (required when services/ changes)
Generated clients/server regenerated with the documented make target
Documentation updated

Skipped checks and reasons:

This change does not affect the relevant components.

Risks and reviewer notes

Election: growth accounting is per cache dir; a write 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.
Allocator: an earlier iteration exported dirty_decay_ms:1000,muzzy_decay_ms:1000,background_thread:true via a malloc_conf symbol, 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.
Key files: storage/overlaybd/src/lsmt/file/helper.rs, storage/ublk-daemon/src/main.rs.

Checklist

The PR contains one coherent change and no unrelated formatting or refactoring.
New behavior is covered by tests, or I explained why testing is impractical.
Logs and examples contain no credentials, tokens, or private registry information.
I did not manually edit generated code without updating its source and regenerating it.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/AgentENV/pull/179 **Author:** [@YXalix](https://github.com/YXalix) **Created:** 8/17/2026 **Status:** ✅ Merged **Merged:** 8/18/2026 **Merged by:** [@yingdi-shan](https://github.com/yingdi-shan) **Base:** `main` ← **Head:** `fix_leak` --- ### 📝 Commits (2) - [`6a7800c`](https://github.com/kvcache-ai/AgentENV/commit/6a7800c857a6aa2719e4194a0d8703b6455a594a) fix(overlaybd): offload premerged index prune scan and trigger it by cache growth - [`c4c6dfd`](https://github.com/kvcache-ai/AgentENV/commit/c4c6dfd03e3965f8cd7c85ed8c1ceca8ce4d2759) fix(ublk-daemon): adopt jemalloc ### 📊 Changes **5 files changed** (+216 additions, -24 deletions) <details> <summary>View changed files</summary> 📝 `Cargo.lock` (+1 -0) 📝 `storage/overlaybd/src/lsmt/file/helper.rs` (+156 -24) 📝 `storage/overlaybd/src/lsmt/file/tests.rs` (+54 -0) 📝 `storage/ublk-daemon/Cargo.toml` (+1 -0) 📝 `storage/ublk-daemon/src/main.rs` (+4 -0) </details> ### 📄 Description ## What Bound `uvm-ublk-daemon` memory growth under pause/resume churn, via two commits: - `8939f2b` **fix(overlaybd): offload premerged index prune scan and trigger it by cache growth** — the prune scan runs as one `spawn_blocking` task over `std::fs` (one blocking op per scan, not per directory entry), triggers by cache growth (one scan per `max_dir_bytes / 8` of new writes instead of one queued scan per write), and no longer pins the merged index while queued. - `355692b` **fix(ublk-daemon): adopt jemalloc** — global `tikv-jemallocator` with the `background_threads` feature (compile-time `background_thread:true`), mirroring `src/bin/server.rs`, so freed pages return to the OS instead of lingering in glibc arenas. No run-time `malloc_conf` export 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) | Build | daemon anon_mmap | pause P50/P95 | resume P50/P95 | total P50/avg | | ---------------------- | ---------------- | --------------- | --------------- | ---------------- | | baseline (no fixes) | 28.39 GB | 1.123 / 1.723 s | 0.812 / 1.144 s | 8.043 / 13.406 s | | `8939f2b` only (prune) | 1.54 GB | 0.322 / 0.606 s | 0.193 / 0.388 s | 4.172 / 10.754 s | | both commits | **0.72 GB** | 0.338 / 0.669 s | 0.190 / 0.443 s | 4.303 / 10.826 s | ## 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-bpfcc` points at the premerged-index write path. Two compounding mechanisms: 1. **Prune amplification** — every artifact write spawned a task whose tail ran a full cache-dir prune via `tokio::fs::ReadDir` (one blocking-pool round-trip per entry — the 101k `read_dir` allocations in the trace; a large cache dir becomes tens of thousands of queued blocking ops). Each queued task also pinned its `Arc<ReadOnlyIndex>` (can be hundreds of MB) and a manifest `Vec<(PathBuf, u64, SystemTime)>` of the whole directory (the ~27.5 MB `grow_one` stack in the trace). 2. **Allocator retention** — glibc's dynamic mmap threshold climbs to 32 MB under multi-MB alloc/free churn, after which large buffers come off arena heaps and freed arena pages are never returned to the OS. Part of the reported "leak" is memory that was freed but retained. 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) and `storage/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 - **Scan offload**: the scan/remove loop moved into a single `spawn_blocking` task using `std::fs`; entries deleted concurrently with the scan are skipped instead of aborting the prune. - **Growth-based trigger**: per-cache-dir state (keyed by `cache_dir`, like the other premerged helpers) charges written artifact bytes; once growth since the last scan reaches `max_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). - **Early index release**: `drop(merged)` before the lock-release/prune tail, so queued tasks stop pinning the merged index. - **Allocator**: the daemon now mirrors `src/bin/server.rs` exactly — `tikv-jemallocator` with the `background_threads` feature, which compiles jemalloc with a build-time `background_thread:true` so purging runs off the allocation path; dirty/muzzy decay keeps jemalloc's 10s defaults. A run-time `malloc_conf` export is deliberately not kept (see Risks and reviewer notes). - **Behavior change**: scan frequency moves from once per write to once per `max_dir_bytes / 8` of new bytes per cache dir (≥ 8 MiB given the 64 MiB minimum budget, so it cannot degenerate). ## Compatibility and operations - Public API or generated protocol: none — internal overlaybd helper and daemon allocator only. - Configuration or defaults: none — the 1/8 scan fraction is a code constant; jemalloc defaults are unchanged. - Snapshot manifest, artifact layout, or storage format: none — premerged artifacts keep the existing `PMIDX001` v1 format; only eviction scheduling changes. - Upgrade and rollback: the two fixes are independent; server and daemon binaries can roll forward/back independently with no on-disk coordination. - Host requirements, permissions, ports, or dependencies: adds vendored `tikv-jemallocator` to `uvm-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  - [x] `make fmt` - [x] `make clippy` - [x] `make test-unit` - [ ] Relevant Rust integration tests - [ ] `make -C services test` (required when `services/` changes) - [x] Generated clients/server regenerated with the documented `make` target - [ ] Documentation updated Skipped checks and reasons: - This change does not affect the relevant components. ## Risks and reviewer notes - Election: growth accounting is per cache dir; a write `fetch_add`s 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. - Allocator: an earlier iteration exported `dirty_decay_ms:1000,muzzy_decay_ms:1000,background_thread:true` via a `malloc_conf` symbol, 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. - Key files: `storage/overlaybd/src/lsmt/file/helper.rs`, `storage/ublk-daemon/src/main.rs`. ## Checklist - [x] The PR contains one coherent change and no unrelated formatting or refactoring. - [x] New behavior is covered by tests, or I explained why testing is impractical. - [x] Logs and examples contain no credentials, tokens, or private registry information. - [x] I did not manually edit generated code without updating its source and regenerating it. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-29 14:21:17 +00:00