1
0
forked from wrenn/wrenn
Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev>

Reviewed-on: wrenn/wrenn#50
This commit is contained in:
2026-05-24 21:10:37 +00:00
parent 4707f16c76
commit 05ddf62399
203 changed files with 15815 additions and 9344 deletions

417
envd-rs/Cargo.lock generated
View File

@ -17,6 +17,15 @@ dependencies = [
"memchr",
]
[[package]]
name = "android_system_properties"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
dependencies = [
"libc",
]
[[package]]
name = "anstream"
version = "1.0.0"
@ -113,6 +122,12 @@ version = "1.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
[[package]]
name = "autocfg"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
[[package]]
name = "axum"
version = "0.8.9"
@ -280,6 +295,17 @@ version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
[[package]]
name = "chrono"
version = "0.4.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
dependencies = [
"iana-time-zone",
"num-traits",
"windows-link",
]
[[package]]
name = "clap"
version = "4.6.1"
@ -486,17 +512,6 @@ dependencies = [
"subtle",
]
[[package]]
name = "displaydoc"
version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "either"
version = "1.15.0"
@ -514,7 +529,7 @@ dependencies = [
[[package]]
name = "envd"
version = "0.2.1"
version = "0.3.0"
dependencies = [
"async-stream",
"axum",
@ -522,6 +537,7 @@ dependencies = [
"buffa",
"buffa-types",
"bytes",
"chrono",
"clap",
"connectrpc",
"connectrpc-build",
@ -537,7 +553,6 @@ dependencies = [
"mime_guess",
"nix",
"notify",
"reqwest",
"serde",
"serde_json",
"sha2",
@ -889,7 +904,6 @@ dependencies = [
"pin-project-lite",
"smallvec",
"tokio",
"want",
]
[[package]]
@ -898,103 +912,37 @@ version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
dependencies = [
"base64",
"bytes",
"futures-channel",
"futures-util",
"http",
"http-body",
"hyper",
"ipnet",
"libc",
"percent-encoding",
"pin-project-lite",
"socket2",
"tokio",
"tower-service",
"tracing",
]
[[package]]
name = "icu_collections"
version = "2.2.0"
name = "iana-time-zone"
version = "0.1.65"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c"
checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
dependencies = [
"displaydoc",
"potential_utf",
"utf8_iter",
"yoke",
"zerofrom",
"zerovec",
"android_system_properties",
"core-foundation-sys",
"iana-time-zone-haiku",
"js-sys",
"log",
"wasm-bindgen",
"windows-core",
]
[[package]]
name = "icu_locale_core"
version = "2.2.0"
name = "iana-time-zone-haiku"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29"
checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
dependencies = [
"displaydoc",
"litemap",
"tinystr",
"writeable",
"zerovec",
]
[[package]]
name = "icu_normalizer"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4"
dependencies = [
"icu_collections",
"icu_normalizer_data",
"icu_properties",
"icu_provider",
"smallvec",
"zerovec",
]
[[package]]
name = "icu_normalizer_data"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38"
[[package]]
name = "icu_properties"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de"
dependencies = [
"icu_collections",
"icu_locale_core",
"icu_properties_data",
"icu_provider",
"zerotrie",
"zerovec",
]
[[package]]
name = "icu_properties_data"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14"
[[package]]
name = "icu_provider"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421"
dependencies = [
"displaydoc",
"icu_locale_core",
"writeable",
"yoke",
"zerofrom",
"zerotrie",
"zerovec",
"cc",
]
[[package]]
@ -1003,27 +951,6 @@ version = "2.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954"
[[package]]
name = "idna"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de"
dependencies = [
"idna_adapter",
"smallvec",
"utf8_iter",
]
[[package]]
name = "idna_adapter"
version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714"
dependencies = [
"icu_normalizer",
"icu_properties",
]
[[package]]
name = "indexmap"
version = "2.14.0"
@ -1065,22 +992,6 @@ dependencies = [
"cfg-if",
]
[[package]]
name = "ipnet"
version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "iri-string"
version = "0.7.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "25e659a4bb38e810ebc252e53b5814ff908a8c58c2a9ce2fae1bbec24cbf4e20"
dependencies = [
"memchr",
"serde",
]
[[package]]
name = "is_terminal_polyfill"
version = "1.70.2"
@ -1105,9 +1016,9 @@ dependencies = [
[[package]]
name = "js-sys"
version = "0.3.97"
version = "0.3.98"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1840c94c045fbcf8ba2812c95db44499f7c64910a912551aaaa541decebcacf"
checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08"
dependencies = [
"cfg-if",
"futures-util",
@ -1171,12 +1082,6 @@ version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "litemap"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0"
[[package]]
name = "lock_api"
version = "0.4.14"
@ -1326,6 +1231,15 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "num-traits"
version = "0.2.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
dependencies = [
"autocfg",
]
[[package]]
name = "once_cell"
version = "1.21.4"
@ -1405,15 +1319,6 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6"
[[package]]
name = "potential_utf"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564"
dependencies = [
"zerovec",
]
[[package]]
name = "prettyplease"
version = "0.2.37"
@ -1509,38 +1414,6 @@ version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a"
[[package]]
name = "reqwest"
version = "0.12.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64",
"bytes",
"futures-core",
"http",
"http-body",
"http-body-util",
"hyper",
"hyper-util",
"js-sys",
"log",
"percent-encoding",
"pin-project-lite",
"serde",
"serde_json",
"serde_urlencoded",
"sync_wrapper",
"tokio",
"tower",
"tower-http",
"tower-service",
"url",
"wasm-bindgen",
"wasm-bindgen-futures",
"web-sys",
]
[[package]]
name = "rustix"
version = "1.1.4"
@ -1723,12 +1596,6 @@ version = "0.9.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
[[package]]
name = "stable_deref_trait"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596"
[[package]]
name = "strsim"
version = "0.11.1"
@ -1757,20 +1624,6 @@ name = "sync_wrapper"
version = "1.0.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263"
dependencies = [
"futures-core",
]
[[package]]
name = "synstructure"
version = "0.13.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "sysinfo"
@ -1828,16 +1681,6 @@ dependencies = [
"cfg-if",
]
[[package]]
name = "tinystr"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d"
dependencies = [
"displaydoc",
"zerovec",
]
[[package]]
name = "tokio"
version = "1.52.1"
@ -1911,14 +1754,12 @@ dependencies = [
"http-body-util",
"http-range-header",
"httpdate",
"iri-string",
"mime",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"tokio",
"tokio-util",
"tower",
"tower-layer",
"tower-service",
"tracing",
@ -2011,12 +1852,6 @@ dependencies = [
"tracing-serde",
]
[[package]]
name = "try-lock"
version = "0.2.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
[[package]]
name = "typenum"
version = "1.20.0"
@ -2041,24 +1876,6 @@ version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
[[package]]
name = "url"
version = "2.5.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed"
dependencies = [
"form_urlencoded",
"idna",
"percent-encoding",
"serde",
]
[[package]]
name = "utf8_iter"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
[[package]]
name = "utf8parse"
version = "0.2.2"
@ -2087,15 +1904,6 @@ dependencies = [
"winapi-util",
]
[[package]]
name = "want"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e"
dependencies = [
"try-lock",
]
[[package]]
name = "wasi"
version = "0.11.1+wasi-snapshot-preview1"
@ -2122,9 +1930,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen"
version = "0.2.120"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df52b6d9b87e0c74c9edfa1eb2d9bf85e5d63515474513aa50fa181b3c4f5db1"
checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790"
dependencies = [
"cfg-if",
"once_cell",
@ -2133,21 +1941,11 @@ dependencies = [
"wasm-bindgen-shared",
]
[[package]]
name = "wasm-bindgen-futures"
version = "0.4.70"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "af934872acec734c2d80e6617bbb5ff4f12b052dd8e6332b0817bce889516084"
dependencies = [
"js-sys",
"wasm-bindgen",
]
[[package]]
name = "wasm-bindgen-macro"
version = "0.2.120"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78b1041f495fb322e64aca85f5756b2172e35cd459376e67f2a6c9dffcedb103"
checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578"
dependencies = [
"quote",
"wasm-bindgen-macro-support",
@ -2155,9 +1953,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-macro-support"
version = "0.2.120"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dcd0ff20416988a18ac686d4d4d0f6aae9ebf08a389ff5d29012b05af2a1b41"
checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2"
dependencies = [
"bumpalo",
"proc-macro2",
@ -2168,9 +1966,9 @@ dependencies = [
[[package]]
name = "wasm-bindgen-shared"
version = "0.2.120"
version = "0.2.121"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49757b3c82ebf16c57d69365a142940b384176c24df52a087fb748e2085359ea"
checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441"
dependencies = [
"unicode-ident",
]
@ -2209,16 +2007,6 @@ dependencies = [
"semver",
]
[[package]]
name = "web-sys"
version = "0.3.97"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2eadbac71025cd7b0834f20d1fe8472e8495821b4e9801eb0a60bd1f19827602"
dependencies = [
"js-sys",
"wasm-bindgen",
]
[[package]]
name = "winapi"
version = "0.3.9"
@ -2485,56 +2273,6 @@ dependencies = [
"wasmparser",
]
[[package]]
name = "writeable"
version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4"
[[package]]
name = "yoke"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca"
dependencies = [
"stable_deref_trait",
"yoke-derive",
"zerofrom",
]
[[package]]
name = "yoke-derive"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e"
dependencies = [
"proc-macro2",
"quote",
"syn",
"synstructure",
]
[[package]]
name = "zerofrom"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df"
dependencies = [
"zerofrom-derive",
]
[[package]]
name = "zerofrom-derive"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1"
dependencies = [
"proc-macro2",
"quote",
"syn",
"synstructure",
]
[[package]]
name = "zeroize"
version = "1.8.2"
@ -2555,39 +2293,6 @@ dependencies = [
"syn",
]
[[package]]
name = "zerotrie"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf"
dependencies = [
"displaydoc",
"yoke",
"zerofrom",
]
[[package]]
name = "zerovec"
version = "0.11.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239"
dependencies = [
"yoke",
"zerofrom",
"zerovec-derive",
]
[[package]]
name = "zerovec-derive"
version = "0.11.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "zmij"
version = "1.0.21"

View File

@ -1,8 +1,8 @@
[package]
name = "envd"
version = "0.2.1"
version = "0.3.0"
edition = "2024"
rust-version = "1.88"
rust-version = "1.95"
[dependencies]
# Async runtime
@ -53,12 +53,12 @@ notify = "7"
# Compression
flate2 = "1"
# HTTP client (MMDS polling)
reqwest = { version = "0.12", default-features = false, features = ["json"] }
# Directory walking
walkdir = "2"
# Time parsing (RFC3339 → epoch nanos for clock fix)
chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
# Misc
libc = "0.2"
bytes = "1"

View File

@ -1,6 +1,6 @@
# envd (Rust)
Wrenn guest agent daemon — runs as PID 1 inside Firecracker microVMs. Provides process management, filesystem operations, file transfer, port forwarding, and VM lifecycle control over Connect RPC and HTTP.
Wrenn guest agent daemon — runs as PID 1 inside Cloud Hypervisor microVMs. Provides process management, filesystem operations, file transfer, port forwarding, and VM lifecycle control over Connect RPC and HTTP.
Rust rewrite of `envd/` (Go). Drop-in replacement — same wire protocol, same endpoints, same CLI flags.
@ -50,7 +50,7 @@ cargo build
Run locally (outside a VM):
```bash
./target/debug/envd --isnotfc --port 49983
./target/debug/envd --port 49983
```
### Via Makefile (from repo root)
@ -64,7 +64,6 @@ make build-envd-go # Go version (for comparison)
```
--port <PORT> Listen port [default: 49983]
--isnotfc Not running inside Firecracker (disables MMDS, cgroups)
--version Print version and exit
--commit Print git commit and exit
--cmd <CMD> Spawn a process at startup (e.g. --cmd "/bin/bash")
@ -81,7 +80,7 @@ make build-envd-go # Go version (for comparison)
| GET | `/metrics` | System metrics (CPU, memory, disk) |
| GET | `/envs` | Current environment variables |
| POST | `/init` | Host agent init (token, env, mounts) |
| POST | `/snapshot/prepare` | Quiesce before Firecracker snapshot |
| POST | `/snapshot/prepare` | Quiesce before Cloud Hypervisor snapshot |
| GET | `/files` | Download file (gzip, range support) |
| POST | `/files` | Upload file(s) via multipart |
@ -108,7 +107,7 @@ src/
├── util.rs # AtomicMax
├── auth/ # Token, signing, middleware
├── crypto/ # SHA-256, SHA-512, HMAC
├── host/ # MMDS polling, system metrics
├── host/ # System metrics
├── http/ # Axum handlers (health, init, snapshot, files, encoding)
├── permissions/ # Path resolution, user lookup, chown
├── rpc/ # Connect RPC services
@ -129,13 +128,15 @@ src/
After building the static binary, copy it into the rootfs:
```bash
bash scripts/update-debug-rootfs.sh [rootfs_path]
bash scripts/update-minimal-rootfs.sh [rootfs_path]
```
Or manually:
With no argument it updates all four system base images; pass a path to target one.
Or manually (example path: the minimal-ubuntu image, platform team + template id 0):
```bash
sudo mount -o loop /var/lib/wrenn/images/minimal.ext4 /mnt
sudo cp target/x86_64-unknown-linux-musl/release/envd /mnt/usr/bin/envd
sudo mount -o loop /var/lib/wrenn/images/teams/0000000000000000000000000/0000000000000000000000000/rootfs.ext4 /mnt
sudo cp target/x86_64-unknown-linux-musl/release/envd /mnt/usr/local/bin/envd
sudo umount /mnt
```

View File

@ -9,8 +9,3 @@ pub const WRENN_RUN_DIR: &str = "/run/wrenn";
pub const KILOBYTE: u64 = 1024;
pub const MEGABYTE: u64 = 1024 * KILOBYTE;
pub const MMDS_ADDRESS: &str = "169.254.169.254";
pub const MMDS_POLL_INTERVAL: Duration = Duration::from_millis(50);
pub const MMDS_TOKEN_EXPIRATION_SECS: u64 = 60;
pub const MMDS_ACCESS_TOKEN_CLIENT_TIMEOUT: Duration = Duration::from_secs(10);

View File

@ -1,24 +1,14 @@
use std::collections::HashSet;
use std::sync::Mutex;
/// Tracks active TCP connections for snapshot/restore lifecycle.
///
/// Before snapshot: close idle connections, record active ones.
/// After restore: close all pre-snapshot connections (zombie TCP sockets).
///
/// In Rust/axum, we don't have Go's ConnState callback. Instead we track
/// connections via a tower middleware that registers connection IDs.
/// For the initial implementation, we track by a simple connection counter
/// and rely on axum's graceful shutdown mechanics.
/// Tracks active TCP connections.
pub struct ConnTracker {
inner: Mutex<ConnTrackerInner>,
}
struct ConnTrackerInner {
active: HashSet<u64>,
pre_snapshot: Option<HashSet<u64>>,
next_id: u64,
keepalives_enabled: bool,
}
impl ConnTracker {
@ -26,9 +16,7 @@ impl ConnTracker {
Self {
inner: Mutex::new(ConnTrackerInner {
active: HashSet::new(),
pre_snapshot: None,
next_id: 0,
keepalives_enabled: true,
}),
}
}
@ -44,37 +32,6 @@ impl ConnTracker {
pub fn remove_connection(&self, id: u64) {
let mut inner = self.inner.lock().unwrap();
inner.active.remove(&id);
if let Some(ref mut pre) = inner.pre_snapshot {
pre.remove(&id);
}
}
pub fn prepare_for_snapshot(&self) {
let mut inner = self.inner.lock().unwrap();
inner.keepalives_enabled = false;
inner.pre_snapshot = Some(inner.active.clone());
tracing::info!(
active_connections = inner.active.len(),
"snapshot: recorded pre-snapshot connections, keep-alives disabled"
);
}
pub fn restore_after_snapshot(&self) {
let mut inner = self.inner.lock().unwrap();
if let Some(pre) = inner.pre_snapshot.take() {
let zombie_count = pre.len();
for id in &pre {
inner.active.remove(id);
}
if zombie_count > 0 {
tracing::info!(zombie_count, "restore: closed zombie connections");
}
}
inner.keepalives_enabled = true;
}
pub fn keepalives_enabled(&self) -> bool {
self.inner.lock().unwrap().keepalives_enabled
}
#[cfg(test)]
@ -110,91 +67,4 @@ mod tests {
ct.remove_connection(999);
assert_eq!(ct.active_count(), 0);
}
#[test]
fn prepare_disables_keepalives() {
let ct = ConnTracker::new();
assert!(ct.keepalives_enabled());
ct.register_connection();
ct.prepare_for_snapshot();
assert!(!ct.keepalives_enabled());
}
#[test]
fn restore_removes_zombies_and_reenables_keepalives() {
let ct = ConnTracker::new();
let id0 = ct.register_connection();
let id1 = ct.register_connection();
ct.prepare_for_snapshot();
ct.restore_after_snapshot();
assert!(ct.keepalives_enabled());
// Both pre-snapshot connections removed as zombies
assert_eq!(ct.active_count(), 0);
// IDs don't matter anymore, but remove shouldn't panic
ct.remove_connection(id0);
ct.remove_connection(id1);
}
#[test]
fn restore_without_prepare_is_noop() {
let ct = ConnTracker::new();
let _id = ct.register_connection();
ct.restore_after_snapshot();
assert!(ct.keepalives_enabled());
assert_eq!(ct.active_count(), 1);
}
#[test]
fn connection_closed_before_restore_not_zombie() {
let ct = ConnTracker::new();
let id0 = ct.register_connection();
let _id1 = ct.register_connection();
ct.prepare_for_snapshot();
// Close id0 during snapshot window
ct.remove_connection(id0);
assert_eq!(ct.active_count(), 1);
ct.restore_after_snapshot();
// id1 was zombie (still active at restore), id0 already gone
assert_eq!(ct.active_count(), 0);
}
#[test]
fn post_snapshot_connection_survives_restore() {
let ct = ConnTracker::new();
ct.register_connection();
ct.prepare_for_snapshot();
// New connection after snapshot
let _post = ct.register_connection();
ct.restore_after_snapshot();
// Pre-snapshot connection removed, post-snapshot survives
assert_eq!(ct.active_count(), 1);
}
#[test]
fn full_lifecycle() {
let ct = ConnTracker::new();
let _a = ct.register_connection();
let b = ct.register_connection();
let _c = ct.register_connection();
assert_eq!(ct.active_count(), 3);
assert!(ct.keepalives_enabled());
ct.prepare_for_snapshot();
assert!(!ct.keepalives_enabled());
let d = ct.register_connection();
ct.remove_connection(b);
ct.restore_after_snapshot();
assert!(ct.keepalives_enabled());
// a and c were zombies, b removed before restore, d is post-snapshot
assert_eq!(ct.active_count(), 1);
ct.remove_connection(d);
assert_eq!(ct.active_count(), 0);
// Can reuse tracker after restore
let e = ct.register_connection();
assert_eq!(ct.active_count(), 1);
assert!(e > d);
}
}

View File

@ -1,73 +0,0 @@
use std::ffi::CString;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::Serialize;
#[derive(Serialize)]
pub struct Metrics {
pub ts: i64,
pub cpu_count: u32,
pub cpu_used_pct: f32,
pub mem_total_mib: u64,
pub mem_used_mib: u64,
pub mem_total: u64,
pub mem_used: u64,
pub disk_used: u64,
pub disk_total: u64,
}
pub fn get_metrics() -> Result<Metrics, String> {
use sysinfo::System;
let mut sys = System::new();
sys.refresh_memory();
sys.refresh_cpu_all();
std::thread::sleep(std::time::Duration::from_millis(100));
sys.refresh_cpu_all();
let cpu_count = sys.cpus().len() as u32;
let cpu_used_pct = sys.global_cpu_usage();
let cpu_used_pct_rounded = if cpu_used_pct > 0.0 {
(cpu_used_pct * 100.0).round() / 100.0
} else {
0.0
};
let mem_total = sys.total_memory();
let mem_used = sys.used_memory();
let (disk_total, disk_used) = disk_stats("/")?;
let ts = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs() as i64;
Ok(Metrics {
ts,
cpu_count,
cpu_used_pct: cpu_used_pct_rounded,
mem_total_mib: mem_total / 1024 / 1024,
mem_used_mib: mem_used / 1024 / 1024,
mem_total,
mem_used,
disk_used,
disk_total,
})
}
fn disk_stats(path: &str) -> Result<(u64, u64), String> {
let c_path = CString::new(path).unwrap();
let mut stat: libc::statfs = unsafe { std::mem::zeroed() };
let ret = unsafe { libc::statfs(c_path.as_ptr(), &mut stat) };
if ret != 0 {
return Err(format!("statfs failed: {}", std::io::Error::last_os_error()));
}
let block = stat.f_bsize as u64;
let total = stat.f_blocks * block;
let available = stat.f_bavail * block;
Ok((total, total - available))
}

View File

@ -1,120 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use dashmap::DashMap;
use serde::Deserialize;
use tokio_util::sync::CancellationToken;
use crate::config::{MMDS_ADDRESS, MMDS_POLL_INTERVAL, MMDS_TOKEN_EXPIRATION_SECS, WRENN_RUN_DIR};
#[derive(Debug, Clone, Deserialize)]
pub struct MMDSOpts {
#[serde(rename = "instanceID")]
pub sandbox_id: String,
#[serde(rename = "envID")]
pub template_id: String,
#[serde(rename = "address", default)]
pub logs_collector_address: String,
#[serde(rename = "accessTokenHash", default)]
pub access_token_hash: String,
}
async fn get_mmds_token(client: &reqwest::Client) -> Result<String, String> {
let resp = client
.put(format!("http://{MMDS_ADDRESS}/latest/api/token"))
.header(
"X-metadata-token-ttl-seconds",
MMDS_TOKEN_EXPIRATION_SECS.to_string(),
)
.send()
.await
.map_err(|e| format!("mmds token request failed: {e}"))?;
let token = resp.text().await.map_err(|e| format!("mmds token read: {e}"))?;
if token.is_empty() {
return Err("mmds token is an empty string".into());
}
Ok(token)
}
async fn get_mmds_opts(client: &reqwest::Client, token: &str) -> Result<MMDSOpts, String> {
let resp = client
.get(format!("http://{MMDS_ADDRESS}"))
.header("X-metadata-token", token)
.header("Accept", "application/json")
.send()
.await
.map_err(|e| format!("mmds opts request failed: {e}"))?;
resp.json::<MMDSOpts>()
.await
.map_err(|e| format!("mmds opts parse: {e}"))
}
pub async fn get_access_token_hash() -> Result<String, String> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.no_proxy()
.build()
.map_err(|e| format!("http client: {e}"))?;
let token = get_mmds_token(&client).await?;
let opts = get_mmds_opts(&client, &token).await?;
Ok(opts.access_token_hash)
}
/// Polls MMDS every 50ms until metadata is available.
/// Stores sandbox_id and template_id in env_vars and writes to /run/wrenn/ files.
pub async fn poll_for_opts(
env_vars: Arc<DashMap<String, String>>,
cancel: CancellationToken,
) -> Option<MMDSOpts> {
let client = reqwest::Client::builder()
.no_proxy()
.build()
.ok()?;
let mut interval = tokio::time::interval(MMDS_POLL_INTERVAL);
loop {
tokio::select! {
_ = cancel.cancelled() => {
tracing::warn!("context cancelled while waiting for mmds opts");
return None;
}
_ = interval.tick() => {
let token = match get_mmds_token(&client).await {
Ok(t) => t,
Err(e) => {
tracing::debug!(error = %e, "mmds token poll");
continue;
}
};
let opts = match get_mmds_opts(&client, &token).await {
Ok(o) => o,
Err(e) => {
tracing::debug!(error = %e, "mmds opts poll");
continue;
}
};
env_vars.insert("WRENN_SANDBOX_ID".into(), opts.sandbox_id.clone());
env_vars.insert("WRENN_TEMPLATE_ID".into(), opts.template_id.clone());
let run_dir = std::path::Path::new(WRENN_RUN_DIR);
if let Err(e) = std::fs::create_dir_all(run_dir) {
tracing::error!(error = %e, "mmds: failed to create run dir");
}
if let Err(e) = std::fs::write(run_dir.join(".WRENN_SANDBOX_ID"), &opts.sandbox_id) {
tracing::error!(error = %e, "mmds: failed to write .WRENN_SANDBOX_ID");
}
if let Err(e) = std::fs::write(run_dir.join(".WRENN_TEMPLATE_ID"), &opts.template_id) {
tracing::error!(error = %e, "mmds: failed to write .WRENN_TEMPLATE_ID");
}
return Some(opts);
}
}
}
}

View File

@ -1,2 +0,0 @@
pub mod metrics;
pub mod mmds;

View File

@ -1,5 +1,4 @@
use std::sync::Arc;
use std::sync::atomic::Ordering;
use axum::Json;
use axum::extract::State;
@ -10,14 +9,6 @@ use serde_json::json;
use crate::state::AppState;
pub async fn get_health(State(state): State<Arc<AppState>>) -> impl IntoResponse {
if state
.needs_restore
.compare_exchange(true, false, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
post_restore_recovery(&state);
}
tracing::trace!("health check");
(
@ -25,17 +16,3 @@ pub async fn get_health(State(state): State<Arc<AppState>>) -> impl IntoResponse
Json(json!({ "version": state.version })),
)
}
fn post_restore_recovery(state: &AppState) {
tracing::info!("restore: post-restore recovery (no GC needed in Rust)");
state.snapshot_in_progress.store(false, std::sync::atomic::Ordering::Release);
state.conn_tracker.restore_after_snapshot();
tracing::info!("restore: zombie connections closed");
if let Some(ref ps) = state.port_subsystem {
ps.restart();
tracing::info!("restore: port subsystem restarted");
}
}

View File

@ -1,6 +1,5 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use axum::Json;
use axum::extract::State;
@ -8,20 +7,29 @@ use axum::http::{StatusCode, header};
use axum::response::IntoResponse;
use serde::Deserialize;
use crate::crypto;
use crate::host::mmds;
use crate::state::AppState;
#[derive(Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct InitRequest {
#[serde(rename = "access_token")]
pub access_token: Option<String>,
#[serde(rename = "defaultUser")]
pub default_user: Option<String>,
#[serde(rename = "defaultWorkdir")]
pub default_workdir: Option<String>,
#[serde(rename = "envVars")]
pub env_vars: Option<HashMap<String, String>>,
#[serde(rename = "hyperloop_ip")]
pub hyperloop_ip: Option<String>,
pub timestamp: Option<String>,
#[serde(rename = "volume_mounts")]
pub volume_mounts: Option<Vec<VolumeMount>>,
pub sandbox_id: Option<String>,
pub template_id: Option<String>,
/// New lifecycle identifier for this resume. When it changes between
/// /init calls, envd treats the call as a post-resume hook: port
/// forwarder is restarted and NFS mounts are refreshed.
pub lifecycle_id: Option<String>,
}
#[derive(Deserialize)]
@ -30,7 +38,7 @@ pub struct VolumeMount {
pub path: String,
}
/// POST /init — called by host agent after boot and after every resume.
/// POST /init — called by host agent after boot.
pub async fn post_init(
State(state): State<Arc<AppState>>,
body: Option<Json<InitRequest>>,
@ -45,12 +53,71 @@ pub async fn post_init(
}
}
// Idempotent timestamp check
// Post-resume lifecycle hook: restart port forwarder so socat children
// are reaped + respawned against the new wall clock and any rotated
// listeners. Must run BEFORE the stale-timestamp early-return so a
// resume with an out-of-order timestamp still refreshes the subsystem.
let lifecycle_changed = if let Some(ref new_id) = init_req.lifecycle_id {
state.bump_lifecycle(new_id)
} else {
false
};
if lifecycle_changed {
// Each new lifecycle (i.e. a snapshot restore) requires a fresh memory
// preload pass — pages materialised before the previous pause are now
// back in the source memory-ranges file as the host re-restored them
// lazily. Reset the flags so the next POST /memory/preload kicks off
// a new loader instead of returning the stale "already-done".
use std::sync::atomic::Ordering;
state.mem_preload_cancel.store(false, Ordering::SeqCst);
state.mem_preload_done.store(false, Ordering::SeqCst);
state.mem_preload_started.store(false, Ordering::SeqCst);
state.mem_preload_regions.store(0, Ordering::SeqCst);
state.mem_preload_pages.store(0, Ordering::SeqCst);
state.mem_preload_bytes.store(0, Ordering::SeqCst);
state.mem_preload_elapsed_us.store(0, Ordering::SeqCst);
state.mem_preload_source.store(0, Ordering::SeqCst);
*state.mem_preload_error.lock().unwrap() = None;
if let Some(ref port_sub) = state.port_subsystem {
tracing::info!("lifecycle changed, restarting port subsystem");
port_sub.restart();
}
// Force chrony to step the clock immediately. chronyd is launched by
// wrenn-init.sh and disciplines against PHC (/dev/ptp0), so the host
// wall time is already available — `makestep` just bypasses chrony's
// normal slewing and snaps the clock in one go. Best effort.
tokio::spawn(async {
match tokio::process::Command::new("chronyc")
.args(["makestep"])
.output()
.await
{
Ok(out) if out.status.success() => {
tracing::info!("chronyc makestep ok");
}
Ok(out) => {
let stderr = String::from_utf8_lossy(&out.stderr);
tracing::warn!(stderr = %stderr, "chronyc makestep failed");
}
Err(e) => {
tracing::warn!(error = %e, "chronyc makestep spawn failed");
}
}
});
}
// Idempotent timestamp check. Run after lifecycle handling so a
// stale-timestamp /init still gets to refresh ports + step clock.
// No userspace clock_settime here — chrony owns time discipline.
if let Some(ref ts_str) = init_req.timestamp {
if let Ok(ts) = chrono_parse_to_nanos(ts_str) {
if let Ok(ts) = parse_timestamp_to_nanos(ts_str) {
if !state.last_set_time.set_to_greater(ts) {
// Stale request, skip data updates
return trigger_restore_and_respond(&state).await;
return (
StatusCode::NO_CONTENT,
[(header::CACHE_CONTROL, "no-store")],
)
.into_response();
}
}
}
@ -90,56 +157,40 @@ pub async fn post_init(
}
}
// Hyperloop /etc/hosts setup
// Hyperloop /etc/hosts setup. Awaited so callers that immediately
// resolve events.wrenn.local see the entry. Cheap (two file ops).
if let Some(ref ip) = init_req.hyperloop_ip {
let ip = ip.clone();
let env_vars = Arc::clone(&state.defaults.env_vars);
tokio::spawn(async move {
setup_hyperloop(&ip, &env_vars).await;
});
setup_hyperloop(ip, &state.defaults.env_vars).await;
}
// NFS mounts
// NFS mounts. Awaited in parallel so callers that immediately access the
// mount path don't race the mount(2). Previously these were detached via
// tokio::spawn, which let /init return success before mounts existed.
if let Some(ref mounts) = init_req.volume_mounts {
for mount in mounts {
let target = mount.nfs_target.clone();
let path = mount.path.clone();
tokio::spawn(async move {
let futs = mounts.iter().map(|m| {
let target = m.nfs_target.clone();
let path = m.path.clone();
async move {
setup_nfs(&target, &path).await;
});
}
}
// Re-poll MMDS in background
if state.is_fc {
let env_vars = Arc::clone(&state.defaults.env_vars);
let cancel = tokio_util::sync::CancellationToken::new();
let cancel_clone = cancel.clone();
tokio::spawn(async move {
tokio::time::timeout(std::time::Duration::from_secs(60), async {
mmds::poll_for_opts(env_vars, cancel_clone).await;
})
.await
.ok();
}
});
futures::future::join_all(futs).await;
}
trigger_restore_and_respond(&state).await
}
async fn trigger_restore_and_respond(state: &AppState) -> axum::response::Response {
// Safety net: if health check's postRestoreRecovery hasn't run yet
if state
.needs_restore
.compare_exchange(true, false, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
post_restore_recovery(state);
// Set sandbox/template metadata from request body.
if let Some(ref id) = init_req.sandbox_id {
tracing::debug!(sandbox_id = %id, "setting sandbox ID from init request");
// SAFETY: envd is single-threaded at init time; no concurrent env reads.
unsafe { std::env::set_var("WRENN_SANDBOX_ID", id) };
write_run_file(".WRENN_SANDBOX_ID", id);
state.defaults.env_vars.insert("WRENN_SANDBOX_ID".into(), id.clone());
}
state.conn_tracker.restore_after_snapshot();
if let Some(ref ps) = state.port_subsystem {
ps.restart();
if let Some(ref id) = init_req.template_id {
tracing::debug!(template_id = %id, "setting template ID from init request");
// SAFETY: envd is single-threaded at init time; no concurrent env reads.
unsafe { std::env::set_var("WRENN_TEMPLATE_ID", id) };
write_run_file(".WRENN_TEMPLATE_ID", id);
state.defaults.env_vars.insert("WRENN_TEMPLATE_ID".into(), id.clone());
}
(
@ -149,46 +200,13 @@ async fn trigger_restore_and_respond(state: &AppState) -> axum::response::Respon
.into_response()
}
fn post_restore_recovery(state: &AppState) {
tracing::info!("restore: post-restore recovery (no GC needed in Rust)");
state.snapshot_in_progress.store(false, std::sync::atomic::Ordering::Release);
state.conn_tracker.restore_after_snapshot();
if let Some(ref ps) = state.port_subsystem {
ps.restart();
tracing::info!("restore: port subsystem restarted");
}
}
async fn validate_init_access_token(state: &AppState, request_token: &str) -> Result<(), String> {
// Fast path: matches existing token
if state.access_token.is_set() && !request_token.is_empty() && state.access_token.equals(request_token) {
return Ok(());
}
// Check MMDS hash
if state.is_fc {
if let Ok(mmds_hash) = mmds::get_access_token_hash().await {
if !mmds_hash.is_empty() {
if request_token.is_empty() {
let empty_hash = crypto::sha512::hash_access_token("");
if mmds_hash == empty_hash {
return Ok(());
}
} else {
let token_hash = crypto::sha512::hash_access_token(request_token);
if mmds_hash == token_hash {
return Ok(());
}
}
return Err("access token validation failed".into());
}
}
}
// First-time setup: no existing token and no MMDS
// First-time setup: no existing token
if !state.access_token.is_set() {
return Ok(());
}
@ -268,14 +286,27 @@ async fn setup_nfs(nfs_target: &str, path: &str) {
}
}
fn chrono_parse_to_nanos(ts: &str) -> Result<i64, ()> {
// Parse RFC3339 timestamp to nanoseconds since epoch
// Simple approach: parse as seconds + fractional
let secs = ts.parse::<f64>().ok();
if let Some(s) = secs {
return Ok((s * 1_000_000_000.0) as i64);
fn write_run_file(name: &str, value: &str) {
let dir = std::path::Path::new("/run/wrenn");
if let Err(e) = std::fs::create_dir_all(dir) {
tracing::warn!(error = %e, "failed to create /run/wrenn");
return;
}
if let Err(e) = std::fs::write(dir.join(name), value) {
tracing::warn!(error = %e, name, "failed to write run file");
}
}
/// Parses a host-provided timestamp into nanoseconds since the Unix epoch.
/// Accepts either RFC3339 (`2026-05-17T16:13:03.123456Z`) or a float-seconds
/// string (legacy callers).
fn parse_timestamp_to_nanos(ts: &str) -> Result<i64, ()> {
if let Ok(parsed) = chrono::DateTime::parse_from_rfc3339(ts) {
return Ok(parsed.timestamp_nanos_opt().ok_or(())?);
}
if let Ok(secs) = ts.parse::<f64>() {
return Ok((secs * 1_000_000_000.0) as i64);
}
// Try RFC3339 format
// For now, fall back to allowing the update
Err(())
}

350
envd-rs/src/http/memory.rs Normal file
View File

@ -0,0 +1,350 @@
// POST /memory/preload — guest-side helper that materialises every physical
// RAM page so a subsequent ch.snapshot writes a self-contained memory-ranges
// file. Required after a restore with memory_restore_mode=ondemand: pages
// that were never demand-faulted live only in the source memory-ranges file
// and would become holes (read back as zero) in the new snapshot.
//
// Trigger is one-byte-per-page reads through /dev/mem (fallback /proc/kcore
// PT_LOAD segments). The guest kernel walks its direct map → accesses the
// physical page → host kernel handles the EPT entry → CH's userfaultfd
// handler fills the page from the source memory-ranges file.
//
// Wire protocol:
// POST /memory/preload — starts the loader (idempotent) and returns
// the current status JSON immediately
// GET /memory/preload — returns the current status JSON
// POST /memory/preload/cancel — signals the loader to stop early
//
// Returning immediately avoids any HTTP-level header timeout in the caller
// while materialisation (hundreds of MiB at one byte per page) runs in a
// background blocking thread.
use std::fs;
use std::io::{Read, Seek, SeekFrom};
use std::os::unix::fs::FileExt;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Instant;
use axum::Json;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::Serialize;
use crate::state::AppState;
const PAGE_SIZE: u64 = 4096;
#[derive(Serialize, Clone)]
pub struct PreloadStatus {
/// One of: "idle", "running", "done", "failed", "cancelled".
pub state: &'static str,
pub regions: u64,
pub pages: u64,
pub bytes: u64,
pub elapsed_sec: f64,
pub source: &'static str,
pub error: Option<String>,
}
pub async fn post_memory_preload(State(state): State<Arc<AppState>>) -> impl IntoResponse {
// First caller wins the CAS and spawns the loader; subsequent callers
// just report the existing status.
let we_start = state
.mem_preload_started
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok();
if we_start {
let state_clone = Arc::clone(&state);
// Detached blocking thread — no axum task lifetime ties it to the
// request, so the connection can close immediately without aborting
// materialisation. Lifecycle bump on the next /init clears the flags
// for a fresh run after a restore.
std::thread::spawn(move || {
let started = Instant::now();
match preload_blocking(&state_clone) {
Ok((source, regions, pages, bytes)) => {
let elapsed = started.elapsed().as_secs_f64();
state_clone
.mem_preload_regions
.store(regions, Ordering::SeqCst);
state_clone.mem_preload_pages.store(pages, Ordering::SeqCst);
state_clone.mem_preload_bytes.store(bytes, Ordering::SeqCst);
state_clone
.mem_preload_elapsed_us
.store((elapsed * 1_000_000.0) as u64, Ordering::SeqCst);
set_source(&state_clone, source);
*state_clone.mem_preload_error.lock().unwrap() = None;
state_clone.mem_preload_done.store(true, Ordering::SeqCst);
tracing::info!(
regions,
pages,
bytes,
elapsed_sec = elapsed,
source,
"memory preload complete"
);
}
Err(e) => {
let elapsed = started.elapsed().as_secs_f64();
state_clone
.mem_preload_elapsed_us
.store((elapsed * 1_000_000.0) as u64, Ordering::SeqCst);
*state_clone.mem_preload_error.lock().unwrap() = Some(e.clone());
state_clone.mem_preload_done.store(true, Ordering::SeqCst);
tracing::warn!(error = %e, "memory preload failed");
}
}
});
}
let status = read_status(&state);
(StatusCode::OK, Json(status))
}
pub async fn get_memory_preload(State(state): State<Arc<AppState>>) -> impl IntoResponse {
(StatusCode::OK, Json(read_status(&state)))
}
pub async fn post_memory_preload_cancel(State(state): State<Arc<AppState>>) -> impl IntoResponse {
state.mem_preload_cancel.store(true, Ordering::SeqCst);
StatusCode::NO_CONTENT
}
fn read_status(state: &AppState) -> PreloadStatus {
let started = state.mem_preload_started.load(Ordering::SeqCst);
let done = state.mem_preload_done.load(Ordering::SeqCst);
let cancelled = state.mem_preload_cancel.load(Ordering::SeqCst);
let error = state.mem_preload_error.lock().unwrap().clone();
let lane = if !started {
"idle"
} else if !done {
"running"
} else if let Some(_) = &error {
"failed"
} else if cancelled {
"cancelled"
} else {
"done"
};
PreloadStatus {
state: lane,
regions: state.mem_preload_regions.load(Ordering::SeqCst),
pages: state.mem_preload_pages.load(Ordering::SeqCst),
bytes: state.mem_preload_bytes.load(Ordering::SeqCst),
elapsed_sec: state.mem_preload_elapsed_us.load(Ordering::SeqCst) as f64 / 1_000_000.0,
source: get_source(state),
error,
}
}
fn set_source(state: &AppState, src: &'static str) {
let code: u8 = match src {
"/dev/mem" => 1,
"/proc/kcore" => 2,
_ => 0,
};
state.mem_preload_source.store(code, Ordering::SeqCst);
}
fn get_source(state: &AppState) -> &'static str {
match state.mem_preload_source.load(Ordering::SeqCst) {
1 => "/dev/mem",
2 => "/proc/kcore",
_ => "",
}
}
fn preload_blocking(state: &AppState) -> Result<(&'static str, u64, u64, u64), String> {
let ranges = parse_system_ram_ranges().map_err(|e| format!("iomem: {e}"))?;
if ranges.is_empty() {
return Err("no System RAM ranges found in /proc/iomem".into());
}
let mut pages: u64 = 0;
let mut bytes: u64 = 0;
match preload_via_devmem(&ranges, state, &mut pages, &mut bytes) {
Ok(()) => Ok(("/dev/mem", ranges.len() as u64, pages, bytes)),
Err(devmem_err) => {
tracing::warn!(
error = %devmem_err,
"/dev/mem preload failed, falling back to /proc/kcore"
);
pages = 0;
bytes = 0;
preload_via_kcore(state, &mut pages, &mut bytes)
.map_err(|e| format!("/dev/mem: {devmem_err}; /proc/kcore: {e}"))?;
Ok(("/proc/kcore", ranges.len() as u64, pages, bytes))
}
}
}
fn parse_system_ram_ranges() -> std::io::Result<Vec<(u64, u64)>> {
let data = fs::read_to_string("/proc/iomem")?;
let mut out = Vec::new();
for line in data.lines() {
if line.starts_with(|c: char| c.is_whitespace()) {
continue;
}
let Some((range, label)) = line.split_once(" : ") else {
continue;
};
if label.trim() != "System RAM" {
continue;
}
let Some((start, end)) = range.split_once('-') else {
continue;
};
let start = u64::from_str_radix(start.trim(), 16).ok();
let end = u64::from_str_radix(end.trim(), 16).ok();
if let (Some(s), Some(e)) = (start, end) {
out.push((s, e.saturating_add(1)));
}
}
Ok(out)
}
fn preload_via_devmem(
ranges: &[(u64, u64)],
state: &AppState,
pages: &mut u64,
bytes: &mut u64,
) -> std::io::Result<()> {
let f = fs::File::open("/dev/mem")?;
let mut buf = [0u8; 1];
for (start, end) in ranges {
let mut off = *start;
while off < *end {
if state.mem_preload_cancel.load(Ordering::SeqCst) {
return Ok(());
}
f.read_at(&mut buf, off)?;
*pages += 1;
*bytes += PAGE_SIZE;
// Publish progress so GET /memory/preload reports useful numbers
// while the loader is still running.
if *pages % 1024 == 0 {
state.mem_preload_pages.store(*pages, Ordering::SeqCst);
state.mem_preload_bytes.store(*bytes, Ordering::SeqCst);
}
off = off.saturating_add(PAGE_SIZE);
}
}
Ok(())
}
// Read /proc/kcore's direct-map segment to materialise physical RAM. The
// direct map's PT_LOAD covers the kernel's *maximum* possible direct-map
// region (64TB on x86_64), not just the present physical RAM — iterating the
// whole segment would loop for billions of pages. Bound the walk to the sum
// of System RAM ranges from /proc/iomem; sequential reads through the
// segment touch consecutive physical pages 1:1, which is what we need.
fn preload_via_kcore(state: &AppState, pages: &mut u64, bytes: &mut u64) -> std::io::Result<()> {
let ram_ranges = parse_system_ram_ranges()?;
if ram_ranges.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"no System RAM ranges to bound kcore walk",
));
}
let total_ram_bytes: u64 = ram_ranges.iter().map(|(s, e)| e - s).sum();
let mut f = fs::File::open("/proc/kcore")?;
let segments = parse_kcore_pt_load(&mut f)?;
if segments.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"no PT_LOAD segments in /proc/kcore",
));
}
// Pick the direct map: highest-vaddr kernel-space segment large enough
// to plausibly cover RAM. KASLR randomises the base, so don't hardcode it.
// Kernel virtual addresses start at 0xffff800000000000 on x86_64; vmalloc
// / modules sit above the direct map and are usually smaller.
const KERNEL_SPACE_MIN: u64 = 0xffff_8000_0000_0000;
let direct_map = segments
.iter()
.filter(|s| s.vaddr >= KERNEL_SPACE_MIN && s.file_size >= total_ram_bytes)
.min_by_key(|s| s.vaddr)
.ok_or_else(|| {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"no PT_LOAD segment large enough for {} bytes of RAM in /proc/kcore",
total_ram_bytes
),
)
})?;
let mut buf = [0u8; 1];
let read_bytes = total_ram_bytes.min(direct_map.file_size);
let start = direct_map.file_offset;
let end = start.saturating_add(read_bytes);
let mut off = start;
while off < end {
if state.mem_preload_cancel.load(Ordering::SeqCst) {
return Ok(());
}
// Reads into MMIO holes within the direct map can fail; ignore so the
// loop keeps making progress over the present RAM ranges either side.
if f.read_at(&mut buf, off).is_ok() {
*pages += 1;
*bytes += PAGE_SIZE;
if *pages % 256 == 0 {
state.mem_preload_pages.store(*pages, Ordering::SeqCst);
state.mem_preload_bytes.store(*bytes, Ordering::SeqCst);
}
}
off = off.saturating_add(PAGE_SIZE);
}
Ok(())
}
struct KcoreSegment {
file_offset: u64,
file_size: u64,
vaddr: u64,
}
fn parse_kcore_pt_load(f: &mut fs::File) -> std::io::Result<Vec<KcoreSegment>> {
let mut hdr = [0u8; 64];
f.seek(SeekFrom::Start(0))?;
f.read_exact(&mut hdr)?;
if &hdr[0..4] != b"\x7fELF" || hdr[4] != 2 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"not an ELF64 file",
));
}
let e_phoff = u64::from_le_bytes(hdr[32..40].try_into().unwrap());
let e_phentsize = u16::from_le_bytes(hdr[54..56].try_into().unwrap()) as u64;
let e_phnum = u16::from_le_bytes(hdr[56..58].try_into().unwrap()) as u64;
let mut out = Vec::new();
let mut entry = vec![0u8; e_phentsize as usize];
for i in 0..e_phnum {
f.seek(SeekFrom::Start(e_phoff + i * e_phentsize))?;
f.read_exact(&mut entry)?;
let p_type = u32::from_le_bytes(entry[0..4].try_into().unwrap());
if p_type != 1 {
continue;
}
let p_offset = u64::from_le_bytes(entry[8..16].try_into().unwrap());
let p_vaddr = u64::from_le_bytes(entry[16..24].try_into().unwrap());
let p_filesz = u64::from_le_bytes(entry[32..40].try_into().unwrap());
out.push(KcoreSegment {
file_offset: p_offset,
file_size: p_filesz,
vaddr: p_vaddr,
});
}
Ok(out)
}

View File

@ -4,6 +4,7 @@ pub mod error;
pub mod files;
pub mod health;
pub mod init;
pub mod memory;
pub mod metrics;
pub mod snapshot;
@ -50,6 +51,14 @@ pub fn router(state: Arc<AppState>) -> Router {
.route("/envs", get(envs::get_envs))
.route("/init", post(init::post_init))
.route("/snapshot/prepare", post(snapshot::post_snapshot_prepare))
.route(
"/memory/preload",
get(memory::get_memory_preload).post(memory::post_memory_preload),
)
.route(
"/memory/preload/cancel",
post(memory::post_memory_preload_cancel),
)
.route("/files", get(files::get_files).post(files::post_files))
.layer(cors)
.with_state(state)

View File

@ -1,47 +1,58 @@
use std::sync::Arc;
use std::sync::atomic::Ordering;
use axum::extract::State;
use axum::http::{StatusCode, header};
use axum::response::IntoResponse;
use nix::unistd::sync;
use crate::state::AppState;
/// POST /snapshot/prepare — quiesce subsystems before Firecracker snapshot.
///
/// In Rust there is no GC dance. We just:
/// 1. Drop page cache to shrink snapshot size
/// 2. Stop port subsystem
/// 3. Close idle connections via conntracker
/// 4. Set needs_restore flag
/// POST /snapshot/prepare — called by the host agent immediately before it
/// invokes vm.pause + vm.snapshot. The handler quiesces guest state so the
/// resulting snapshot is clean: outstanding writes are flushed to disk, the
/// VFS page cache is dropped (so the dm-snapshot CoW is the source of truth),
/// and the port forwarder is stopped to prevent socat children from being
/// frozen mid-handshake.
pub async fn post_snapshot_prepare(State(state): State<Arc<AppState>>) -> impl IntoResponse {
// Drop page cache BEFORE blocking the reclaimer — avoids snapshotting
// gigabytes of stale cache that inflates the memory dump on disk.
// "1" = pagecache only (keep dentries/inodes for faster resume).
if let Err(e) = std::fs::write("/proc/sys/vm/drop_caches", "1") {
tracing::warn!(error = %e, "snapshot/prepare: drop_caches failed");
} else {
tracing::info!("snapshot/prepare: page cache dropped");
// Stop port forwarder + scanner so no socat process is captured in the
// snapshot with a half-open TCP connection. /init on resume restarts it.
if let Some(ref port_sub) = state.port_subsystem {
port_sub.stop();
}
// Block memory reclaimer — prevents drop_caches from running mid-freeze
// which would corrupt kernel page table state.
state.snapshot_in_progress.store(true, Ordering::Release);
// sync(2) flushes the in-memory FS state. Done before drop_caches so the
// pages we drop are clean.
sync();
if let Some(ref ps) = state.port_subsystem {
ps.stop();
tracing::info!("snapshot/prepare: port subsystem stopped");
// Drop the VFS page cache + dentries/inodes. Reduces snapshot size by
// ensuring CH only persists memory pages that the guest actually needs.
if let Err(e) = std::fs::write("/proc/sys/vm/drop_caches", "3") {
tracing::warn!(error = %e, "drop_caches (first pass) failed (continuing)");
}
state.conn_tracker.prepare_for_snapshot();
tracing::info!("snapshot/prepare: connections prepared");
// Best-effort fstrim on the rootfs so unused blocks are returned to the
// dm-snapshot, keeping CoW size minimal.
let _ = tokio::process::Command::new("fstrim")
.arg("/")
.output()
.await;
// Sync filesystem buffers so dirty pages are flushed before freeze.
unsafe { libc::sync(); }
// Second drop_caches pass after fstrim: fstrim re-reads superblock /
// group descriptor pages that we just evicted, putting them back in the
// page cache. A second pass drops those and any other late readers (e.g.
// sync flushers).
sync();
if let Err(e) = std::fs::write("/proc/sys/vm/drop_caches", "3") {
tracing::warn!(error = %e, "drop_caches (second pass) failed (continuing)");
}
state.needs_restore.store(true, Ordering::Release);
tracing::info!("snapshot/prepare: ready for freeze");
// Free-page reporting drains asynchronously: the balloon driver hands
// freed pages to the host in batches and CH punches holes in the backing
// memfile. Without a brief settle window most of the pages freed by the
// drop_caches passes above would still be present in the snapshot.
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
tracing::info!("snapshot/prepare: quiesced");
(
StatusCode::NO_CONTENT,
[(header::CACHE_CONTROL, "no-store")],

View File

@ -6,7 +6,6 @@ mod config;
mod conntracker;
mod crypto;
mod execcontext;
mod host;
mod http;
mod logging;
mod permissions;
@ -22,7 +21,6 @@ use std::sync::Arc;
use clap::Parser;
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
use config::{DEFAULT_PORT, DEFAULT_USER, WRENN_RUN_DIR};
use execcontext::Defaults;
@ -44,9 +42,6 @@ struct Cli {
#[arg(long, default_value_t = DEFAULT_PORT)]
port: u16,
#[arg(long = "isnotfc", default_value_t = false)]
is_not_fc: bool,
#[arg(long)]
version: bool,
@ -73,35 +68,22 @@ async fn main() {
return;
}
let use_json = !cli.is_not_fc;
logging::init(use_json);
logging::init(true);
if let Err(e) = fs::create_dir_all(WRENN_RUN_DIR) {
tracing::error!(error = %e, "failed to create wrenn run directory");
}
let defaults = Defaults::new(DEFAULT_USER);
let is_fc_str = if cli.is_not_fc { "false" } else { "true" };
defaults
.env_vars
.insert("WRENN_SANDBOX".into(), is_fc_str.into());
.insert("WRENN_SANDBOX".into(), "true".into());
let wrenn_sandbox_path = Path::new(WRENN_RUN_DIR).join(".WRENN_SANDBOX");
if let Err(e) = fs::write(&wrenn_sandbox_path, is_fc_str.as_bytes()) {
if let Err(e) = fs::write(&wrenn_sandbox_path, b"true") {
tracing::error!(error = %e, "failed to write sandbox file");
}
let cancel = CancellationToken::new();
// MMDS polling (only in FC mode)
if !cli.is_not_fc {
let env_vars = Arc::clone(&defaults.env_vars);
let cancel_clone = cancel.clone();
tokio::spawn(async move {
host::mmds::poll_for_opts(env_vars, cancel_clone).await;
});
}
// Cgroup manager
let cgroup_manager: Arc<dyn cgroups::CgroupManager> =
match cgroups::Cgroup2Manager::new(
@ -143,14 +125,12 @@ async fn main() {
defaults,
VERSION.to_string(),
COMMIT.to_string(),
!cli.is_not_fc,
Some(Arc::clone(&port_subsystem)),
);
// Memory reclaimer — drop page cache when available memory is low.
// Firecracker balloon device can only reclaim pages the guest kernel freed.
// Pauses during snapshot/prepare to avoid corrupting kernel page table state.
if !cli.is_not_fc {
// The balloon device can only reclaim pages the guest kernel freed.
{
let state_for_reclaimer = Arc::clone(&state);
std::thread::spawn(move || memory_reclaimer(state_for_reclaimer));
}
@ -188,7 +168,6 @@ async fn main() {
}
port_subsystem.stop();
cancel.cancel();
}
fn spawn_initial_command(cmd: &str, state: &AppState) {
@ -231,19 +210,15 @@ fn spawn_initial_command(cmd: &str, state: &AppState) {
}
}
fn memory_reclaimer(state: Arc<AppState>) {
use std::sync::atomic::Ordering;
fn memory_reclaimer(_state: Arc<AppState>) {
use std::time::Duration;
const CHECK_INTERVAL: std::time::Duration = std::time::Duration::from_secs(10);
const CHECK_INTERVAL: Duration = Duration::from_secs(10);
const DROP_THRESHOLD_PCT: u64 = 80;
loop {
std::thread::sleep(CHECK_INTERVAL);
if state.snapshot_in_progress.load(Ordering::Acquire) {
continue;
}
let mut sys = sysinfo::System::new();
sys.refresh_memory();
let total = sys.total_memory();
@ -255,10 +230,6 @@ fn memory_reclaimer(state: Arc<AppState>) {
let used_pct = ((total - available) * 100) / total;
if used_pct >= DROP_THRESHOLD_PCT {
if state.snapshot_in_progress.load(Ordering::Acquire) {
continue;
}
if let Err(e) = std::fs::write("/proc/sys/vm/drop_caches", "3") {
tracing::debug!(error = %e, "drop_caches failed");
} else {

View File

@ -57,7 +57,9 @@ impl Scanner {
pub async fn scan_and_broadcast(&self, cancel: CancellationToken) {
loop {
let conns = read_tcp_connections();
let conns = tokio::task::spawn_blocking(read_tcp_connections)
.await
.unwrap_or_default();
{
let subs = self.subs.read().unwrap();

View File

@ -57,7 +57,11 @@ impl ProcessHandle {
}
pub fn send_signal(&self, sig: Signal) -> Result<(), ConnectError> {
signal::kill(Pid::from_raw(self.pid as i32), sig).map_err(|e| {
// Signal the whole process group (negative pid), not just the immediate
// /bin/sh wrapper. Otherwise children the process spawned are orphaned
// and keep running. Both spawn paths make the process a group leader
// (setsid for pty, setpgid for pipe), so pgid == pid.
signal::kill(Pid::from_raw(-(self.pid as i32)), sig).map_err(|e| {
ConnectError::new(ErrorCode::Internal, format!("error sending signal: {e}"))
})
}
@ -165,11 +169,22 @@ pub fn spawn_process(
env.push((k.clone(), v.clone()));
}
// Reset the child's nice value only when envd itself was started at an
// elevated nice value (delta > 0 means raising the nice number / lowering
// priority, which is permitted for non-root processes). A non-root process
// cannot improve its priority, so skip the `nice` wrapper otherwise — it
// would fail with EPERM ("cannot set niceness: permission denied") for
// commands run as a non-root user. Writing 100 to the process's own
// oom_score_adj is always permitted (raising the score).
let nice_delta = 0 - current_nice();
let oom_script = format!(
r#"echo 100 > /proc/$$/oom_score_adj && exec /usr/bin/nice -n {} "${{@}}""#,
nice_delta
);
let oom_script = if nice_delta > 0 {
format!(
r#"echo 100 > /proc/$$/oom_score_adj && exec /usr/bin/nice -n {} "${{@}}""#,
nice_delta
)
} else {
r#"echo 100 > /proc/$$/oom_score_adj && exec "$@""#.to_string()
};
let mut wrapper_args = vec![
"-c".to_string(),
oom_script,
@ -264,7 +279,7 @@ pub fn spawn_process(
let end_rx = handle.subscribe_end();
let data_tx_clone = data_tx.clone();
std::thread::spawn(move || {
let pty_reader = std::thread::spawn(move || {
let mut master = master_clone;
let mut buf = vec![0u8; PTY_CHUNK_SIZE];
loop {
@ -282,7 +297,11 @@ pub fn spawn_process(
let handle_for_waiter = Arc::clone(&handle);
std::thread::spawn(move || {
let mut child = child;
let end_event = match child.wait() {
let status = child.wait();
// Drain the pty to EOF before publishing the end event so trailing
// output is never lost to a process-exit/pty-read race.
let _ = pty_reader.join();
let end_event = match status {
Ok(s) => EndEvent {
exit_code: s.code().unwrap_or(-1),
exited: s.code().is_some(),
@ -320,6 +339,11 @@ pub fn spawn_process(
unsafe {
command.pre_exec(move || {
// Become a process-group leader so SendSignal can kill the
// whole group, not just this wrapper. The pty path gets this
// for free via setsid().
nix::unistd::setpgid(Pid::from_raw(0), Pid::from_raw(0))
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
libc::setgid(gid);
libc::setuid(uid);
Ok(())
@ -349,9 +373,11 @@ pub fn spawn_process(
let data_rx = handle.subscribe_data();
let end_rx = handle.subscribe_end();
let mut output_readers: Vec<std::thread::JoinHandle<()>> = Vec::new();
if let Some(mut out) = stdout {
let tx = data_tx.clone();
std::thread::spawn(move || {
output_readers.push(std::thread::spawn(move || {
let mut buf = vec![0u8; STD_CHUNK_SIZE];
loop {
match out.read(&mut buf) {
@ -362,12 +388,12 @@ pub fn spawn_process(
Err(_) => break,
}
}
});
}));
}
if let Some(mut err_pipe) = stderr {
let tx = data_tx.clone();
std::thread::spawn(move || {
output_readers.push(std::thread::spawn(move || {
let mut buf = vec![0u8; STD_CHUNK_SIZE];
loop {
match err_pipe.read(&mut buf) {
@ -378,13 +404,19 @@ pub fn spawn_process(
Err(_) => break,
}
}
});
}));
}
let end_tx_clone = end_tx.clone();
let handle_for_waiter = Arc::clone(&handle);
std::thread::spawn(move || {
let end_event = match child.wait() {
let status = child.wait();
// Drain stdout/stderr to EOF before publishing the end event so
// trailing output is never lost to a process-exit/pipe-read race.
for reader in output_readers {
let _ = reader.join();
}
let end_event = match status {
Ok(s) => EndEvent {
exit_code: s.code().unwrap_or(-1),
exited: s.code().is_some(),
@ -414,6 +446,8 @@ fn current_nice() -> i32 {
if *libc::__errno_location() != 0 {
return 0;
}
20 - prio
// getpriority(PRIO_PROCESS, 0) returns the nice value directly,
// in the range [-20, 19]; the normal default is 0.
prio
}
}

View File

@ -14,14 +14,14 @@ use crate::state::AppState;
pub struct ProcessServiceImpl {
state: Arc<AppState>,
processes: DashMap<u32, Arc<ProcessHandle>>,
processes: Arc<DashMap<u32, Arc<ProcessHandle>>>,
}
impl ProcessServiceImpl {
pub fn new(state: Arc<AppState>) -> Self {
Self {
state,
processes: DashMap::new(),
processes: Arc::new(DashMap::new()),
}
}
@ -131,11 +131,21 @@ impl ProcessServiceImpl {
self.processes.insert(spawned.handle.pid, Arc::clone(&spawned.handle));
let processes = self.processes.clone();
let processes = Arc::clone(&self.processes);
let pid = spawned.handle.pid;
// Subscribe before checking cached_end so the prune cannot be lost to a
// race: a short-lived process can exit and broadcast its end event
// before this task runs. A broadcast receiver only sees messages sent
// after subscribe(), so a late subscribe would miss the event forever
// (recv() never returns Closed either — the handle keeps end_tx alive
// until it leaves the map, which only this task does). The waiter sets
// ended before sending end_tx, so cached_end() is a reliable fallback.
let mut cleanup_end_rx = spawned.handle.subscribe_end();
let already_ended = spawned.handle.cached_end().is_some();
tokio::spawn(async move {
let _ = cleanup_end_rx.recv().await;
if !already_ended {
let _ = cleanup_end_rx.recv().await;
}
processes.remove(&pid);
});
@ -199,12 +209,28 @@ impl Process for ProcessServiceImpl {
match data {
Ok(ev) => yield Ok(make_data_start_response(ev)),
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
// Data channel closed: the process ended and its
// handle was dropped. The end event is published
// before the handle drop, so it is still buffered
// — emit it rather than losing the exit code.
if let Ok(end) = end_rx.try_recv() {
yield Ok(make_end_start_response(end));
}
break;
}
}
}
end = end_rx.recv() => {
while let Ok(ev) = data_rx.try_recv() {
yield Ok(make_data_start_response(ev));
// Process ended. The waiter joins the output readers
// before sending this event, so every byte is already
// in the data channel — drain it fully before the end.
loop {
match data_rx.try_recv() {
Ok(ev) => yield Ok(make_data_start_response(ev)),
Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
Err(_) => break,
}
}
if let Ok(end) = end {
yield Ok(make_end_start_response(end));
@ -268,15 +294,35 @@ impl Process for ProcessServiceImpl {
});
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
// Data channel closed: the process ended and
// its handle was dropped. The end event is
// published before the handle drop, so it is
// still buffered — emit it rather than losing
// the exit code.
if let Ok(end) = end_rx.try_recv() {
yield Ok(ConnectResponse {
event: buffa::MessageField::some(make_end_event(end)),
..Default::default()
});
}
break;
}
}
}
end = end_rx.recv() => {
while let Ok(ev) = data_rx.try_recv() {
yield Ok(ConnectResponse {
event: buffa::MessageField::some(make_data_event(ev)),
..Default::default()
});
// Process ended. The waiter joins the output readers
// before sending this event, so every byte is already
// in the data channel — drain it fully before the end.
loop {
match data_rx.try_recv() {
Ok(ev) => yield Ok(ConnectResponse {
event: buffa::MessageField::some(make_data_event(ev)),
..Default::default()
}),
Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
Err(_) => break,
}
}
if let Ok(end) = end {
yield Ok(ConnectResponse {

View File

@ -1,5 +1,5 @@
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicU8, Ordering};
use std::sync::{Arc, Mutex};
use crate::auth::token::SecureToken;
use crate::conntracker::ConnTracker;
@ -11,15 +11,32 @@ pub struct AppState {
pub defaults: Defaults,
pub version: String,
pub commit: String,
pub is_fc: bool,
pub needs_restore: AtomicBool,
pub last_set_time: AtomicMax,
pub access_token: SecureToken,
pub conn_tracker: ConnTracker,
pub port_subsystem: Option<Arc<PortSubsystem>>,
pub cpu_used_pct: AtomicU32,
pub cpu_count: AtomicU32,
pub snapshot_in_progress: AtomicBool,
/// Memory preload coordination. The host agent POSTs /memory/preload after
/// a snapshot restore to materialise every physical page (so the next
/// ch.snapshot writes a self-contained memory-ranges). `mem_preload_started`
/// ensures only one loader runs; `mem_preload_done` lets concurrent callers
/// rendezvous; `mem_preload_cancel` lets a teardown abort the loader.
pub mem_preload_started: AtomicBool,
pub mem_preload_done: AtomicBool,
pub mem_preload_cancel: AtomicBool,
pub mem_preload_regions: AtomicU64,
pub mem_preload_pages: AtomicU64,
pub mem_preload_bytes: AtomicU64,
pub mem_preload_elapsed_us: AtomicU64,
/// 0 = unset, 1 = /dev/mem, 2 = /proc/kcore.
pub mem_preload_source: AtomicU8,
pub mem_preload_error: Mutex<Option<String>>,
/// Last lifecycle ID seen on /init. Used to detect post-resume calls so
/// envd can refresh port forwarders and remount NFS volumes.
lifecycle_id: Mutex<Option<String>>,
}
impl AppState {
@ -27,22 +44,28 @@ impl AppState {
defaults: Defaults,
version: String,
commit: String,
is_fc: bool,
port_subsystem: Option<Arc<PortSubsystem>>,
) -> Arc<Self> {
let state = Arc::new(Self {
defaults,
version,
commit,
is_fc,
needs_restore: AtomicBool::new(false),
last_set_time: AtomicMax::new(),
access_token: SecureToken::new(),
conn_tracker: ConnTracker::new(),
port_subsystem,
cpu_used_pct: AtomicU32::new(0),
cpu_count: AtomicU32::new(0),
snapshot_in_progress: AtomicBool::new(false),
mem_preload_started: AtomicBool::new(false),
mem_preload_done: AtomicBool::new(false),
mem_preload_cancel: AtomicBool::new(false),
mem_preload_regions: AtomicU64::new(0),
mem_preload_pages: AtomicU64::new(0),
mem_preload_bytes: AtomicU64::new(0),
mem_preload_elapsed_us: AtomicU64::new(0),
mem_preload_source: AtomicU8::new(0),
mem_preload_error: Mutex::new(None),
lifecycle_id: Mutex::new(None),
});
let state_clone = Arc::clone(&state);
@ -60,6 +83,20 @@ impl AppState {
pub fn cpu_count(&self) -> u32 {
self.cpu_count.load(Ordering::Relaxed)
}
/// Records a new lifecycle ID, returning true if it changed (i.e. this
/// is the first /init since a resume). First-ever call returns false:
/// boot-time /init doesn't need port-subsystem restart since the
/// subsystem hasn't been started yet by anything else.
pub fn bump_lifecycle(&self, new_id: &str) -> bool {
let mut guard = self.lifecycle_id.lock().unwrap();
let changed = match guard.as_deref() {
Some(existing) => existing != new_id,
None => false,
};
*guard = Some(new_id.to_owned());
changed
}
}
fn cpu_sampler(state: Arc<AppState>) {
@ -70,6 +107,7 @@ fn cpu_sampler(state: Arc<AppState>) {
loop {
std::thread::sleep(std::time::Duration::from_secs(1));
sys.refresh_cpu_all();
let pct = sys.global_cpu_usage();