[PR #446] feat(ublk): restart and crash-recover the ublk daemon without losing devices #442

New issue

Open

opened 2026-10-29 14:23:26 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: feat/ublk-daemon-recovery

📝 Commits (3)

418a896 feat(ublk): add user recovery primitives
aafbc56 feat(ublk-daemon): recover devices across daemon restarts
9548067 feat(ublk): respawn and recover the ublk daemon from the client

📊 Changes

17 files changed (+2975 additions, -109 deletions)

View changed files

📝 Makefile (+2 -1)
📝 config/default.toml (+6 -0)
📝 docs/src/configuration/reference.md (+2 -0)
📝 docs/src/internals/architecture.md (+11 -0)
📝 src/cfg.rs (+7 -0)
📝 src/sandbox/ublk/device.rs (+4 -0)
📝 storage/ublk-daemon/src/client.rs (+663 -54)
📝 storage/ublk-daemon/src/lib.rs (+2 -1)
📝 storage/ublk-daemon/src/main.rs (+38 -6)
📝 storage/ublk-daemon/src/protocol.rs (+94 -0)
📝 storage/ublk-daemon/src/server.rs (+502 -18)
➕ storage/ublk-daemon/tests/recovery_test.rs (+1043 -0)
📝 storage/ublk-daemon/tests/ublk_daemon_test.rs (+2 -0)
📝 storage/ublk/src/ctrl.rs (+153 -0)
📝 storage/ublk/src/dev.rs (+108 -29)
📝 storage/ublk/src/ublk_caps.rs (+13 -0)
➕ tests/ublk_recovery_integration.rs (+325 -0)

📄 Description

What

Let the ublk daemon restart, deliberately or after a crash, without taking running sandboxes down. Guests see a short I/O stall instead of I/O errors, lose no acknowledged writes, and keep the same dev_id and /dev/ublkbN.

Three commits, best reviewed in order:

feat(ublk): add user recovery primitives: UBLK_F_USER_RECOVERY | UBLK_F_USER_RECOVERY_REISSUE on the ctrl builder, an owner tag in ublksrv_flags, START_USER_RECOVERY / END_USER_RECOVERY / GET_PARAMS, and UVMUblkDevBuilder::recover() / UVMUblkDev::resume().
feat(ublk-daemon): recover devices across daemon restarts: daemon Detach (flush images, exit without deleting devices) and Recover { devices } (reattach each DeviceSpec into the map that tracked it, then sweep orphaned devices), plus [ublk] recovery_enabled.
feat(ublk): respawn and recover the ublk daemon from the client: UblkDaemonClient keeps an in-memory registry of handed-over devices, restart(), automatic replacement when the daemon dies, and stress tests with real sandboxes.

Why

Today any ublk daemon exit kills every sandbox on the node: devices are created without recovery flags, so the kernel removes them, and the client marks the daemon dead for good. That blocks both crash tolerance and the online node upgrades proposed in #440.

Refs #440. This implements the ublk handover part (§4) plus crash recovery; it does not close the issue.

Scope and non-goals

In scope: whole-daemon restart, both deliberate (UblkDaemonClient::restart(), internal) and after a crash (automatic replacement), for raw, pool-exclusive, and pool-shared devices.

Not in scope:

An operator-facing trigger (admin endpoint, signal, A/B containers). restart() is internal and exercised by tests.
Persisting the device registry. It lives only in the node server's memory, so a node-server crash still tears everything down as today. Persisting DeviceSpec in the per-sandbox RuntimeRecord (#440 §2) is the follow-up that adds server crash consistency.
Per-device QUIESCE handover; this stops all devices once for a whole-daemon restart.
Crash-loop backoff: a daemon that crashes right after every recovery is replaced every time.

Design and behavior changes

No persistent ublk state. The client's registry of DeviceSpec { dev_id, image_config, global_config, role } is the only record of what to recover. It is updated inside the RPC that hands a device over or releases it, under a call gate (tokio::sync::RwLock): RPCs hold it shared, replacing the daemon holds it exclusively, so a restart never snapshots a registry that misses a device.
Restart: take the gate, Detach (abort pack recordings, sync() every image, exit without STOP/DEL). Closing /dev/ublkcN quiesces the devices and the kernel requeues in-flight I/O. Then spawn a new daemon and Recover the registry.
Crash: each watchdog records which spawn it watches. After its daemon exits (other than at shutdown) it waits for the gate; if its daemon is still the current one, it replaces it, retrying with backoff. This also covers a replacement that dies right after recovering, even in the middle of restart(). If no replacement starts, the client force-deletes every registered device so guests get EIO rather than hanging.
Recover: devices are recovered in parallel, one per CPU by default ([ublk] recovery_concurrency overrides it), and each device resumes as soon as its own recovery finishes. A device that cannot be recovered (for example one created without recovery flags) is deleted and reported in failed. A device that does not carry this instance's owner tag is reported in failed but never touched, so a stale entry whose id the kernel reused cannot delete another instance's device.
Orphan sweep: after recovery, a background sweep deletes devices that carry this daemon's owner tag (an FNV-1a hash of the socket path; the kernel resets ublksrv_pid on exit, so the pid is useless), are not live, and were not recovered. Examples are idle pool devices and creates whose response never reached the client. Only devices that existed when the daemon started are candidates, so devices it is creating are never swept, and a graceful shutdown waits for the sweep.
Durability invariant: acknowledged writes survive a kill -9 only because the LSMT group_commit_size stays 0. This is documented in docs/src/internals/architecture.md.

Compatibility and operations

Public API or generated protocol: no public API change. The daemon's private socket protocol gains Detach, Recover, and Recovered; the client and daemon ship together.
Configuration or defaults: new [ublk] recovery_concurrency, which sets how many devices a replacement daemon recovers at once; unset or 0 means one per CPU. New [ublk] recovery_enabled, default true. Recovery is used only when the kernel supports USER_RECOVERY_REISSUE (Linux 6.0+), which is detected at runtime. With recovery enabled, I/O to a device whose daemon is gone queues instead of failing until the daemon is replaced or the client gives up and deletes the device.
Snapshot manifest, artifact layout, or storage format: N/A.
Upgrade and rollback: devices created by a daemon without this change cannot be recovered, so the first rollout of this change needs nodes drained before any daemon restart. Rolling back is safe: an older daemon ignores the flags on existing devices.
Host requirements, permissions, ports, or dependencies: none new. Recovery needs Linux 6.0+ and is otherwise off automatically.

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
Benchmarks or performance comparison completed

Commands and results (Linux 6.17, KVM, 16 cores, 31 GiB):

make fmt / make clippy                                   clean (HEAD; commit 2 also checked on its own)
make test-unit                                           all suites pass
cargo test -p uvm-ublk -p uvm-ublk-daemon --lib          pass
cargo test -p uvm-ublk-daemon --test recovery_test       8 passed (privileged)
  O_DIRECT writers, longest single write across a restart:
    graceful detach ~35 ms, kill -9 ~52-60 ms,
    client restart() ~40-65 ms, crash + automatic respawn ~54 ms
random_crash_sequences_keep_every_device soak            64/64 seeds x 40 steps pass
  kills, double kills (killing the replacement mid-recovery), restarts racing
  kills, creates racing kills, add/remove; up to 8 devices under writes plus a
  shared device under verifying reads; asserts no EIO, no lost writes, exact
  shared refcount, and no leaked owner-tagged devices
ublk_recovery_integration (real Firecracker sandboxes)   pass at 4, 8, 10, 24 and 48 sandboxes
  48-sandbox run: seed 0x20, 20 rounds, grew to 51 sandboxes / 56 ublk devices,
  every replacement recovered all devices (failed=0), 4 replacements killed
  mid-recovery and retried, crash-to-recovered 233 / 316 / 1367 ms
  (min / median / max), every guest's checksums verified from the device with
  O_DIRECT and no I/O or ext4 errors in any guest's dmesg

Skipped checks and reasons:

make -C services test: no services/ changes.
No generated code changed.
Not run in full: the other suites of make test-agent-integration and make test-ublk. The new tests in those targets were run directly with the capability runner.

Risks and reviewer notes

The concurrency is in storage/ublk-daemon/src/client.rs: the call gate, RegistryUpdate, spawn_watchdog's generation check, restart(), respawn_and_recover_once(), and fail_permanently().
In the daemon, see handle_recover, recover_device (role placement), and sweep_orphaned_devices in storage/ublk-daemon/src/server.rs.
The random tests found two bugs that are fixed here: a replacement that died right after recovering was never replaced, and the orphan sweep could be cut short by shutdown and could have deleted a device the new daemon was still creating.
At 48 sandboxes with an unpaced guest workload on a 16-core host (load average ~100), a guest exec occasionally never completed. A control run with no daemon crashes hit the same hang, so it is unrelated to recovery. The integration test now paces its guest loops and bounds every sandbox operation with a timeout, so a hang fails the test with the operation named instead of stalling it.
ublk_recovery_integration is a separate test binary because it kills the process-wide daemon. It is added to make test-agent-integration.

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.

🤖 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/446 **Author:** [@yingdi-shan](https://github.com/yingdi-shan) **Created:** 10/28/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/ublk-daemon-recovery` --- ### 📝 Commits (3) - [`418a896`](https://github.com/kvcache-ai/AgentENV/commit/418a896f37c29794445fcbb505cdbf116d6675a0) feat(ublk): add user recovery primitives - [`aafbc56`](https://github.com/kvcache-ai/AgentENV/commit/aafbc566532fad37ef6f069deb8e0b8345f2bbfd) feat(ublk-daemon): recover devices across daemon restarts - [`9548067`](https://github.com/kvcache-ai/AgentENV/commit/954806746a9972809a72afe67c9965173385670a) feat(ublk): respawn and recover the ublk daemon from the client ### 📊 Changes **17 files changed** (+2975 additions, -109 deletions) <details> <summary>View changed files</summary> 📝 `Makefile` (+2 -1) 📝 `config/default.toml` (+6 -0) 📝 `docs/src/configuration/reference.md` (+2 -0) 📝 `docs/src/internals/architecture.md` (+11 -0) 📝 `src/cfg.rs` (+7 -0) 📝 `src/sandbox/ublk/device.rs` (+4 -0) 📝 `storage/ublk-daemon/src/client.rs` (+663 -54) 📝 `storage/ublk-daemon/src/lib.rs` (+2 -1) 📝 `storage/ublk-daemon/src/main.rs` (+38 -6) 📝 `storage/ublk-daemon/src/protocol.rs` (+94 -0) 📝 `storage/ublk-daemon/src/server.rs` (+502 -18) ➕ `storage/ublk-daemon/tests/recovery_test.rs` (+1043 -0) 📝 `storage/ublk-daemon/tests/ublk_daemon_test.rs` (+2 -0) 📝 `storage/ublk/src/ctrl.rs` (+153 -0) 📝 `storage/ublk/src/dev.rs` (+108 -29) 📝 `storage/ublk/src/ublk_caps.rs` (+13 -0) ➕ `tests/ublk_recovery_integration.rs` (+325 -0) </details> ### 📄 Description ## What Let the ublk daemon restart, deliberately or after a crash, without taking running sandboxes down. Guests see a short I/O stall instead of I/O errors, lose no acknowledged writes, and keep the same `dev_id` and `/dev/ublkbN`. Three commits, best reviewed in order: 1. **`feat(ublk): add user recovery primitives`**: `UBLK_F_USER_RECOVERY | UBLK_F_USER_RECOVERY_REISSUE` on the ctrl builder, an owner tag in `ublksrv_flags`, `START_USER_RECOVERY` / `END_USER_RECOVERY` / `GET_PARAMS`, and `UVMUblkDevBuilder::recover()` / `UVMUblkDev::resume()`. 2. **`feat(ublk-daemon): recover devices across daemon restarts`**: daemon `Detach` (flush images, exit without deleting devices) and `Recover { devices }` (reattach each `DeviceSpec` into the map that tracked it, then sweep orphaned devices), plus `[ublk] recovery_enabled`. 3. **`feat(ublk): respawn and recover the ublk daemon from the client`**: `UblkDaemonClient` keeps an in-memory registry of handed-over devices, `restart()`, automatic replacement when the daemon dies, and stress tests with real sandboxes. ## Why Today any ublk daemon exit kills every sandbox on the node: devices are created without recovery flags, so the kernel removes them, and the client marks the daemon dead for good. That blocks both crash tolerance and the online node upgrades proposed in #440. ## Related issue Refs #440. This implements the ublk handover part (§4) plus crash recovery; it does not close the issue. ## Scope and non-goals In scope: whole-daemon restart, both deliberate (`UblkDaemonClient::restart()`, internal) and after a crash (automatic replacement), for raw, pool-exclusive, and pool-shared devices. Not in scope: - An operator-facing trigger (admin endpoint, signal, A/B containers). `restart()` is internal and exercised by tests. - Persisting the device registry. It lives only in the node server's memory, so a node-server crash still tears everything down as today. Persisting `DeviceSpec` in the per-sandbox `RuntimeRecord` (#440 §2) is the follow-up that adds server crash consistency. - Per-device `QUIESCE` handover; this stops all devices once for a whole-daemon restart. - Crash-loop backoff: a daemon that crashes right after every recovery is replaced every time. ## Design and behavior changes - **No persistent ublk state.** The client's registry of `DeviceSpec { dev_id, image_config, global_config, role }` is the only record of what to recover. It is updated inside the RPC that hands a device over or releases it, under a call gate (`tokio::sync::RwLock`): RPCs hold it shared, replacing the daemon holds it exclusively, so a restart never snapshots a registry that misses a device. - **Restart:** take the gate, `Detach` (abort pack recordings, `sync()` every image, exit without STOP/DEL). Closing `/dev/ublkcN` quiesces the devices and the kernel requeues in-flight I/O. Then spawn a new daemon and `Recover` the registry. - **Crash:** each watchdog records which spawn it watches. After its daemon exits (other than at shutdown) it waits for the gate; if its daemon is still the current one, it replaces it, retrying with backoff. This also covers a replacement that dies right after recovering, even in the middle of `restart()`. If no replacement starts, the client force-deletes every registered device so guests get EIO rather than hanging. - **Recover:** devices are recovered in parallel, one per CPU by default (`[ublk] recovery_concurrency` overrides it), and each device resumes as soon as its own recovery finishes. A device that cannot be recovered (for example one created without recovery flags) is deleted and reported in `failed`. A device that does not carry this instance's owner tag is reported in `failed` but never touched, so a stale entry whose id the kernel reused cannot delete another instance's device. - **Orphan sweep:** after recovery, a background sweep deletes devices that carry this daemon's owner tag (an FNV-1a hash of the socket path; the kernel resets `ublksrv_pid` on exit, so the pid is useless), are not live, and were not recovered. Examples are idle pool devices and creates whose response never reached the client. Only devices that existed when the daemon started are candidates, so devices it is creating are never swept, and a graceful shutdown waits for the sweep. - **Durability invariant:** acknowledged writes survive a `kill -9` only because the LSMT `group_commit_size` stays 0. This is documented in `docs/src/internals/architecture.md`. ## Compatibility and operations - Public API or generated protocol: no public API change. The daemon's private socket protocol gains `Detach`, `Recover`, and `Recovered`; the client and daemon ship together. - Configuration or defaults: new `[ublk] recovery_concurrency`, which sets how many devices a replacement daemon recovers at once; unset or `0` means one per CPU. New `[ublk] recovery_enabled`, default `true`. Recovery is used only when the kernel supports `USER_RECOVERY_REISSUE` (Linux 6.0+), which is detected at runtime. With recovery enabled, I/O to a device whose daemon is gone queues instead of failing until the daemon is replaced or the client gives up and deletes the device. - Snapshot manifest, artifact layout, or storage format: N/A. - Upgrade and rollback: devices created by a daemon without this change cannot be recovered, so the first rollout of this change needs nodes drained before any daemon restart. Rolling back is safe: an older daemon ignores the flags on existing devices. - Host requirements, permissions, ports, or dependencies: none new. Recovery needs Linux 6.0+ and is otherwise off automatically. ## Validation - [x] `make fmt` - [x] `make clippy` - [x] `make test-unit` - [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 (Linux 6.17, KVM, 16 cores, 31 GiB): ```text make fmt / make clippy clean (HEAD; commit 2 also checked on its own) make test-unit all suites pass cargo test -p uvm-ublk -p uvm-ublk-daemon --lib pass cargo test -p uvm-ublk-daemon --test recovery_test 8 passed (privileged) O_DIRECT writers, longest single write across a restart: graceful detach ~35 ms, kill -9 ~52-60 ms, client restart() ~40-65 ms, crash + automatic respawn ~54 ms random_crash_sequences_keep_every_device soak 64/64 seeds x 40 steps pass kills, double kills (killing the replacement mid-recovery), restarts racing kills, creates racing kills, add/remove; up to 8 devices under writes plus a shared device under verifying reads; asserts no EIO, no lost writes, exact shared refcount, and no leaked owner-tagged devices ublk_recovery_integration (real Firecracker sandboxes) pass at 4, 8, 10, 24 and 48 sandboxes 48-sandbox run: seed 0x20, 20 rounds, grew to 51 sandboxes / 56 ublk devices, every replacement recovered all devices (failed=0), 4 replacements killed mid-recovery and retried, crash-to-recovered 233 / 316 / 1367 ms (min / median / max), every guest's checksums verified from the device with O_DIRECT and no I/O or ext4 errors in any guest's dmesg ``` Skipped checks and reasons: - `make -C services test`: no `services/` changes. - No generated code changed. - Not run in full: the other suites of `make test-agent-integration` and `make test-ublk`. The new tests in those targets were run directly with the capability runner. ## Risks and reviewer notes - The concurrency is in `storage/ublk-daemon/src/client.rs`: the call gate, `RegistryUpdate`, `spawn_watchdog`'s generation check, `restart()`, `respawn_and_recover_once()`, and `fail_permanently()`. - In the daemon, see `handle_recover`, `recover_device` (role placement), and `sweep_orphaned_devices` in `storage/ublk-daemon/src/server.rs`. - The random tests found two bugs that are fixed here: a replacement that died right after recovering was never replaced, and the orphan sweep could be cut short by shutdown and could have deleted a device the new daemon was still creating. - At 48 sandboxes with an unpaced guest workload on a 16-core host (load average ~100), a guest exec occasionally never completed. A control run with no daemon crashes hit the same hang, so it is unrelated to recovery. The integration test now paces its guest loops and bounds every sandbox operation with a timeout, so a hang fails the test with the operation named instead of stalling it. - `ublk_recovery_integration` is a separate test binary because it kills the process-wide daemon. It is added to `make test-agent-integration`. ## 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. 🤖 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:23:26 +00:00