From f366c38f17f4f6380776cc102a718249adc125ee Mon Sep 17 00:00:00 2001 From: Benjamin Arntzen Date: Tue, 30 Jun 2026 12:46:35 +0100 Subject: [PATCH] =?UTF-8?q?artifact=20cache:=20zero-copy=20in-memory=20ser?= =?UTF-8?q?ve=20to=20kill=20N=C3=97=20concurrent=20disk=20reads?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under concurrent PXE boot, N machines pull the same large artifact (modloop ~295 MiB, initramfs, vmlinuz, apkovl) at once. read_file_as_stream served each from disk independently, so N concurrent clients re-read the same file N times off one PVE-backed store (loop0 over vm-1000-disk-0.raw) — disk throughput collapsed and stalled every consumer. The intended single-flight pullthrough (stream_download_with_caching) is dead code (0 fires; the iPXE script points clients at /boot/{arch}/* local files, not /ipxe-artifacts), so it never deduped the reads. Add ArtifactCache (on AppState): load each file into memory ONCE into a shared Bytes; every concurrent client streams Bytes::slice zero-copy views of that one allocation (refcounted — CoW-on-write, never written) in 64 KiB chunks. Population is single-flight (per-path tokio::Mutex, double-checked) so the cold read happens once and concurrent readers await then slice the shared buffer — no duplicate loads, no torn reads. Bounded LRU (default 1 GiB, DRAGONFLY_ARTIFACT_CACHE_BYTES); files over the cap stream from disk unchanged. read_file_as_stream now serves from the cache when AppState is available and the file fits (early-return on Cached, fall through to the existing disk path on TooLarge/error). This also unifies the range + full branches for cached files: range requests slice the shared buffer instead of read_to_end into a per-request Vec — the "well-designed range" the per-request buffering was supposed to be. Wired serve_boot_asset (already passed Some(state)), serve_cached_image, and serve_os_image (signature + the /os route closures) to reach the cache. Tests (artifact_cache, 5): load-once-then-memory, single-flight (16 concurrent callers share one allocation — proven by Bytes pointer equality), zero-copy slice, LRU eviction, oversized-not-cached. dragonfly-server lib 198/198 green; fmt clean; no new clippy. Co-Authored-By: Claude --- crates/dragonfly-server/src/api.rs | 79 ++++- crates/dragonfly-server/src/artifact_cache.rs | 313 ++++++++++++++++++ crates/dragonfly-server/src/lib.rs | 34 +- crates/dragonfly-server/src/test_helpers.rs | 1 + 4 files changed, 417 insertions(+), 10 deletions(-) create mode 100644 crates/dragonfly-server/src/artifact_cache.rs diff --git a/crates/dragonfly-server/src/api.rs b/crates/dragonfly-server/src/api.rs index e439e217..4af07884 100644 --- a/crates/dragonfly-server/src/api.rs +++ b/crates/dragonfly-server/src/api.rs @@ -2697,6 +2697,35 @@ async fn stream_artifact( } } +/// Resolve a Range header against a known total size into +/// `(start, end, response_length, content_range_header)`. Whole file when the +/// header is absent or invalid. Used by the in-memory cache path (which knows +/// the size from the cached buffer, without a disk `metadata` call). +async fn resolve_byte_range( + range_header: Option<&HeaderValue>, + total_size: u64, +) -> (u64, u64, u64, Option) { + let whole = (0, total_size.saturating_sub(1), total_size, None); + let Some(rv) = range_header else { + return whole; + }; + let Ok(range_str) = rv.to_str() else { + return whole; + }; + match parse_range_header(range_str, total_size, None, None).await { + Some((start, end)) => { + let length = end - start + 1; + ( + start, + end, + length, + Some(format!("bytes {}-{}/{}", start, end, total_size)), + ) + } + None => whole, + } +} + async fn read_file_as_stream( path: &StdPath, range_header: Option<&HeaderValue>, // Add parameter for Range header @@ -2718,6 +2747,45 @@ async fn read_file_as_stream( machine_id ); + // Zero-copy in-memory cache: one disk read per file, N concurrent clients + // stream `Bytes::slice` views of a single shared allocation (refcounted — + // no copy, no per-request disk read). Falls through to the disk path below + // for files larger than the cache cap (`TooLarge`) or if the load fails. + if let Some(app) = state { + match app.artifact_cache.get_or_load(path).await { + Ok(crate::artifact_cache::GetResult::Cached(bytes)) => { + let total_size = bytes.len() as u64; + let (start, end, response_length, content_range_header) = + resolve_byte_range(range_header, total_size).await; + let slice = if total_size == 0 { + bytes.slice(0..0) + } else { + bytes.slice(start as usize..(end as usize + 1)) + }; + let (tx, rx) = mpsc::channel::>(32); + crate::artifact_cache::spawn_zero_copy_stream(slice, tx); + return Ok(( + tokio_stream::wrappers::ReceiverStream::new(rx), + Some(response_length), + content_range_header, + )); + } + Ok(crate::artifact_cache::GetResult::TooLarge(_)) => { + info!( + "[STREAM_READ] {} exceeds cache cap; streaming from disk", + path.display() + ); + } + Err(e) => { + warn!( + "[STREAM_READ] cache load failed for {}: {}; streaming from disk", + path.display(), + e + ); + } + } + } + let mut file = fs::File::open(path) .await .map_err(|e| Error::Internal(format!("Failed to open file {}: {}", path.display(), e)))?; // Added mut back @@ -5661,7 +5729,12 @@ pub async fn download_ipxe_binaries() -> anyhow::Result<()> { const OS_IMAGES_DIR: &str = "/var/lib/dragonfly/os-images"; /// Serve OS image file -pub async fn serve_os_image(os: &str, arch: &str, range: Option<&HeaderValue>) -> Response { +pub async fn serve_os_image( + os: &str, + arch: &str, + range: Option<&HeaderValue>, + state: &AppState, +) -> Response { let path = match os { "debian-13" => { let filename = format!("debian-13-generic-{}.tar.xz", arch); @@ -5683,7 +5756,7 @@ pub async fn serve_os_image(os: &str, arch: &str, range: Option<&HeaderValue>) - // OS images are multi-GB and every imaging machine pulls one concurrently; // stream in 64 KiB chunks (constant memory + Range) instead of buffering // the whole file. See `stream_artifact`. - stream_artifact(&path, "application/x-xz", range, None, None).await + stream_artifact(&path, "application/x-xz", range, Some(state), None).await } /// Serve cached image file (JIT-converted QCOW2 to raw) @@ -5713,7 +5786,7 @@ pub async fn serve_cached_image( &canonical, content_type, headers.get(axum::http::header::RANGE), - None, + Some(&state), None, ) .await diff --git a/crates/dragonfly-server/src/artifact_cache.rs b/crates/dragonfly-server/src/artifact_cache.rs new file mode 100644 index 00000000..dae25ba3 --- /dev/null +++ b/crates/dragonfly-server/src/artifact_cache.rs @@ -0,0 +1,313 @@ +//! Zero-copy in-memory cache of served artifacts (boot assets, OS/cached images). +//! +//! Under concurrent PXE boot, N machines pull the same large artifact (e.g. +//! modloop, ~295 MiB) at once. Serving each from disk independently re-reads +//! the file N times off one PVE-backed store, collapsing throughput and +//! stalling every consumer. This cache loads each file into memory ONCE into a +//! single shared `Bytes`; every concurrent client then streams zero-copy +//! `Bytes::slice` views of that one allocation — refcounted (CoW-on-write, +//! never written), so N clients pay for one disk read, not N. +//! +//! Population is single-flight: the first caller loads; concurrent callers +//! await a per-path lock, double-check, then clone the shared buffer (no +//! duplicate loads, no torn reads). The cache is size-bounded: when an insert +//! would exceed the cap, least-recently-used entries are evicted. Files larger +//! than the cap are reported as [`GetResult::TooLarge`] so the caller streams +//! them from disk via the existing path rather than caching. + +use bytes::Bytes; +use std::collections::HashMap; +use std::io; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +/// Default cache cap (1 GiB). Fits the boot-asset set (~350 MiB) with headroom +/// on a 2 GiB host; override with `DRAGONFLY_ARTIFACT_CACHE_BYTES`. +pub const DEFAULT_MAX_BYTES: u64 = 1024 * 1024 * 1024; +/// Streaming chunk size for the zero-copy fan-out (matches the prior disk path). +pub const STREAM_CHUNK: usize = 65_536; + +/// One cached artifact: its shared bytes + LRU bookkeeping. +struct Entry { + bytes: Bytes, + last_used: Instant, + len: u64, +} + +/// Outcome of [`ArtifactCache::get_or_load`]. +pub enum GetResult { + /// The file's bytes, shared (clone the `Bytes` for a zero-copy view). + Cached(Bytes), + /// File is larger than the cache cap (carries its length); caller should + /// stream it from disk instead of caching. + TooLarge(u64), +} + +pub struct ArtifactCache { + /// Completed loads, keyed by canonical path. + entries: Mutex>, + /// Per-path single-flight locks. A path is present while a load is + /// rendezvousing; concurrent callers clone the same `Arc` and await. + /// Cardinality is bounded by the distinct-file set actually served, so + /// lingering entries (a few dozen boot assets/images) are negligible. + in_flight: Mutex>>>, + max_bytes: u64, + current_bytes: AtomicU64, +} + +impl ArtifactCache { + pub fn new(max_bytes: u64) -> Self { + Self { + entries: Mutex::new(HashMap::new()), + in_flight: Mutex::new(HashMap::new()), + max_bytes, + current_bytes: AtomicU64::new(0), + } + } + + /// Read the cap from `DRAGONFLY_ARTIFACT_CACHE_BYTES`, else [`DEFAULT_MAX_BYTES`]. + pub fn max_bytes_from_env() -> u64 { + std::env::var("DRAGONFLY_ARTIFACT_CACHE_BYTES") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|&v| v > 0) + .unwrap_or(DEFAULT_MAX_BYTES) + } + + /// Fast path: clone the cached bytes (refcount bump — zero-copy) and bump + /// last-used. `None` if absent. + fn get_cached(&self, path: &Path) -> Option { + let mut entries = self.entries.lock().ok()?; + let e = entries.get_mut(path)?; + e.last_used = Instant::now(); + Some(e.bytes.clone()) + } + + /// Get-or-create the per-path single-flight lock. + fn inflight_for(&self, path: &Path) -> Arc> { + let mut inflight = self.in_flight.lock().expect("in_flight mutex poisoned"); + inflight + .entry(path.to_path_buf()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + } + + /// Get the file's bytes, loading it once (single-flight) if absent. Returns + /// [`GetResult::TooLarge`] (uncached) when the file exceeds the cap. + pub async fn get_or_load(&self, path: &Path) -> Result { + // Fast path: already cached. + if let Some(b) = self.get_cached(path) { + return Ok(GetResult::Cached(b)); + } + // Single-flight: serialize concurrent first-loaders on this path. + let lock = self.inflight_for(path); + let _guard = lock.lock().await; + // Double-check after acquiring — another waiter may have just loaded it. + if let Some(b) = self.get_cached(path) { + return Ok(GetResult::Cached(b)); + } + // Load once. + let vec = tokio::fs::read(path).await?; + let len = vec.len() as u64; + if len > self.max_bytes { + return Ok(GetResult::TooLarge(len)); + } + let bytes = Bytes::from(vec); + self.insert(path, bytes.clone(), len); + Ok(GetResult::Cached(bytes)) + } + + /// Insert an entry, evicting least-recently-used entries until it fits. + fn insert(&self, path: &Path, bytes: Bytes, len: u64) { + let mut entries = self.entries.lock().expect("entries mutex poisoned"); + // If refreshing an existing path, account for the replaced length. + if let Some(old) = entries.insert( + path.to_path_buf(), + Entry { + bytes, + last_used: Instant::now(), + len, + }, + ) { + self.current_bytes.fetch_sub(old.len, Ordering::Relaxed); + } + self.current_bytes.fetch_add(len, Ordering::Relaxed); + // Evict LRU until within cap. Never evict the just-inserted entry. + while self.current_bytes.load(Ordering::Relaxed) > self.max_bytes { + let victim = entries + .iter() + .filter(|(k, _)| *k != path) + .min_by_key(|(_, e)| e.last_used) + .map(|(k, _)| k.clone()); + match victim { + Some(k) => { + if let Some(e) = entries.remove(&k) { + self.current_bytes.fetch_sub(e.len, Ordering::Relaxed); + } + } + None => break, // only the just-inserted entry remains + } + } + } + + /// Current total cached bytes (for tests/observability). + pub fn current_bytes(&self) -> u64 { + self.current_bytes.load(Ordering::Relaxed) + } +} + +/// Stream `slice` in [`STREAM_CHUNK`]-sized zero-copy chunks through `tx`. +/// Each chunk is a `Bytes::split_to` view of the shared allocation — no copy, +/// no disk. The receiver drains at the client's read rate (channel backpressure). +pub fn spawn_zero_copy_stream( + slice: Bytes, + tx: tokio::sync::mpsc::Sender>, +) { + tokio::spawn(async move { + let mut remaining = slice; + while !remaining.is_empty() { + let take = std::cmp::min(remaining.len(), STREAM_CHUNK); + let chunk = remaining.split_to(take); + if tx.send(Ok(chunk)).await.is_err() { + return; // receiver dropped (client disconnected) + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::fs; + + async fn write_tmp(name: &str, contents: &[u8]) -> PathBuf { + let dir = std::env::temp_dir().join(format!( + "dragonfly-artifact-cache-test-{}", + std::process::id() + )); + fs::create_dir_all(&dir).await.unwrap(); + let p = dir.join(name); + fs::write(&p, contents).await.unwrap(); + p + } + + // Two Bytes share one allocation iff their data pointers are equal. A + // re-read would allocate fresh -> different pointer. So pointer equality + // is the proof of "served from the cache, not re-read from disk". + fn share_allocation(a: &Bytes, b: &Bytes) -> bool { + a.as_ptr() == b.as_ptr() + } + + #[tokio::test] + async fn loads_on_miss_then_serves_subsequent_calls_from_memory() { + let cache = ArtifactCache::new(1024); + let path = write_tmp("load_once.bin", b"hello artifact cache").await; + let first = match cache.get_or_load(&path).await.unwrap() { + GetResult::Cached(b) => b, + _ => panic!("expected Cached"), + }; + let second = match cache.get_or_load(&path).await.unwrap() { + GetResult::Cached(b) => b, + _ => panic!("expected Cached"), + }; + assert!( + share_allocation(&first, &second), + "second call re-read from disk" + ); + let _ = fs::remove_file(&path).await; + } + + #[tokio::test] + async fn single_flight_loads_once_for_n_concurrent_callers() { + let cache = Arc::new(ArtifactCache::new(65_536)); + let path = Arc::new(write_tmp("singleflight.bin", &[7u8; 4096]).await); + + // N concurrent first-loaders on the same cold path. + let results = futures::future::join_all((0..16).map(|_| { + let cache = Arc::clone(&cache); + let path = Arc::clone(&path); + async move { cache.get_or_load(&path).await.unwrap() } + })) + .await; + + // Every caller must have received the SAME allocation — one load, not 16. + let first_ptr = match &results[0] { + GetResult::Cached(b) => b.as_ptr(), + _ => panic!("expected Cached"), + }; + for r in &results { + match r { + GetResult::Cached(b) => assert_eq!( + b.as_ptr(), + first_ptr, + "concurrent caller got an independent allocation (no single-flight)" + ), + _ => panic!("expected Cached"), + } + } + let _ = fs::remove_file(&*path).await; + } + + #[tokio::test] + async fn slice_is_zero_copy() { + let cache = ArtifactCache::new(1_048_576); + let path = write_tmp("slice.bin", &[9u8; 100_000]).await; + let full = match cache.get_or_load(&path).await.unwrap() { + GetResult::Cached(b) => b, + _ => panic!("expected Cached"), + }; + // A slice shares the underlying allocation (offset pointer from base). + let mid = full.slice(40_000..60_000); + assert_eq!(mid.as_ptr() as usize, full.as_ptr() as usize + 40_000); + assert_eq!(&mid[..3], &full[40_000..40_003], "slice content differs"); + let _ = fs::remove_file(&path).await; + } + + #[tokio::test] + async fn evicts_lru_when_cap_exceeded() { + // Cap holds exactly one file of this size. + const SIZE: usize = 1024; + let cache = ArtifactCache::new(SIZE as u64); + let a = write_tmp("a.bin", &[1u8; SIZE]).await; + let b = write_tmp("b.bin", &[2u8; SIZE]).await; + + let first_a = match cache.get_or_load(&a).await.unwrap() { + GetResult::Cached(b) => b, + _ => panic!("expected Cached"), + }; + // Loading B (same size) exceeds the cap -> evicts A (the LRU). + let _ = match cache.get_or_load(&b).await.unwrap() { + GetResult::Cached(b) => b, + _ => panic!("expected Cached"), + }; + // A was evicted, so re-loading it must allocate fresh (new pointer). + let second_a = match cache.get_or_load(&a).await.unwrap() { + GetResult::Cached(b) => b, + _ => panic!("expected Cached"), + }; + assert!( + !share_allocation(&first_a, &second_a), + "A was not evicted (re-load returned the cached allocation)" + ); + let _ = fs::remove_file(&a).await; + let _ = fs::remove_file(&b).await; + } + + #[tokio::test] + async fn oversized_file_is_not_cached() { + let cache = ArtifactCache::new(100); + let path = write_tmp("oversize.bin", &[0u8; 200]).await; + match cache.get_or_load(&path).await.unwrap() { + GetResult::TooLarge(len) => assert_eq!(len, 200), + GetResult::Cached(_) => panic!("oversized file was cached"), + } + assert_eq!( + cache.current_bytes(), + 0, + "oversized file counted toward cap" + ); + let _ = fs::remove_file(&path).await; + } +} diff --git a/crates/dragonfly-server/src/lib.rs b/crates/dragonfly-server/src/lib.rs index b3d26757..1d42a769 100644 --- a/crates/dragonfly-server/src/lib.rs +++ b/crates/dragonfly-server/src/lib.rs @@ -1,7 +1,10 @@ use anyhow::{Context, anyhow}; use axum::extract::{MatchedPath, Path}; use axum::middleware::from_fn; -use axum::{Router, extract::Extension, http::StatusCode, response::IntoResponse, routing::get}; +use axum::{ + Router, extract::Extension, extract::State, http::StatusCode, response::IntoResponse, + routing::get, +}; use axum_login::AuthManagerLayerBuilder; use listenfd::ListenFd; use std::net::SocketAddr; @@ -42,6 +45,7 @@ use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt}; pub mod agent_ws; mod api; pub mod api_token; +pub mod artifact_cache; mod auth; pub mod cluster; pub mod dns_sync; @@ -339,6 +343,9 @@ pub struct AppState { pub network_services_started: Arc, // Image cache for JIT QCOW2 conversion pub image_cache: Arc, + // Zero-copy in-memory cache of served artifacts (boot assets, OS/cached + // images): one disk read per file, N clients stream zero-copy slices. + pub artifact_cache: Arc, // Shutdown sender for network services (allows independent restart) pub services_shutdown_tx: Arc>>>, // Shared DHCP lease table — survives service restarts, queryable from API @@ -1284,6 +1291,9 @@ pub async fn run() -> anyhow::Result<()> { network_services_started: Arc::new(AtomicBool::new(false)), // Image cache for JIT QCOW2 conversion image_cache: image_cache.clone(), + artifact_cache: Arc::new(artifact_cache::ArtifactCache::new( + artifact_cache::ArtifactCache::max_bytes_from_env(), + )), // Services shutdown sender for independent restart services_shutdown_tx: Arc::new(Mutex::new(None)), // Shared DHCP lease table @@ -1356,16 +1366,26 @@ pub async fn run() -> anyhow::Result<()> { // OS images (served during provisioning) .route( "/os/debian-13/amd64", - get(|headers: axum::http::HeaderMap| async move { - api::serve_os_image("debian-13", "amd64", headers.get(axum::http::header::RANGE)) - .await + get(|State(app_state): State, headers: axum::http::HeaderMap| async move { + api::serve_os_image( + "debian-13", + "amd64", + headers.get(axum::http::header::RANGE), + &app_state, + ) + .await }), ) .route( "/os/debian-13/arm64", - get(|headers: axum::http::HeaderMap| async move { - api::serve_os_image("debian-13", "arm64", headers.get(axum::http::header::RANGE)) - .await + get(|State(app_state): State, headers: axum::http::HeaderMap| async move { + api::serve_os_image( + "debian-13", + "arm64", + headers.get(axum::http::header::RANGE), + &app_state, + ) + .await }), ) // Cached images (JIT-converted QCOW2 to raw) diff --git a/crates/dragonfly-server/src/test_helpers.rs b/crates/dragonfly-server/src/test_helpers.rs index d69c1e05..b45ce58e 100644 --- a/crates/dragonfly-server/src/test_helpers.rs +++ b/crates/dragonfly-server/src/test_helpers.rs @@ -63,6 +63,7 @@ pub async fn create_test_app_state() -> AppState { store_proxy, network_services_started: Arc::new(AtomicBool::new(false)), image_cache, + artifact_cache: Arc::new(crate::artifact_cache::ArtifactCache::new(1024 * 1024)), services_shutdown_tx: Arc::new(Mutex::new(None)), dhcp_lease_table: Arc::new(tokio::sync::RwLock::new(dragonfly_dhcp::LeaseTable::new())), ha_manager: Arc::new(ha::HaManager::new("test-node".to_string())),