From 77ecaf63b37dec22eea5b9d8fee17b5621fc406c Mon Sep 17 00:00:00 2001 From: Adrien Date: Mon, 28 Sep 2026 17:43:47 +0200 Subject: [PATCH 1/2] fix(xet): give each read stream its own download buffer (#234) Every read stream competed FIFO for xet-core's single global download buffer, which hf-mount capped at 256 MiB. Buffer held by downloaded but unconsumed data is only released when a FUSE worker thread consumes it, and under many concurrent cold readers those threads were all blocked waiting on streams that could not get their first term. Reads timed out at cursor=0, retried at the back of the queue, and ended in EIO. Each stream now gets a private AdjustableSemaphore via FileReconstructor::with_buffer_semaphore, sized as a fair share of the download buffer limit (1 GiB by default) clamped to [one xorb, 256 MiB]. The floor cannot go below one xorb: xet-core clamps a term acquire only when it is issued, so a shrink below an in-flight acquire would leave it unsatisfiable. Measured on c7gn.16xlarge, 128 x 1 GiB files, parallel -j64 sha512sum -c: main 996 s at 1.2 Gbps with 108 stream timeouts; fix 59 s at 19.3 Gbps with none. Single-stream reads go from 0.8-1.1 Gbps to 1.2-1.35 Gbps. Adds test_fuse_parallel_cold_reads (24 cold readers on 4 FUSE threads), which fails on main with 13 of 24 readers in EIO. --- src/setup.rs | 3 +- src/xet.rs | 193 +++++++++++++++++++++++++++++++++++++++++--- tests/common/mod.rs | 116 +++++++++++++++++++++++++- tests/fuse_ops.rs | 103 +++++++++++++++++++++++ 4 files changed, 404 insertions(+), 11 deletions(-) diff --git a/src/setup.rs b/src/setup.rs index 2ec7446e..2a0d4f05 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -364,7 +364,8 @@ pub fn init_tracing(daemon: bool) { ("HF_XET_RECONSTRUCTION_MIN_PREFETCH_BUFFER", "8388608"), ("HF_XET_RECONSTRUCTION_TARGET_BLOCK_COMPLETION_TIME", "30"), ("HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_SIZE", "134217728"), - ("HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_LIMIT", "268435456"), + // Also the budget split between per-stream read buffers (see xet.rs). + ("HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_LIMIT", "1073741824"), // Per-read inactivity timeout for CAS/CDN transfers (resets on every byte // received, so slow-but-progressing reads are fine). This governs the // DOWNLOAD/reconstruction path (term fetches and whole-file downloads); diff --git a/src/xet.rs b/src/xet.rs index a3019f8f..1d1e7d79 100644 --- a/src/xet.rs +++ b/src/xet.rs @@ -1,16 +1,18 @@ use std::path::{Path, PathBuf}; -use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use bytes::Bytes; use xet_client::cas_client::Client; use xet_client::cas_types::FileRange; use xet_client::chunk_cache::ChunkCache; use xet_core_structures::merklehash::MerkleHash; +use xet_core_structures::xorb_object::constants::MAX_XORB_BYTES; use xet_data::file_reconstruction::{DownloadStream, FileReconstructor}; use xet_data::processing::configurations::TranslatorConfig; use xet_data::processing::{FileDownloadSession, FileUploadSession, Sha256Policy, SingleFileCleaner, XetFileInfo}; use xet_runtime::core::XetContext; +use xet_runtime::utils::adjustable_semaphore::AdjustableSemaphore; use crate::error::{Error, Result}; @@ -48,6 +50,110 @@ pub trait DownloadStreamOps: Send { async fn next(&mut self) -> Result>; } +// ── Per-stream download buffers ─────────────────────────────────────── + +/// Ceiling of a single stream's buffer: pipelining depth of a lone reader. +const STREAM_BUFFER_MAX: u64 = 256 * 1_048_576; + +/// Floor of a single stream's buffer: one full xorb, the largest possible +/// term. xet-core clamps a term acquire against the semaphore total only +/// when it is issued (`AdjustableSemaphore::to_physical_acquire`) and does +/// not re-clamp pending acquires on shrink, so shrinking below the largest +/// term would leave an in-flight acquire that can never be satisfied. The +/// floor can become memory-driven once that is fixed upstream. +fn stream_buffer_min() -> u64 { + *MAX_XORB_BYTES as u64 +} + +/// Gives every read stream a private download buffer. +/// +/// Streams are consumed by FUSE `read()` calls that each pin a worker thread. +/// With xet-core's single global buffer (FIFO), a stream that already holds +/// buffer cannot release it until its consumer gets a worker thread, which +/// may itself be blocked waiting for buffer on another stream. Under many +/// concurrent cold readers this starves new streams at their first byte and +/// wedges the mount (#234). Private buffers mean streams never wait on each +/// other. Each is sized `budget / active_streams`, clamped to +/// `[stream_buffer_min(), STREAM_BUFFER_MAX]`, so memory in flight is bounded +/// by `max(budget, streams * stream_buffer_min())`. +struct StreamBufferPool { + budget: u64, + inner: Mutex, +} + +#[derive(Default)] +struct StreamBuffers { + active: Vec>, + /// Share currently applied to every active buffer. + share: u64, +} + +impl StreamBufferPool { + fn new(budget: u64) -> Arc { + Arc::new(Self { + budget, + inner: Mutex::new(StreamBuffers::default()), + }) + } + + /// Create a buffer for a new stream, already sized to the new fair + /// share, and resize the other active buffers to match. The guard + /// unregisters the buffer on drop. + fn register(self: &Arc) -> StreamBufferGuard { + let mut inner = self.inner.lock().expect("stream buffers poisoned"); + let share = self.share_for(inner.active.len() + 1); + let buffer = AdjustableSemaphore::new(share, (stream_buffer_min(), STREAM_BUFFER_MAX)); + inner.active.push(buffer.clone()); + Self::apply_share(&mut inner, share); + StreamBufferGuard { + pool: self.clone(), + buffer, + } + } + + fn unregister(&self, buffer: &Arc) { + let mut inner = self.inner.lock().expect("stream buffers poisoned"); + if let Some(index) = inner.active.iter().position(|other| Arc::ptr_eq(other, buffer)) { + inner.active.swap_remove(index); + } + let share = self.share_for(inner.active.len()); + Self::apply_share(&mut inner, share); + } + + /// Resize every active buffer to `share`; a no-op when it is unchanged. + /// Shrinks apply lazily: permits already held are reclaimed as they return. + fn apply_share(inner: &mut StreamBuffers, share: u64) { + if inner.share == share { + return; + } + inner.share = share; + for buffer in &inner.active { + // Each call is a no-op when the target is on the other side of + // the current total; the increment's virtual permit releases the + // added capacity on drop. + drop(buffer.increment_permits_to_target(share)); + buffer.decrement_permits_to_target(share); + } + } + + fn share_for(&self, active_streams: usize) -> u64 { + (self.budget / active_streams.max(1) as u64).clamp(stream_buffer_min(), STREAM_BUFFER_MAX) + } +} + +/// Keeps a stream's buffer registered in its pool; dropping it (with the +/// stream) hands the freed share back to the remaining streams. +struct StreamBufferGuard { + pool: Arc, + buffer: Arc, +} + +impl Drop for StreamBufferGuard { + fn drop(&mut self) { + self.pool.unregister(&self.buffer); + } +} + // ── XetSessions ─────────────────────────────────────────────────────── /// Core xet-core sessions for CAS downloads and uploads. @@ -61,6 +167,7 @@ pub struct XetSessions { /// Chunk cache attached to unbounded streams; bounded range downloads skip it /// to avoid pulling whole xorbs for small range requests. chunk_cache: Option>, + stream_buffers: Arc, } impl XetSessions { @@ -71,34 +178,43 @@ impl XetSessions { cas_client: Arc, chunk_cache: Option>, ) -> Arc { + // The same knob xet-core uses for its global buffer, so the existing + // HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_LIMIT override keeps working. + let stream_buffers = StreamBufferPool::new(ctx.config.reconstruction.download_buffer_limit.as_u64()); Arc::new(Self { ctx, session, upload_config, cas_client, chunk_cache, + stream_buffers, }) } /// Start a streaming download for a byte range. /// When `end` is `Some`, only bytes `[offset, end)` are fetched (bounded range). /// When `end` is `None`, fetches from `offset` to end of file (unbounded stream). - pub fn download_stream(&self, file_info: &XetFileInfo, offset: u64, end: Option) -> Result { + fn download_stream(&self, file_info: &XetFileInfo, offset: u64, end: Option) -> Result { let hash = file_info .merkle_hash() .map_err(|e| Error::Xet(format!("invalid hash: {e}")))?; let is_unbounded = end.is_none(); let file_size = file_info.file_size().unwrap_or(u64::MAX); let end = end.unwrap_or(file_size); - let mut reconstructor = - FileReconstructor::new(&self.ctx, &self.cas_client, hash).with_byte_range(FileRange::new(offset, end)); + let buffer = self.stream_buffers.register(); + let mut reconstructor = FileReconstructor::new(&self.ctx, &self.cas_client, hash) + .with_byte_range(FileRange::new(offset, end)) + .with_buffer_semaphore(buffer.buffer.clone()); // Attach chunk cache only to the unbounded stream path: the xorb disk // cache pulls full xorbs (~64MB) even for small range requests, which // is wasteful for random reads. Sequential reads (unbounded) benefit. if is_unbounded && let Some(cache) = self.chunk_cache.as_ref() { reconstructor = reconstructor.with_chunk_cache(cache.clone()); } - Ok(reconstructor.reconstruct_to_stream()) + Ok(DownloadStreamWrapper { + stream: reconstructor.reconstruct_to_stream(), + _buffer: buffer, + }) } } @@ -146,8 +262,7 @@ impl XetOps for XetSessions { offset: u64, end: Option, ) -> Result> { - let stream = self.download_stream(file_info, offset, end)?; - Ok(Box::new(DownloadStreamWrapper(stream))) + Ok(Box::new(self.download_stream(file_info, offset, end)?)) } async fn warm_reconstruction_cache(&self, xet_hash: &str) { @@ -159,12 +274,17 @@ impl XetOps for XetSessions { // ── DownloadStreamWrapper ───────────────────────────────────────────── -struct DownloadStreamWrapper(DownloadStream); +struct DownloadStreamWrapper { + stream: DownloadStream, + /// Declared after `stream` so the stream (and its reconstruction task) is + /// cancelled before the buffer share is handed back. + _buffer: StreamBufferGuard, +} #[async_trait::async_trait] impl DownloadStreamOps for DownloadStreamWrapper { async fn next(&mut self) -> Result> { - Ok(self.0.next().await?) + Ok(self.stream.next().await?) } } @@ -346,3 +466,58 @@ impl StreamingWriterOps for StreamingWriter { self.bytes_written == 0 } } + +#[cfg(test)] +mod stream_buffer_tests { + use super::*; + + const MIB: u64 = 1_048_576; + const BUDGET: u64 = 1024 * MIB; + + fn total(guard: &StreamBufferGuard) -> u64 { + guard.buffer.total_permits() + } + + #[test] + fn lone_stream_gets_the_ceiling() { + let pool = StreamBufferPool::new(BUDGET); + let stream = pool.register(); + assert_eq!(total(&stream), 256 * MIB); + } + + #[test] + fn shares_shrink_to_the_floor_and_grow_back() { + let pool = StreamBufferPool::new(BUDGET); + let mut streams: Vec<_> = (0..4).map(|_| pool.register()).collect(); + // 1 GiB / 4 = 256 MiB: still at the ceiling. + assert!(streams.iter().all(|stream| total(stream) == 256 * MIB)); + + streams.extend((0..4).map(|_| pool.register())); + // 1 GiB / 8 = 128 MiB for everyone, including the early streams. + assert!(streams.iter().all(|stream| total(stream) == 128 * MIB)); + + streams.extend((0..56).map(|_| pool.register())); + // 1 GiB / 64 = 16 MiB, clamped up to the floor. + assert!(streams.iter().all(|stream| total(stream) == 64 * MIB)); + + streams.truncate(2); + // 1 GiB / 2 = 512 MiB, clamped down to the ceiling. + assert!(streams.iter().all(|stream| total(stream) == 256 * MIB)); + assert_eq!(pool.inner.lock().unwrap().active.len(), 2); + } + + #[tokio::test] + async fn shrink_applies_once_held_permits_return() { + let pool = StreamBufferPool::new(BUDGET); + let first = pool.register(); + let held = first.buffer.acquire_many(200 * MIB).await.unwrap(); + + let _others: Vec<_> = (0..7).map(|_| pool.register()).collect(); + // Target is 128 MiB but 200 MiB are out: the total is updated now, + // the part that cannot be reclaimed yet stays pending. + assert_eq!(total(&first), 128 * MIB); + assert!(first.buffer.available_permits() <= 56 * MIB); + drop(held); + assert_eq!(first.buffer.available_permits(), 128 * MIB); + } +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs index aaeb8e05..37b73c28 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -4,7 +4,7 @@ pub mod bench; pub mod fs_tests; use std::path::Path; -use std::process::{Child, Command}; +use std::process::{Child, Command, Stdio}; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -281,6 +281,46 @@ pub async fn upload_file(config: Arc, staged_path: &Path) -> X /// Spawn hf-mount-fuse as a child process, wait until the mountpoint is live. /// `extra_args` are appended to the command (e.g. `&["--read-only"]`). pub fn mount_bucket(bucket_id: &str, mount_point: &str, cache_dir: &str, extra_args: &[&str]) -> Child { + spawn_mount_bucket( + bucket_id, + mount_point, + cache_dir, + extra_args, + Stdio::inherit(), + Stdio::inherit(), + ) +} + +/// Like `mount_bucket`, but with the daemon's output written to `log_path` +/// so a test can assert on what it logged. Tracing goes to stdout, so both +/// stdout and stderr are captured. +pub fn mount_bucket_logged( + bucket_id: &str, + mount_point: &str, + cache_dir: &str, + extra_args: &[&str], + log_path: &str, +) -> Child { + let log = std::fs::File::create(log_path).expect("create daemon log file"); + let log_err = log.try_clone().expect("clone daemon log file"); + spawn_mount_bucket( + bucket_id, + mount_point, + cache_dir, + extra_args, + log.into(), + log_err.into(), + ) +} + +fn spawn_mount_bucket( + bucket_id: &str, + mount_point: &str, + cache_dir: &str, + extra_args: &[&str], + stdout: Stdio, + stderr: Stdio, +) -> Child { let token = std::env::var("HF_TOKEN").unwrap(); let binary = std::env::current_exe() @@ -314,6 +354,8 @@ pub fn mount_bucket(bucket_id: &str, mount_point: &str, cache_dir: &str, extra_a ]) .args(extra_args) .args(["bucket", bucket_id, mount_point]) + .stdout(stdout) + .stderr(stderr) .spawn() .expect("Failed to spawn hf-mount-fuse"); @@ -595,3 +637,75 @@ pub fn generate_pattern(size: usize) -> Vec { pub fn verify_pattern(data: &[u8], offset: usize) -> bool { data.iter().enumerate().all(|(i, &b)| b == ((offset + i) % 251) as u8) } + +/// Deterministic pseudo-random byte generator (xorshift64). Each seeded file +/// has distinct content so nothing dedups across files, and readers can +/// regenerate the expected bytes chunk by chunk instead of holding them. +pub struct PseudoRandomBytes(u64); + +impl PseudoRandomBytes { + pub fn new(seed: u64) -> Self { + Self(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1) + } + + pub fn fill(&mut self, buf: &mut [u8]) { + for chunk in buf.chunks_mut(8) { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + let bytes = x.to_le_bytes(); + chunk.copy_from_slice(&bytes[..chunk.len()]); + } + } +} + +/// Seed a bucket with `count` files of `size` bytes each under `par/`, all +/// uploaded in one session. Returns the relative paths; file `i` regenerates +/// from `PseudoRandomBytes::new(i)`. +pub async fn seed_parallel_read_files( + hub: &Arc, + tmp_dir_tag: &str, + count: usize, + size: usize, +) -> Vec { + let write_config = build_write_config(hub).await; + let tmp_dir = std::env::temp_dir().join(format!("hf-mount-{}-{}", tmp_dir_tag, std::process::id())); + std::fs::create_dir_all(&tmp_dir).ok(); + + let mut buf = vec![0u8; size]; + let staged: Vec<_> = (0..count) + .map(|i| { + let path = tmp_dir.join(format!("f_{i:02}.bin")); + PseudoRandomBytes::new(i as u64).fill(&mut buf); + std::fs::write(&path, &buf).unwrap(); + path + }) + .collect(); + + let upload_session = FileUploadSession::new(write_config) + .await + .expect("FileUploadSession::new failed"); + let infos = upload_session + .upload_files(staged.iter().map(|p| (p.clone(), Sha256Policy::Skip))) + .await + .expect("upload_files failed"); + upload_session.finalize().await.expect("finalize failed"); + assert_eq!(infos.len(), count); + + let rel_paths: Vec = (0..count).map(|i| format!("par/f_{i:02}.bin")).collect(); + let ops: Vec<_> = rel_paths + .iter() + .zip(&infos) + .map(|(rel, info)| hf_mount::hub_api::BatchOp::AddFile { + path: rel.clone(), + xet_hash: info.hash().to_string(), + mtime: 0, + content_type: None, + }) + .collect(); + hub.batch_operations(&ops).await.expect("batch add failed"); + std::fs::remove_dir_all(&tmp_dir).ok(); + rel_paths +} diff --git a/tests/fuse_ops.rs b/tests/fuse_ops.rs index 44e95b84..6a0cfebf 100644 --- a/tests/fuse_ops.rs +++ b/tests/fuse_ops.rs @@ -249,3 +249,106 @@ async fn test_fuse_revalidation() { panic!("FUSE revalidation test failed: {}", e); } } + +/// Many cold sequential readers at once, more readers than FUSE worker +/// threads, files larger than the initial prefetch window so every stream +/// outlives its first read. Regression test for #234: streams must not +/// starve each other for download buffer, and no read may time out or fail. +#[tokio::test] +async fn test_fuse_parallel_cold_reads() { + const FILE_COUNT: usize = 24; + const CHUNK: usize = 1 << 20; + const FILE_SIZE: usize = 64 * CHUNK; + + let guard = match common::setup_bucket("fuse-par-read").await { + Some(g) => g, + None => return, + }; + let rel_paths = common::seed_parallel_read_files(&guard.hub, "fuse-pr", FILE_COUNT, FILE_SIZE).await; + + let mount_point = format!("/tmp/hf-mount-pr-mnt-{}", std::process::id()); + let cache_dir = format!("/tmp/hf-mount-pr-cache-{}", std::process::id()); + let log_path = format!("/tmp/hf-mount-pr-log-{}.txt", std::process::id()); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let child = common::mount_bucket_logged( + &guard.bucket_id, + &mount_point, + &cache_dir, + // Few workers and a short fetch timeout so a wedge shows up as + // EIO within seconds instead of minutes. + &["--read-only", "--max-threads", "4", "--read-fetch-timeout-ms", "5000"], + &log_path, + ); + + let started = std::time::Instant::now(); + let readers: Vec<_> = rel_paths + .iter() + .enumerate() + .map(|(i, rel)| { + let path = format!("{}/{}", mount_point, rel); + std::thread::spawn(move || -> Result<(), String> { + use std::io::Read; + let mut file = std::fs::File::open(&path).map_err(|e| format!("{path}: open: {e}"))?; + let mut expected_gen = common::PseudoRandomBytes::new(i as u64); + let mut got = vec![0u8; CHUNK]; + let mut expected = vec![0u8; CHUNK]; + for chunk in 0..FILE_SIZE / CHUNK { + file.read_exact(&mut got) + .map_err(|e| format!("{path}: read at {}: {e}", chunk * CHUNK))?; + expected_gen.fill(&mut expected); + if got != expected { + return Err(format!("{path}: content mismatch at {}", chunk * CHUNK)); + } + } + if file.read(&mut got).map_err(|e| format!("{path}: read at eof: {e}"))? != 0 { + return Err(format!("{path}: longer than {FILE_SIZE}")); + } + Ok(()) + }) + }) + .collect(); + + let mut errors = Vec::new(); + for reader in readers { + match reader.join() { + Ok(Ok(())) => {} + Ok(Err(e)) => errors.push(e), + Err(_) => errors.push("reader thread panicked".to_string()), + } + } + eprintln!( + "{} readers x {} MiB done in {:?}", + FILE_COUNT, + FILE_SIZE / CHUNK, + started.elapsed() + ); + + common::unmount(&mount_point, child, 30); + + // A single timeout is retried and would otherwise be invisible. + let log_text = std::fs::read_to_string(&log_path).unwrap_or_default(); + let timeouts = log_text.matches("stream read timed out").count(); + if !errors.is_empty() || timeouts > 0 { + let log_tail: Vec<&str> = log_text.lines().rev().take(20).collect(); + return Err(format!( + "{} read errors, {} stream timeouts\nerrors: {:#?}\ndaemon log tail:\n{}", + errors.len(), + timeouts, + errors, + log_tail.into_iter().rev().collect::>().join("\n"), + )); + } + Ok(()) + })); + + std::fs::remove_dir_all(&mount_point).ok(); + std::fs::remove_dir_all(&cache_dir).ok(); + std::fs::remove_file(&log_path).ok(); + + match result { + Ok(Ok(())) => {} + Ok(Err(e)) => panic!("parallel cold read test failed: {}", e), + Err(e) => std::panic::resume_unwind(e), + } +} From 129f76bfc6620a32b720aa02dcae3f7a9ea8fd20 Mon Sep 17 00:00:00 2001 From: Adrien Date: Mon, 28 Sep 2026 17:43:47 +0200 Subject: [PATCH 2/2] perf(xet): raise the adaptive download concurrency cap to 124 With per-stream buffers, many concurrent readers keep more terms in flight than xet-core's default cap of 64 connections allows. 124 is the value of xet-core's own high-performance preset. Not measured in isolation; the #234 bench numbers were taken with this cap. --- src/setup.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/setup.rs b/src/setup.rs index 2a0d4f05..ad628864 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -359,6 +359,9 @@ pub fn init_tracing(daemon: bool) { let upload_cap = XET_UPLOAD_CONCURRENCY_CAP.to_string(); for (k, v) in [ ("HF_XET_CLIENT_AC_INITIAL_DOWNLOAD_CONCURRENCY", "16"), + // Let the adaptive controller scale far enough to saturate fast links + // under many concurrent readers (xet-core's default cap is 64). + ("HF_XET_CLIENT_AC_MAX_DOWNLOAD_CONCURRENCY", "124"), ("HF_XET_CLIENT_AC_MIN_BYTES_REQUIRED_FOR_ADJUSTMENT", "4194304"), ("HF_XET_RECONSTRUCTION_MIN_RECONSTRUCTION_FETCH_SIZE", "8388608"), ("HF_XET_RECONSTRUCTION_MIN_PREFETCH_BUFFER", "8388608"),