kvcache-ai/AgentENV

Fork 0

mirror of https://github.com/kvcache-ai/AgentENV.git synced 2026-11-05 01:17:36 +00:00

[PR #92] [MERGED] feat(sandbox): add optional ZFile compression for memory snapshots #175

New issue

Closed

opened 2026-10-29 14:20:39 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:20:39 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/AgentENV/pull/92
Author: @huajq
Created: 7/31/2026
Status: ✅ Merged
Merged: 8/4/2026
Merged by: @yingdi-shan

Base: main ← Head: zfile

📝 Commits (3)

a7f24b4 feat(sandbox): add optional ZFile compression for memory snapshots
ccbae02 feat(overlaybd): add persistent worker pool and sync pwrite for ZFile compression
4b7501d test(sandbox): add memory snapshot compression integration tests

📊 Changes

22 files changed (+2072 additions, -796 deletions)

View changed files

📝 Cargo.lock (+2 -16)
📝 config/default.toml (+5 -0)
📝 config/oss_default.toml (+5 -0)
📝 crates/e2e-tests/Cargo.toml (+1 -0)
📝 crates/e2e-tests/tests/snapshot_oss_e2e_test.rs (+69 -6)
📝 docs/src/configuration/reference.md (+2 -0)
📝 src/cfg.rs (+64 -0)
📝 src/sandbox/firecracker/overlaybd_snapshot.rs (+241 -50)
📝 src/sandbox/firecracker/sandbox.rs (+20 -9)
📝 src/sandbox/ublk/mod.rs (+3 -1)
📝 src/sandbox/ublk/overlaybd.rs (+284 -12)
📝 src/snapshot/repository/backends/posixfs/artifacts.rs (+114 -0)
📝 storage/overlaybd/Cargo.toml (+2 -1)
📝 storage/overlaybd/src/backend/local.rs (+48 -0)
📝 storage/overlaybd/src/compression/zfile.rs (+941 -638)
📝 storage/overlaybd/src/dense_export.rs (+4 -0)
📝 storage/overlaybd/src/io/virtual_file.rs (+8 -0)
📝 storage/overlaybd/src/lsmt/file/helper.rs (+74 -29)
📝 storage/overlaybd/src/tools/mod.rs (+3 -1)
📝 storage/overlaybd/src/tools/packaging.rs (+26 -19)

...and 2 more files

📄 Description

What

Add optional ZFile compression support for Firecracker memory snapshots. When enabled, dirty memory pages are compressed during snapshot creation using LZ4 (default) or Zstd, reducing snapshot artifact size by ~70-88% at the cost of ~40-90% longer pause latency.

Why

Memory snapshots for large VMs produce multi-GB artifacts. Compressing them during creation reduces storage and transfer costs, which is especially valuable for snapshot templates shared across nodes.

Closes #

Scope and non-goals

Included:

[memory_snapshot] config: compression_enabled (default false), compression_algorithm (lz4/zstd), compression_workers (default 1, max 64)
Direct memory snapshot path: Firecracker dirty ranges → ProcessVmReader (process_vm_readv) → compact_to → raw or ZFileCompactWriter → temp → rename
ZFile compression core: persistent WorkerPool (std::thread + crossbeam-channel), CompressedSegment (contiguous output, no per-block Vec), 512 KiB compression batch, ordered committer, cancellation-safe poisoning
LocalFile::write_at_sync_pwrite for >4KiB batches (bypasses io_uring overhead on local page-cache writes)
compact_to ordered writer support, zeroed-index sort fix (Critical), MAX_LENGTH clamp
LZ4 switched from lz4_flex to lz4-sys (C liblz4, ~25% faster)
Integration tests for direct and legacy LZ4 compressed pause/resume

Excluded:

No changes to rootfs compression (rootfs layers remain raw)
No changes to snapshot manifest or artifact layout (ZFile is transparent to the repository)
No changes to resume path (switch file transparently decompresses ZFile)
No CI pipeline changes (integration tests run manually with compression configs)

Design and behavior changes

Data flow

Firecracker pause
  → create_state_only_snapshot (vm_state.bin)
  → GET /vm/dirty-memory-ranges
  → dirty_ranges_to_segment_mappings
  → ProcessVmReader (process_vm_readv, 512 KiB chunks)
  → compact_to
  → Raw VirtualFileWriter (unordered, io_uring)
    or ZFileCompactWriter (ordered, WorkerPool compression)
  → overlaybd.commit.tmp
  → rename overlaybd.commit

WorkerPool architecture

compact_to (async, Tokio worker)
  │
  │  512 KiB chunk via ProcessVmReader
  │
  ▼
crossbeam::channel::bounded<WorkItem>(workers * 2)
  │
  ├──► Worker Thread 1 (std::thread, owns compressor)
  ├──► Worker Thread 2
  ├──► Worker Thread 3
  └──► Worker Thread 4
         │
         ▼
  crossbeam::channel::bounded<CompressedBatch>(workers * 2)
         │
         ▼
  Ordered Committer (BTreeMap by seq)
    → commit_compressed_batch
    → LocalFile::write_at_sync_pwrite (>4KiB)
      or write_all_at (non-local fallback)

Key design decisions

workers=1 remains synchronous: no thread overhead for single-worker configs
512 KiB compression batch: amortizes worker dispatch (32 KiB → 512 KiB)
sync pwrite for large batches: bypasses io_uring submit/completion round trips on local page-cache writes (~2-3 µs vs ~15 µs per call)
Cancellation-safe poisoning: any builder write/finish await pre-marks failed=true, cleared only on success; future cancellation or errors leave the writer permanently poisoned
catch_unwind in workers: prevents worker panics from deadlocking the ordered committer drain loop

Performance (1 GiB structured text payload, 1536 MiB VM)

Mode	Pause (mean)	Artifact size	Ratio
Raw	801 ms	1,242 MB	1.00
LZ4 workers=4	1,127 ms	367 MB	0.30
Zstd workers=4	1,533 ms	149 MB	0.12

LZ4 compression adds ~40% pause latency for ~70% size reduction. Zstd adds ~90% for ~88% reduction.

Zeroed-index sort fix (Critical bug fix included)

compact_to previously pushed zeroed SegmentMapping entries into compact_index during phase 1, then appended data entries in phase 2. This produced an unsorted index ([all zeros..., all data...]), causing partition_point lookups in ReadOnlyIndex to silently miss mappings and read committed data as zeros. Fixed by sorting compact_index before compress_raw_index.

Compatibility and operations

Public API or generated protocol: N/A (no API changes)
Configuration or defaults: New [memory_snapshot] keys with safe defaults (disabled, lz4, 1 worker)
Snapshot manifest, artifact layout, or storage format: N/A (ZFile is transparent to the repository; resume path unchanged)
Upgrade and rollback: Old snapshots without compression continue to work. New compressed snapshots require the same code version to resume (ZFile read support is in the same crate)
Host requirements, permissions, ports, or dependencies: Adds crossbeam-channel dependency. No new system requirements.

Validation

make fmt (cargo fmt --all -- --check)
make clippy (cargo clippy -p overlaybd -p agentenv --all-targets -- -D warnings, zero warnings)
make test-unit (overlaybd lib 304/0/1 ignored, agentenv 6 core tests 6/6)
Relevant Rust integration tests (fc::memory_snapshot_format_matches_config_and_resumes, fc::legacy_memory_snapshot_format_matches_config_and_resumes, 2/2 with LZ4 workers=4)
make -C services test (N/A, no services/ changes)
Generated clients/server regenerated with the documented make target (N/A, no generated code changes)
Documentation updated (docs/src/configuration/reference.md)
Benchmarks or performance comparison completed (see Performance section)

Commands and results:

cargo test -p overlaybd --lib
  → 304 passed; 0 failed; 1 ignored

cargo test -p agentenv --lib -- memory_snapshot_compression_config_is_valid   compact_layers_mixed_input_and_configured_output   memory_runtime_suffix_compaction_uses_configured_compression   rootfs_runtime_suffix_compaction_remains_raw   convert_sparse_mem_supports_all_outputs_and_failure_cleanup   import_built_artifacts_preserves_zfile_memory_lower_bytes
  → 6 passed; 0 failed

cargo test -p agentenv --test integration -- --exact --test-threads=1   fc::memory_snapshot_format_matches_config_and_resumes   fc::legacy_memory_snapshot_format_matches_config_and_resumes
  → 2 passed; 0 failed (LZ4 workers=4, AENV_CONFIG_PATH with compression enabled)

cargo clippy -p overlaybd -p agentenv --all-targets -- -D warnings
  → zero warnings

Skipped checks and reasons:

make -C services test: no services/ changes
Generated code regeneration: no generated code changes

Risks and reviewer notes

Correctness:

WorkerPool uses bounded crossbeam channels with seq-numbered results and BTreeMap reordering. Submit+drain has no .await, so cancellation cannot strand stale results. Worker panics are caught via catch_unwind and converted to per-seq errors.
Byte-identity: LZ4/Zstd bulk compression is deterministic per block; test_verify_builder asserts workers=4 vs workers=1 produce identical files.
Cancellation-safe poisoning prevents partial writes from being reused after errors or aborts.

Performance:

pool.recv() blocks the calling Tokio task for up to one 512 KiB batch compression time (~ms). This is bounded and documented with a TODO to wrap in spawn_blocking if latency-sensitive.
process_vm_readv remains synchronous on the Tokio worker thread (existing TODO). Raw pause is unaffected; compression pause is dominated by codec CPU.

Compatibility:

LZ4 switched from lz4_flex to lz4-sys (C liblz4). Both produce standard LZ4 block format; old layers remain readable. Content-addressed caches may see different digests for the same input (LZ4 output is not canonical).
compression_workers is clamped to [1, 64] to prevent absurd configs from spawning excessive OS threads.

Key files for review:

storage/overlaybd/src/compression/zfile.rs (WorkerPool, CompressedSegment, commit_compressed_batch, tests)
src/sandbox/firecracker/overlaybd_snapshot.rs (direct memory snapshot)
src/sandbox/ublk/overlaybd.rs (compression config routing)
storage/overlaybd/src/lsmt/file/helper.rs (compact_to ordered writer, zeroed-index fix)
storage/overlaybd/src/backend/local.rs (write_at_sync_pwrite)

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/92 **Author:** [@huajq](https://github.com/huajq) **Created:** 7/31/2026 **Status:** ✅ Merged **Merged:** 8/4/2026 **Merged by:** [@yingdi-shan](https://github.com/yingdi-shan) **Base:** `main` ← **Head:** `zfile` --- ### 📝 Commits (3) - [`a7f24b4`](https://github.com/kvcache-ai/AgentENV/commit/a7f24b473a17f2feb5494d9a69d9524544937f67) feat(sandbox): add optional ZFile compression for memory snapshots - [`ccbae02`](https://github.com/kvcache-ai/AgentENV/commit/ccbae02a3c3354f47efd1b15b5a95f51a160d687) feat(overlaybd): add persistent worker pool and sync pwrite for ZFile compression - [`4b7501d`](https://github.com/kvcache-ai/AgentENV/commit/4b7501df13f5db4190c88011cded5869d06b12f2) test(sandbox): add memory snapshot compression integration tests ### 📊 Changes **22 files changed** (+2072 additions, -796 deletions) <details> <summary>View changed files</summary> 📝 `Cargo.lock` (+2 -16) 📝 `config/default.toml` (+5 -0) 📝 `config/oss_default.toml` (+5 -0) 📝 `crates/e2e-tests/Cargo.toml` (+1 -0) 📝 `crates/e2e-tests/tests/snapshot_oss_e2e_test.rs` (+69 -6) 📝 `docs/src/configuration/reference.md` (+2 -0) 📝 `src/cfg.rs` (+64 -0) 📝 `src/sandbox/firecracker/overlaybd_snapshot.rs` (+241 -50) 📝 `src/sandbox/firecracker/sandbox.rs` (+20 -9) 📝 `src/sandbox/ublk/mod.rs` (+3 -1) 📝 `src/sandbox/ublk/overlaybd.rs` (+284 -12) 📝 `src/snapshot/repository/backends/posixfs/artifacts.rs` (+114 -0) 📝 `storage/overlaybd/Cargo.toml` (+2 -1) 📝 `storage/overlaybd/src/backend/local.rs` (+48 -0) 📝 `storage/overlaybd/src/compression/zfile.rs` (+941 -638) 📝 `storage/overlaybd/src/dense_export.rs` (+4 -0) 📝 `storage/overlaybd/src/io/virtual_file.rs` (+8 -0) 📝 `storage/overlaybd/src/lsmt/file/helper.rs` (+74 -29) 📝 `storage/overlaybd/src/tools/mod.rs` (+3 -1) 📝 `storage/overlaybd/src/tools/packaging.rs` (+26 -19) _...and 2 more files_ </details> ### 📄 Description ## What Add optional ZFile compression support for Firecracker memory snapshots. When enabled, dirty memory pages are compressed during snapshot creation using LZ4 (default) or Zstd, reducing snapshot artifact size by ~70-88% at the cost of ~40-90% longer pause latency. ## Why Memory snapshots for large VMs produce multi-GB artifacts. Compressing them during creation reduces storage and transfer costs, which is especially valuable for snapshot templates shared across nodes. ## Related issue Closes # ## Scope and non-goals **Included:** - `[memory_snapshot]` config: `compression_enabled` (default false), `compression_algorithm` (lz4/zstd), `compression_workers` (default 1, max 64) - Direct memory snapshot path: Firecracker dirty ranges → `ProcessVmReader` (process_vm_readv) → `compact_to` → raw or `ZFileCompactWriter` → temp → rename - ZFile compression core: persistent `WorkerPool` (std::thread + crossbeam-channel), `CompressedSegment` (contiguous output, no per-block Vec), 512 KiB compression batch, ordered committer, cancellation-safe poisoning - `LocalFile::write_at_sync_pwrite` for >4KiB batches (bypasses io_uring overhead on local page-cache writes) - `compact_to` ordered writer support, zeroed-index sort fix (Critical), MAX_LENGTH clamp - LZ4 switched from lz4_flex to lz4-sys (C liblz4, ~25% faster) - Integration tests for direct and legacy LZ4 compressed pause/resume **Excluded:** - No changes to rootfs compression (rootfs layers remain raw) - No changes to snapshot manifest or artifact layout (ZFile is transparent to the repository) - No changes to resume path (switch file transparently decompresses ZFile) - No CI pipeline changes (integration tests run manually with compression configs) ## Design and behavior changes ### Data flow ``` Firecracker pause → create_state_only_snapshot (vm_state.bin) → GET /vm/dirty-memory-ranges → dirty_ranges_to_segment_mappings → ProcessVmReader (process_vm_readv, 512 KiB chunks) → compact_to → Raw VirtualFileWriter (unordered, io_uring) or ZFileCompactWriter (ordered, WorkerPool compression) → overlaybd.commit.tmp → rename overlaybd.commit ``` ### WorkerPool architecture ``` compact_to (async, Tokio worker) │ │ 512 KiB chunk via ProcessVmReader │ ▼ crossbeam::channel::bounded<WorkItem>(workers * 2) │ ├──► Worker Thread 1 (std::thread, owns compressor) ├──► Worker Thread 2 ├──► Worker Thread 3 └──► Worker Thread 4 │ ▼ crossbeam::channel::bounded<CompressedBatch>(workers * 2) │ ▼ Ordered Committer (BTreeMap by seq) → commit_compressed_batch → LocalFile::write_at_sync_pwrite (>4KiB) or write_all_at (non-local fallback) ``` ### Key design decisions - **workers=1 remains synchronous**: no thread overhead for single-worker configs - **512 KiB compression batch**: amortizes worker dispatch (32 KiB → 512 KiB) - **sync pwrite for large batches**: bypasses io_uring submit/completion round trips on local page-cache writes (~2-3 µs vs ~15 µs per call) - **Cancellation-safe poisoning**: any builder write/finish await pre-marks `failed=true`, cleared only on success; future cancellation or errors leave the writer permanently poisoned - **catch_unwind in workers**: prevents worker panics from deadlocking the ordered committer drain loop ### Performance (1 GiB structured text payload, 1536 MiB VM) | Mode | Pause (mean) | Artifact size | Ratio | |------|-------------|---------------|-------| | Raw | 801 ms | 1,242 MB | 1.00 | | LZ4 workers=4 | 1,127 ms | 367 MB | 0.30 | | Zstd workers=4 | 1,533 ms | 149 MB | 0.12 | LZ4 compression adds ~40% pause latency for ~70% size reduction. Zstd adds ~90% for ~88% reduction. ### Zeroed-index sort fix (Critical bug fix included) `compact_to` previously pushed zeroed `SegmentMapping` entries into `compact_index` during phase 1, then appended data entries in phase 2. This produced an unsorted index (`[all zeros..., all data...]`), causing `partition_point` lookups in `ReadOnlyIndex` to silently miss mappings and read committed data as zeros. Fixed by sorting `compact_index` before `compress_raw_index`. ## Compatibility and operations - Public API or generated protocol: N/A (no API changes) - Configuration or defaults: New `[memory_snapshot]` keys with safe defaults (disabled, lz4, 1 worker) - Snapshot manifest, artifact layout, or storage format: N/A (ZFile is transparent to the repository; resume path unchanged) - Upgrade and rollback: Old snapshots without compression continue to work. New compressed snapshots require the same code version to resume (ZFile read support is in the same crate) - Host requirements, permissions, ports, or dependencies: Adds `crossbeam-channel` dependency. No new system requirements. ## Validation - [x] `make fmt` (cargo fmt --all -- --check) - [x] `make clippy` (cargo clippy -p overlaybd -p agentenv --all-targets -- -D warnings, zero warnings) - [x] `make test-unit` (overlaybd lib 304/0/1 ignored, agentenv 6 core tests 6/6) - [x] Relevant Rust integration tests (fc::memory_snapshot_format_matches_config_and_resumes, fc::legacy_memory_snapshot_format_matches_config_and_resumes, 2/2 with LZ4 workers=4) - [ ] `make -C services test` (N/A, no services/ changes) - [ ] Generated clients/server regenerated with the documented `make` target (N/A, no generated code changes) - [x] Documentation updated (docs/src/configuration/reference.md) - [x] Benchmarks or performance comparison completed (see Performance section) Commands and results: ```text cargo test -p overlaybd --lib → 304 passed; 0 failed; 1 ignored cargo test -p agentenv --lib -- memory_snapshot_compression_config_is_valid compact_layers_mixed_input_and_configured_output memory_runtime_suffix_compaction_uses_configured_compression rootfs_runtime_suffix_compaction_remains_raw convert_sparse_mem_supports_all_outputs_and_failure_cleanup import_built_artifacts_preserves_zfile_memory_lower_bytes → 6 passed; 0 failed cargo test -p agentenv --test integration -- --exact --test-threads=1 fc::memory_snapshot_format_matches_config_and_resumes fc::legacy_memory_snapshot_format_matches_config_and_resumes → 2 passed; 0 failed (LZ4 workers=4, AENV_CONFIG_PATH with compression enabled) cargo clippy -p overlaybd -p agentenv --all-targets -- -D warnings → zero warnings ``` Skipped checks and reasons: - `make -C services test`: no `services/` changes - Generated code regeneration: no generated code changes ## Risks and reviewer notes **Correctness:** - WorkerPool uses bounded crossbeam channels with seq-numbered results and BTreeMap reordering. Submit+drain has no `.await`, so cancellation cannot strand stale results. Worker panics are caught via `catch_unwind` and converted to per-seq errors. - Byte-identity: LZ4/Zstd bulk compression is deterministic per block; `test_verify_builder` asserts workers=4 vs workers=1 produce identical files. - Cancellation-safe poisoning prevents partial writes from being reused after errors or aborts. **Performance:** - `pool.recv()` blocks the calling Tokio task for up to one 512 KiB batch compression time (~ms). This is bounded and documented with a TODO to wrap in `spawn_blocking` if latency-sensitive. - `process_vm_readv` remains synchronous on the Tokio worker thread (existing TODO). Raw pause is unaffected; compression pause is dominated by codec CPU. **Compatibility:** - LZ4 switched from `lz4_flex` to `lz4-sys` (C liblz4). Both produce standard LZ4 block format; old layers remain readable. Content-addressed caches may see different digests for the same input (LZ4 output is not canonical). - `compression_workers` is clamped to [1, 64] to prevent absurd configs from spawning excessive OS threads. **Key files for review:** - `storage/overlaybd/src/compression/zfile.rs` (WorkerPool, CompressedSegment, commit_compressed_batch, tests) - `src/sandbox/firecracker/overlaybd_snapshot.rs` (direct memory snapshot) - `src/sandbox/ublk/overlaybd.rs` (compression config routing) - `storage/overlaybd/src/lsmt/file/helper.rs` (compact_to ordered writer, zeroed-index fix) - `storage/overlaybd/src/backend/local.rs` (write_at_sync_pwrite) ## 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:20:39 +00:00