kvcache-ai/AgentENV

Fork 0

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

[PR #389] feat(ublk): add an NBD block device backend for hosts without ublk #398

New issue

Open

opened 2026-10-29 14:22:59 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:22:59 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/AgentENV/pull/389
Author: @yingdi-shan
Created: 10/18/2026
Status: 🔄 Open

Base: main ← Head: feat/nbd-device-backend

📝 Commits (2)

88e2c2a fix(ublk-daemon): drop cached pages when reusing a pooled device
3f9eb1b feat(ublk): add an NBD block device backend for hosts without ublk

📊 Changes

34 files changed (+2661 additions, -458 deletions)

View changed files

📝 Cargo.lock (+14 -0)
📝 Cargo.toml (+1 -0)
📝 Makefile (+2 -2)
📝 config/default.toml (+3 -0)
📝 deploy/docker-compose.yml (+1 -0)
📝 docs/src/configuration/env-vars.md (+1 -0)
📝 docs/src/configuration/reference.md (+11 -0)
📝 docs/src/deployment/docker-compose.md (+5 -1)
📝 docs/src/deployment/docker.md (+5 -2)
📝 docs/src/getting-started/quickstart.md (+8 -4)
📝 scripts/docker-setup.sh (+64 -19)
📝 scripts/install.sh (+57 -7)
📝 src/cfg.rs (+23 -0)
📝 src/sandbox/ublk/device.rs (+4 -0)
📝 src/setup/mod.rs (+17 -5)
➕ src/setup/nbd.rs (+61 -0)
📝 src/setup/ublk.rs (+16 -62)
➕ src/setup/udev.rs (+66 -0)
➕ storage/nbd/Cargo.toml (+14 -0)
➕ storage/nbd/src/kernel.rs (+576 -0)

...and 14 more files

📄 Description

What

Adds NBD as a second kernel frontend for OverlayBD devices, next to ublk, and lets the installer and docker setup script choose between them.

This PR is based directly on main and is independent of #388. It has two commits:

88e2c2a fix(ublk-daemon): drop cached pages when reusing a pooled device (Refs #302). This is a standalone 5-line commit that can be cherry-picked on its own.
89ab4cc feat(ublk): add an NBD block device backend for hosts without ublk

Why

OverlayBD devices currently require ublk_drv (Linux 6.0+, and in practice 6.8+). That rules out older distro kernels. The kernel NBD client is available almost everywhere and is enough to serve the same OverlaybdTarget. FUSE adds no benefit over NBD here, and vhost-user-blk / NVMe-oF would break Firecracker's MAP_PRIVATE memory restore and drive model.

Refs #302. The pool fix in 88e2c2a addresses the stale placeholder pages that #342 reproduces. NBD made the bug easy to hit. It narrows the race but does not prove it closed; see the Risks section.

Scope and non-goals

Included:

New crate storage/nbd (uvm-nbd). It implements the NBD transmission phase (READ/WRITE/TRIM/DISC) and attaches devices through the generic-netlink nbd family over a Unix socketpair. There is no handshake and no TCP.
OverlaybdTarget implements NbdExport. Reads still feed the startup-pack recorder.
The daemon selects the frontend with --device-backend ublk|nbd through a shared DeviceBackend / BlockDevice. The ublk-specific helpers moved from server.rs to device.rs without behavior changes.
New config key [ublk] device_backend / AENV_UBLK_DEVICE_BACKEND, wired into the daemon spawn and into server --setup-host / startup checks (src/setup/nbd.rs).
scripts/install.sh and scripts/docker-setup.sh take AENV_UBLK_DEVICE_BACKEND=auto|ublk|nbd, defaulting to auto.

Not included:

macOS. The protocol and server modules are portable, but the kernel attach is Linux-only.
NBD over TCP or multiple connections, device resize, and a server-side auto mode. The server default stays ublk; the scripts make the choice.

Design and behavior changes

Write-through semantics. ublk sets no UBLK_ATTR_VOLATILE_CACHE, so the kernel never sends it FLUSH. NBD exports are likewise write-through: the server never advertises SEND_FLUSH/SEND_FUA and rejects FLUSH with EINVAL. Read-only OverlayBD images reject sync(), and advertising flush caused EIO on pool release (BLKFLSBUF).
Device lifecycle. The server task is spawned before NBD_CMD_CONNECT. disconnect sends NBD_CMD_DISCONNECT and waits up to 5 s for the server, then aborts it. Drop disconnects best-effort. No I/O timeout is set, matching ublk's behavior of never timing out against a slow lazy-loading image.
Warm pool. NBD reports no UBLK_F_UPDATE_SIZE, so GetFeatures returns 0. The pool therefore reuses only same-size devices, and it does not prewarm unless startup_prewarm is set.
Scripts.
- auto keeps ublk when ublk_drv is loaded or installable (modinfo, then linux-modules-extra via apt) and otherwise falls back to nbd.
- install.sh records the result in /etc/default/aenv next to AENV_VIRTUALIZATION_MODE and passes it to --setup-only / --setup-host.
- docker-setup.sh loads the matching module and prints the -e AENV_UBLK_DEVICE_BACKEND=nbd / AENV_UBLK_DEVICE_BACKEND=nbd make deploy-up hint. deploy/docker-compose.yml forwards the variable.

Compatibility and operations

Public API or generated protocol: N/A. The REST API is unchanged. The daemon gains a --device-backend flag, and its existing constructors keep their signatures.
Configuration or defaults: new [ublk] device_backend (default ublk) and AENV_UBLK_DEVICE_BACKEND. Installers default to auto, which resolves to ublk on every host that works today.
Snapshot manifest, artifact layout, or storage format: N/A. Both frontends serve the same OverlayBD images and layers.
Upgrade and rollback: re-running install.sh on an existing ublk host selects ublk and adds AENV_UBLK_DEVICE_BACKEND="ublk" to /etc/default/aenv. To roll back, drop that line.
Host requirements, permissions, ports, or dependencies: the NBD backend needs the nbd module (CONFIG_BLK_DEV_NBD) and CAP_SYS_ADMIN for netlink, which the service already has. --setup-host installs /etc/udev/rules.d/99-agentenv-nbd.rules and /etc/modules-load.d/aenv-nbd.conf.

Validation

make fmt
make clippy
make test-unit (ran the affected crates individually; see below)
Relevant Rust integration tests
make -C services test (required when services/ changes)
Generated clients/server regenerated with the documented make target
Documentation updated
Benchmarks or performance comparison completed

Commands and results:

# Linux 6.17, x86_64, non-root user with scripts/run-with-capabilities.sh
CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER=$PWD/scripts/run-with-capabilities.sh \
  cargo test -p uvm-nbd -p uvm-ublk -p uvm-ublk-daemon
  uvm-nbd: 17 unit + 2 real-kernel device tests passed
  uvm-ublk: 19 lib + 4 overlaybd tests passed (incl. NBD export round trip)
  uvm-ublk-daemon: 49 lib + 2 NBD backend tests passed (raw device; pool reuse across images)
cargo test -p agentenv --lib                       1002 passed, 4 ignored (incl. device_backend config test)

# Firecracker integration with the NBD backend
AENV_UBLK_DEVICE_BACKEND=nbd cargo test -p agentenv --test integration -- --test-threads=1 \
  fc::memory_snapshot_format_matches_config_and_resumes \
  fc::microvm_lifecycle_and_snapshot_preserve_disk_state \
  fc::multiple_resumes_have_independent_disk_state \
  fc::oci_tools_share_a_readonly_device_across_concurrent_launches
  4 passed; daemon log shows device_backend="nbd" and no failed NBD requests
# Same 4 tests with the default backend (device_backend="ublk")
  4 passed

# Scripts
bash -n + shellcheck scripts/install.sh scripts/docker-setup.sh   clean
bash scripts/tests/verify-install-service.sh                       ok
install.sh auto selection with stubbed modinfo/apt-get: ublk present -> ublk; absent -> nbd
docker-setup.sh with stubbed modprobe: auto -> nbd fallback; ublk -> error; nbd -> nbd

Skipped checks and reasons:

make test-unit / make test as aggregate targets were not run. The affected crates and the agentenv lib suite were run directly.
The NBD pool test failed intermittently before 88e2c2a (1 in 12 runs) and passed 25 of 25 after it. The ublk path is covered only by the existing suites. #342's reproducer was not run against this branch.
The full integration suite was not run on this branch. Only the 4 Firecracker tests above were run, once with each backend.
Not tested on a real pre-6.0 kernel; this box is Linux 6.17 with both modules. Installer and docker-setup scripts were tested with stubs, not by running them end to end on a fresh host.
services/ is unchanged.

Risks and reviewer notes

#302 fix is partial. clear_page_cache after the target swap drops pages cached while the device was idle. A placeholder read already in flight during the swap, or a reader holding the device open across reuse, could still repopulate the cache. Treat the fix as narrowing the window.
Where to look first. The most important files are storage/nbd/src/kernel.rs (netlink attach/detach and Drop), storage/nbd/src/server.rs (request concurrency, disconnect draining) and storage/ublk-daemon/src/device.rs.
NBD costs. It adds a socket copy per request compared with ublk's shared buffers. Expect lower throughput; no benchmarks were done.
Env overrides config. On installed hosts, the env var in /etc/default/aenv takes precedence over config.toml, the same as AENV_VIRTUALIZATION_MODE. This is documented in the configuration reference.

Checklist

The PR contains one coherent change and no unrelated formatting or refactoring. (The #302 fix is a separate commit; the ublk helper move in device.rs is required by the backend split.)
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.

🤖 Generated with Claude Code

_{🔄 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/389 **Author:** [@yingdi-shan](https://github.com/yingdi-shan) **Created:** 10/18/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/nbd-device-backend` --- ### 📝 Commits (2) - [`88e2c2a`](https://github.com/kvcache-ai/AgentENV/commit/88e2c2a0c04a31c6e5d1e759dc1a5128d2f24a98) fix(ublk-daemon): drop cached pages when reusing a pooled device - [`3f9eb1b`](https://github.com/kvcache-ai/AgentENV/commit/3f9eb1b7cfc9e7bb3bf4380929dbbb481d551471) feat(ublk): add an NBD block device backend for hosts without ublk ### 📊 Changes **34 files changed** (+2661 additions, -458 deletions) <details> <summary>View changed files</summary> 📝 `Cargo.lock` (+14 -0) 📝 `Cargo.toml` (+1 -0) 📝 `Makefile` (+2 -2) 📝 `config/default.toml` (+3 -0) 📝 `deploy/docker-compose.yml` (+1 -0) 📝 `docs/src/configuration/env-vars.md` (+1 -0) 📝 `docs/src/configuration/reference.md` (+11 -0) 📝 `docs/src/deployment/docker-compose.md` (+5 -1) 📝 `docs/src/deployment/docker.md` (+5 -2) 📝 `docs/src/getting-started/quickstart.md` (+8 -4) 📝 `scripts/docker-setup.sh` (+64 -19) 📝 `scripts/install.sh` (+57 -7) 📝 `src/cfg.rs` (+23 -0) 📝 `src/sandbox/ublk/device.rs` (+4 -0) 📝 `src/setup/mod.rs` (+17 -5) ➕ `src/setup/nbd.rs` (+61 -0) 📝 `src/setup/ublk.rs` (+16 -62) ➕ `src/setup/udev.rs` (+66 -0) ➕ `storage/nbd/Cargo.toml` (+14 -0) ➕ `storage/nbd/src/kernel.rs` (+576 -0) _...and 14 more files_ </details> ### 📄 Description ## What Adds NBD as a second kernel frontend for OverlayBD devices, next to ublk, and lets the installer and docker setup script choose between them. This PR is based directly on `main` and is independent of #388. It has two commits: - `88e2c2a fix(ublk-daemon): drop cached pages when reusing a pooled device` (Refs #302). This is a standalone 5-line commit that can be cherry-picked on its own. - `89ab4cc feat(ublk): add an NBD block device backend for hosts without ublk` ## Why OverlayBD devices currently require `ublk_drv` (Linux 6.0+, and in practice 6.8+). That rules out older distro kernels. The kernel NBD client is available almost everywhere and is enough to serve the same `OverlaybdTarget`. FUSE adds no benefit over NBD here, and vhost-user-blk / NVMe-oF would break Firecracker's `MAP_PRIVATE` memory restore and drive model. ## Related issue Refs #302. The pool fix in `88e2c2a` addresses the stale placeholder pages that #342 reproduces. NBD made the bug easy to hit. It narrows the race but does not prove it closed; see the Risks section. ## Scope and non-goals Included: - New crate `storage/nbd` (`uvm-nbd`). It implements the NBD transmission phase (READ/WRITE/TRIM/DISC) and attaches devices through the generic-netlink `nbd` family over a Unix socketpair. There is no handshake and no TCP. - `OverlaybdTarget` implements `NbdExport`. Reads still feed the startup-pack recorder. - The daemon selects the frontend with `--device-backend ublk|nbd` through a shared `DeviceBackend` / `BlockDevice`. The ublk-specific helpers moved from `server.rs` to `device.rs` without behavior changes. - New config key `[ublk] device_backend` / `AENV_UBLK_DEVICE_BACKEND`, wired into the daemon spawn and into `server --setup-host` / startup checks (`src/setup/nbd.rs`). - `scripts/install.sh` and `scripts/docker-setup.sh` take `AENV_UBLK_DEVICE_BACKEND=auto|ublk|nbd`, defaulting to `auto`. Not included: - macOS. The protocol and server modules are portable, but the kernel attach is Linux-only. - NBD over TCP or multiple connections, device resize, and a server-side `auto` mode. The server default stays `ublk`; the scripts make the choice. ## Design and behavior changes - **Write-through semantics.** ublk sets no `UBLK_ATTR_VOLATILE_CACHE`, so the kernel never sends it FLUSH. NBD exports are likewise write-through: the server never advertises `SEND_FLUSH`/`SEND_FUA` and rejects FLUSH with EINVAL. Read-only OverlayBD images reject `sync()`, and advertising flush caused EIO on pool release (`BLKFLSBUF`). - **Device lifecycle.** The server task is spawned before `NBD_CMD_CONNECT`. `disconnect` sends `NBD_CMD_DISCONNECT` and waits up to 5 s for the server, then aborts it. `Drop` disconnects best-effort. No I/O timeout is set, matching ublk's behavior of never timing out against a slow lazy-loading image. - **Warm pool.** NBD reports no `UBLK_F_UPDATE_SIZE`, so `GetFeatures` returns 0. The pool therefore reuses only same-size devices, and it does not prewarm unless `startup_prewarm` is set. - **Scripts.** - `auto` keeps ublk when `ublk_drv` is loaded or installable (`modinfo`, then `linux-modules-extra` via apt) and otherwise falls back to nbd. - `install.sh` records the result in `/etc/default/aenv` next to `AENV_VIRTUALIZATION_MODE` and passes it to `--setup-only` / `--setup-host`. - `docker-setup.sh` loads the matching module and prints the `-e AENV_UBLK_DEVICE_BACKEND=nbd` / `AENV_UBLK_DEVICE_BACKEND=nbd make deploy-up` hint. `deploy/docker-compose.yml` forwards the variable. ## Compatibility and operations - Public API or generated protocol: N/A. The REST API is unchanged. The daemon gains a `--device-backend` flag, and its existing constructors keep their signatures. - Configuration or defaults: new `[ublk] device_backend` (default `ublk`) and `AENV_UBLK_DEVICE_BACKEND`. Installers default to `auto`, which resolves to `ublk` on every host that works today. - Snapshot manifest, artifact layout, or storage format: N/A. Both frontends serve the same OverlayBD images and layers. - Upgrade and rollback: re-running `install.sh` on an existing ublk host selects `ublk` and adds `AENV_UBLK_DEVICE_BACKEND="ublk"` to `/etc/default/aenv`. To roll back, drop that line. - Host requirements, permissions, ports, or dependencies: the NBD backend needs the `nbd` module (`CONFIG_BLK_DEV_NBD`) and `CAP_SYS_ADMIN` for netlink, which the service already has. `--setup-host` installs `/etc/udev/rules.d/99-agentenv-nbd.rules` and `/etc/modules-load.d/aenv-nbd.conf`. ## Validation - [x] `make fmt` - [x] `make clippy` - [ ] `make test-unit` (ran the affected crates individually; see below) - [x] Relevant Rust integration tests - [ ] `make -C services test` (required when `services/` changes) - [ ] Generated clients/server regenerated with the documented `make` target - [x] Documentation updated - [ ] Benchmarks or performance comparison completed Commands and results: ```text # Linux 6.17, x86_64, non-root user with scripts/run-with-capabilities.sh CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER=$PWD/scripts/run-with-capabilities.sh \ cargo test -p uvm-nbd -p uvm-ublk -p uvm-ublk-daemon uvm-nbd: 17 unit + 2 real-kernel device tests passed uvm-ublk: 19 lib + 4 overlaybd tests passed (incl. NBD export round trip) uvm-ublk-daemon: 49 lib + 2 NBD backend tests passed (raw device; pool reuse across images) cargo test -p agentenv --lib 1002 passed, 4 ignored (incl. device_backend config test) # Firecracker integration with the NBD backend AENV_UBLK_DEVICE_BACKEND=nbd cargo test -p agentenv --test integration -- --test-threads=1 \ fc::memory_snapshot_format_matches_config_and_resumes \ fc::microvm_lifecycle_and_snapshot_preserve_disk_state \ fc::multiple_resumes_have_independent_disk_state \ fc::oci_tools_share_a_readonly_device_across_concurrent_launches 4 passed; daemon log shows device_backend="nbd" and no failed NBD requests # Same 4 tests with the default backend (device_backend="ublk") 4 passed # Scripts bash -n + shellcheck scripts/install.sh scripts/docker-setup.sh clean bash scripts/tests/verify-install-service.sh ok install.sh auto selection with stubbed modinfo/apt-get: ublk present -> ublk; absent -> nbd docker-setup.sh with stubbed modprobe: auto -> nbd fallback; ublk -> error; nbd -> nbd ``` Skipped checks and reasons: - `make test-unit` / `make test` as aggregate targets were not run. The affected crates and the `agentenv` lib suite were run directly. - The NBD pool test failed intermittently before `88e2c2a` (1 in 12 runs) and passed 25 of 25 after it. The ublk path is covered only by the existing suites. #342's reproducer was not run against this branch. - The full integration suite was not run on this branch. Only the 4 Firecracker tests above were run, once with each backend. - Not tested on a real pre-6.0 kernel; this box is Linux 6.17 with both modules. Installer and docker-setup scripts were tested with stubs, not by running them end to end on a fresh host. - `services/` is unchanged. ## Risks and reviewer notes - **#302 fix is partial.** `clear_page_cache` after the target swap drops pages cached while the device was idle. A placeholder read already in flight during the swap, or a reader holding the device open across reuse, could still repopulate the cache. Treat the fix as narrowing the window. - **Where to look first.** The most important files are `storage/nbd/src/kernel.rs` (netlink attach/detach and Drop), `storage/nbd/src/server.rs` (request concurrency, disconnect draining) and `storage/ublk-daemon/src/device.rs`. - **NBD costs.** It adds a socket copy per request compared with ublk's shared buffers. Expect lower throughput; no benchmarks were done. - **Env overrides config.** On installed hosts, the env var in `/etc/default/aenv` takes precedence over `config.toml`, the same as `AENV_VIRTUALIZATION_MODE`. This is documented in the configuration reference. ## Checklist - [x] The PR contains one coherent change and no unrelated formatting or refactoring. (The #302 fix is a separate commit; the ublk helper move in `device.rs` is required by the backend split.) - [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. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 14:22:59 +00:00

bsakaguchi referenced this issue

2026-10-29 14:23:02 +00:00

[PR #398] [MERGED] feat: attach image drives to warm starts #405

bsakaguchi referenced this issue

2026-10-29 14:23:05 +00:00

[PR #405] [MERGED] fix: reject explicit volume mounts on snapshots with volumes #408

bsakaguchi referenced this issue

2026-10-29 14:23:05 +00:00