From 85ac074c22ed7bc094343f182a1799efae864e92 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 3 Oct 2026 22:49:28 -0700 Subject: [PATCH 1/4] Reduce ulitebox broker startup latency - Resolve overlay lookups without redundant lower-layer probes. - Map the initial-files tar archive instead of reading it into memory. - Copy shared-memory ring and buffer data directly via peer_memory. - Let a runner caller take the response-reader role instead of a dedicated responses thread; notify full-ring waiters only on the full-to-not-full transition. - Run broker requests with a blocking-handoff worker pool that starts workers lazily, only when a request is about to block. - Lease the lowest free shared-buffer slots so pages stay warm. - End notification receipt cleanly when the local side closes its association, removing spurious exit-time errors. - Wait for the in-process runner's control connection instead of sleeping between accept attempts. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b5ed1835-8e32-4a35-aae1-8d8a8a276d09 --- .github/workflows/ci.yml | 1 + Cargo.lock | 1 + dev_tests/src/ratchet.rs | 2 +- litebox/src/broker/shared_buffer.rs | 68 ++- litebox_broker_core/src/fs/backend.rs | 11 + litebox_broker_core/src/fs/in_mem.rs | 41 +- litebox_broker_core/src/fs/overlay.rs | 300 ++++++---- litebox_broker_core/src/fs/tar_ro.rs | 61 +- litebox_broker_core/src/fs/tests.rs | 73 +++ .../src/socket.rs | 28 +- .../src/socket/tests/mod.rs | 1 + litebox_broker_transport/src/channel.rs | 5 +- litebox_broker_transport/src/lib.rs | 5 + litebox_broker_transport/src/peer_memory.rs | 411 +++++++++++++ litebox_broker_transport/src/pending_calls.rs | 171 ++++-- .../src/memfd.rs | 142 ++--- .../src/pending_calls.rs | 122 ++++ .../src/unix_socket/local.rs | 180 +++--- .../src/local.rs | 117 ++-- .../src/shared_memory.rs | 70 +-- litebox_broker_userland/Cargo.toml | 1 + litebox_broker_userland/src/builder.rs | 3 +- litebox_broker_userland/src/lib.rs | 3 + litebox_broker_userland/src/linux.rs | 11 +- litebox_broker_userland/src/main.rs | 58 +- litebox_broker_userland/src/mapped_file.rs | 185 ++++++ .../src/process_launcher.rs | 2 + litebox_broker_userland/src/runner.rs | 18 +- litebox_broker_userland/src/runner/linux.rs | 21 +- litebox_broker_userland/src/runtime.rs | 560 ++++++++++++------ litebox_broker_userland/src/windows.rs | 9 +- 31 files changed, 1928 insertions(+), 753 deletions(-) create mode 100644 litebox_broker_transport/src/peer_memory.rs create mode 100644 litebox_broker_userland/src/mapped_file.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c77dc0a73c..8aeb963b67 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -103,6 +103,7 @@ jobs: AARCH64_CRATES: >- -p litebox -p litebox_broker_core + -p litebox_broker_transport -p litebox_broker_userland -p litebox_common_linux -p litebox_egress_proxy diff --git a/Cargo.lock b/Cargo.lock index cd4061fbd3..4425f36cb6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1586,6 +1586,7 @@ dependencies = [ "litebox_runner_linux_userland", "litebox_runner_windows_userland", "tempfile", + "windows-sys 0.60.2", ] [[package]] diff --git a/dev_tests/src/ratchet.rs b/dev_tests/src/ratchet.rs index e124927068..102412eac4 100644 --- a/dev_tests/src/ratchet.rs +++ b/dev_tests/src/ratchet.rs @@ -70,7 +70,7 @@ fn ratchet_globals() -> Result<()> { ("litebox/", 7), ("litebox_broker_core/", 1), ("litebox_broker_transport_linux_userland/", 1), - ("litebox_broker_userland/", 1), + ("litebox_broker_userland/", 2), ("litebox_platform/", 2), ("litebox_platform_linux_kernel/", 5), ("litebox_platform_linux_userland/", 5), diff --git a/litebox/src/broker/shared_buffer.rs b/litebox/src/broker/shared_buffer.rs index 27e7db23c6..cb37a634e5 100644 --- a/litebox/src/broker/shared_buffer.rs +++ b/litebox/src/broker/shared_buffer.rs @@ -15,13 +15,15 @@ use litebox_platform::sync::RawMutex as _; use crate::sync::{Mutex, RawSyncPrimitivesProvider}; +/// Leases shared-buffer slots, always the lowest free ones, so that a light +/// workload keeps reusing a few slots whose pages stay mapped and cache-warm +/// in both processes. pub(super) struct SlotAllocator { state: Mutex>, } struct AllocatorState { allocated_slots: Vec, - next_slot: usize, failed: bool, waiters: VecDeque>>, } @@ -49,7 +51,6 @@ impl SlotAllocator { Self { state: Mutex::new(AllocatorState { allocated_slots: vec![false; SHARED_BUFFER_SLOT_COUNT as usize], - next_slot: 0, failed: false, waiters: VecDeque::new(), }), @@ -223,40 +224,30 @@ impl Drop for SlotLease<'_, Platform> { impl AllocatorState { fn allocate(&mut self, length: u32, slot_count: usize) -> Option { - if self + let mut slot_indices = [SharedBufferSlotIndex::default(); MAX_SHARED_BUFFER_SEQUENCE_SLOTS]; + let free_slots = self .allocated_slots .iter() - .filter(|allocated| !**allocated) - .count() - < slot_count - { - return None; - } - - let mut slot_indices = [SharedBufferSlotIndex::default(); MAX_SHARED_BUFFER_SEQUENCE_SLOTS]; - let mut next_slot = self.next_slot; - for stored_slot in &mut slot_indices[..slot_count] { - let slot_index = self - .next_free_slot(next_slot) - .expect("validated shared-buffer capacity must contain a free slot"); - self.allocated_slots[slot_index] = true; + .enumerate() + .filter_map(|(slot_index, allocated)| (!*allocated).then_some(slot_index)); + let mut found = 0; + for (stored_slot, slot_index) in slot_indices[..slot_count].iter_mut().zip(free_slots) { *stored_slot = SharedBufferSlotIndex( u32::try_from(slot_index).expect("shared-buffer slot index must fit in u32"), ); - next_slot = (slot_index + 1) % self.allocated_slots.len(); + found += 1; + } + if found < slot_count { + return None; + } + for slot_index in &slot_indices[..slot_count] { + self.allocated_slots[slot_index.0 as usize] = true; } - self.next_slot = next_slot; Some( SharedBufferSequence::new(&slot_indices[..slot_count], length) .expect("allocated shared-buffer sequence must be valid"), ) } - - fn next_free_slot(&self, next_slot: usize) -> Option { - (0..self.allocated_slots.len()) - .map(|offset| (next_slot + offset) % self.allocated_slots.len()) - .find(|slot_index| !self.allocated_slots[*slot_index]) - } } #[cfg(test)] @@ -290,6 +281,32 @@ mod tests { ); } + #[test] + fn leases_reuse_the_lowest_free_slots() { + let allocator = SlotAllocator::::new(); + let first = allocator.acquire(2 * SHARED_BUFFER_SLOT_SIZE).unwrap(); + let second = allocator.acquire(1).unwrap(); + assert_eq!( + first.sequence().slot_indices(), + &[SharedBufferSlotIndex(0), SharedBufferSlotIndex(1)] + ); + assert_eq!( + second.sequence().slot_indices(), + &[SharedBufferSlotIndex(2)] + ); + + drop(first); + let reused = allocator.acquire(3 * SHARED_BUFFER_SLOT_SIZE).unwrap(); + assert_eq!( + reused.sequence().slot_indices(), + &[ + SharedBufferSlotIndex(0), + SharedBufferSlotIndex(1), + SharedBufferSlotIndex(3) + ] + ); + } + #[test] fn oversized_acquisitions_do_not_fail_the_allocator() { let allocator = SlotAllocator::::new(); @@ -305,7 +322,6 @@ mod tests { fn allocator_state_supports_slots_beyond_bitmap_widths() { let mut state = AllocatorState:: { allocated_slots: alloc::vec![true; 65], - next_slot: 64, failed: false, waiters: VecDeque::new(), }; diff --git a/litebox_broker_core/src/fs/backend.rs b/litebox_broker_core/src/fs/backend.rs index 9e4d5917f8..1ca618b30e 100644 --- a/litebox_broker_core/src/fs/backend.rs +++ b/litebox_broker_core/src/fs/backend.rs @@ -93,6 +93,17 @@ pub trait Backend: Send + Sync + Any { /// Read directory entries at `dir`. fn list_dir_at(&self, handle: DirHandle) -> Result, ReadDirError>; + /// Look up the single entry `name` at `dir`, as [`Self::list_dir_at`] would report it. + /// + /// The default implementation scans the full listing; backends that can find one entry more + /// cheaply should override it. + fn lookup_at(&self, dir: &DirHandle, name: &str) -> Result, ReadDirError> { + Ok(self + .list_dir_at(dir.clone())? + .into_iter() + .find(|entry| entry.name == name)) + } + /// Read at `offset` into `buf`, returning the number of bytes read. /// /// Backends do not have an internal notion of offsets; instead the resolver maintains offsets diff --git a/litebox_broker_core/src/fs/in_mem.rs b/litebox_broker_core/src/fs/in_mem.rs index fee5ac14e8..d2e2339a77 100644 --- a/litebox_broker_core/src/fs/in_mem.rs +++ b/litebox_broker_core/src/fs/in_mem.rs @@ -381,20 +381,24 @@ impl super::backend::Backend for InMe .read() .children .iter() - .map(|(name, child)| { - let (file_type, node_info) = match child { - Node::File(file) => (FileType::RegularFile, file.read().node_info), - Node::Dir(dir) => (FileType::Directory, dir.read().node_info), - }; - DirEntry { - name: name.clone(), - file_type, - ino_info: Some(node_info), - } - }) + .map(|(name, child)| child.dir_entry(name)) .collect()) } + fn lookup_at( + &self, + dir: &super::backend::DirHandle, + name: &str, + ) -> Result, ReadDirError> { + Ok(dir + .get_typed::() + .dir + .read() + .children + .get_key_value(name) + .map(|(name, child)| child.dir_entry(name))) + } + fn read( &self, h: &super::backend::FileHandle, @@ -657,6 +661,21 @@ impl Clone for Node { } } +impl Node { + /// The directory entry for this node, named `name` in its parent. + fn dir_entry(&self, name: &str) -> DirEntry { + let (file_type, node_info) = match self { + Node::File(file) => (FileType::RegularFile, file.read().node_info), + Node::Dir(dir) => (FileType::Directory, dir.read().node_info), + }; + DirEntry { + name: name.into(), + file_type, + ino_info: Some(node_info), + } + } +} + type DirNode = Arc>>; struct DirData { perms: Permissions, diff --git a/litebox_broker_core/src/fs/overlay.rs b/litebox_broker_core/src/fs/overlay.rs index 06d73f4d90..80d587b2bf 100644 --- a/litebox_broker_core/src/fs/overlay.rs +++ b/litebox_broker_core/src/fs/overlay.rs @@ -14,12 +14,13 @@ //! changes do not compromise memory safety or the structural integrity of its internal state. use alloc::boxed::Box; +use alloc::collections::BTreeSet; use alloc::string::String; use alloc::sync::Arc; use alloc::vec; use alloc::vec::Vec; -use hashbrown::{HashMap, HashSet}; +use hashbrown::HashMap; use litebox_platform::sync::{Mutex, MutexGuard, RawSyncPrimitivesProvider}; @@ -109,11 +110,15 @@ enum OverlayFileLayer { } /// A logical directory, resolved to the per-layer directories that make it up. +/// +/// Its entries are looked up by name on demand (see [`Overlay::entry`]), so resolving a path never +/// reads whole directories. struct ResolvedDir { upper: Option, /// Per lower backend, in precedence order, its directory at this path if any. lowers: Vec>, - entries: HashMap, + /// Whether the upper directory hides every lower entry. + opaque: bool, } /// An overlay-visible directory entry, plus which layers contribute to it. @@ -166,15 +171,35 @@ impl Overlay { .iter() .map(|lower| lower.owned_dir_at(lower.root(), OFlags::PATH).map(Some)) .collect::, _>>()?; - self.merge(Some(upper), lowers) + self.resolved_dir(Some(upper), lowers) } - /// Resolve the directory `dir_name` within the already-resolved `parent`, along with the - /// [`WalkedComponent`] reported by the layer that owns it. + /// The logical directory made up of the given per-layer directories. + fn resolved_dir( + &self, + upper: Option, + lowers: Vec>, + ) -> Result { + let opaque = match &upper { + Some(upper) => self + .marker_present(upper, OPAQUE_MARKER) + .map_err(|_| OpenError::Io)?, + None => false, + }; + Ok(ResolvedDir { + upper, + lowers, + opaque, + }) + } + + /// Resolve the directory `dir_name` within the already-resolved `parent`, given its `entry` + /// there, along with the [`WalkedComponent`] reported by the layer that owns it. fn resolve_child_dir( &self, parent: &ResolvedDir, dir_name: &str, + entry: &ResolvedEntry, ) -> Result<(ResolvedDir, WalkedComponent), OpenError> { fn walk_into_dir( backend: &dyn Backend, @@ -199,10 +224,6 @@ impl Overlay { Ok((owned, component)) } - let entry = parent - .entries - .get(dir_name) - .ok_or(OpenError::PathError(PathError::MissingComponent))?; if entry.entry.file_type != FileType::Directory { return Err(OpenError::PathError(PathError::ComponentNotADirectory)); } @@ -234,14 +255,17 @@ impl Overlay { let component = owner_component.expect("a merged directory is owned by upper or by a lower directory"); - Ok((self.merge(upper, lowers)?, component)) + Ok((self.resolved_dir(upper, lowers)?, component)) } /// Resolve a logical `path` (relative to the overlay root) to its per-layer directories. fn resolve_dir(&self, path: &[String]) -> Result { let mut current = self.resolve_root()?; for name in path { - current = self.resolve_child_dir(¤t, name)?.0; + let entry = self + .entry(¤t, name)? + .ok_or(OpenError::PathError(PathError::MissingComponent))?; + current = self.resolve_child_dir(¤t, name, &entry)?.0; } Ok(current) } @@ -366,11 +390,7 @@ impl Overlay { } fn marker_present(&self, dir: &DirHandle, marker: &str) -> Result { - Ok(self - .upper - .list_dir_at(dir.clone())? - .iter() - .any(|entry| entry.name == marker)) + Ok(self.upper.lookup_at(dir, marker)?.is_some()) } /// Create `marker` in the upper directory `dir`, if not already there. @@ -575,107 +595,116 @@ impl Overlay { self.state.lock().copied_up.get(node).cloned() } - /// Merge the per-layer directories of one logical directory into its overlay-visible entries. - fn merge( - &self, - upper: Option, - lowers: Vec>, - ) -> Result { - let upper_entries = match &upper { - Some(handle) => self - .upper - .list_dir_at(handle.clone()) - .map_err(|_| OpenError::Io)?, - None => Vec::new(), - }; - // Markers held by this upper directory, which say what it hides from the lowers. - let markers: HashSet = upper_entries - .iter() - .filter(|entry| !valid(&entry.name)) - .map(|entry| entry.name.clone()) - .collect(); - let opaque = markers.contains(OPAQUE_MARKER); - - let mut entries = HashMap::new(); - // Names at which lower entries can no longer be merged in: an entry exists there that is - // not a directory in every layer that contributed to it. - let mut blocked = HashSet::new(); - - for mut entry in upper_entries.into_iter().filter(|entry| valid(&entry.name)) { - if entry.file_type != FileType::Directory { - blocked.insert(entry.name.clone()); - } - entry.ino_info = entry - .ino_info - .take() - .map(|node| self.map_node(&mut self.state.lock().ids, None, node)); - entries.insert( - entry.name.clone(), - ResolvedEntry { + /// The overlay-visible entry `name` of `dir`, if any, merged across the layers that hold it. + fn entry(&self, dir: &ResolvedDir, name: &str) -> Result, OpenError> { + fn lookup( + backend: &dyn Backend, + dir: &DirHandle, + name: &str, + ) -> Result, OpenError> { + backend.lookup_at(dir, name).map_err(|_| OpenError::Io) + } + if !valid(name) { + return Ok(None); + } + + let mut merged: Option = None; + // Whether lower entries can no longer be merged in: the entry is not a directory in every + // layer that contributed to it. + let mut blocked = false; + if let Some(upper) = &dir.upper { + if let Some(mut entry) = lookup(self.upper.as_ref(), upper, name)? { + blocked = entry.file_type != FileType::Directory; + entry.ino_info = entry + .ino_info + .take() + .map(|node| self.map_node(&mut self.state.lock().ids, None, node)); + merged = Some(ResolvedEntry { entry, upper: true, lower: None, lower_directories: vec![false; self.lowers.len()], - }, - ); + }); + } + // Markers held by the upper directory say what it hides from the lowers. + if dir.opaque + || self + .marker_present(upper, &whiteout(name)) + .map_err(|_| OpenError::Io)? + { + return Ok(merged); + } } - if !opaque { - for (layer, handle) in lowers.iter().enumerate() { - let Some(handle) = handle else { - continue; - }; - let layer_entries = self.lowers[layer] - .list_dir_at(handle.clone()) - .map_err(|_| OpenError::Io)?; - for mut lower_entry in layer_entries { - let name = lower_entry.name.clone(); - if !valid(&name) || markers.contains(&whiteout(&name)) { - continue; - } - let directory = lower_entry.file_type == FileType::Directory; - let lower_node = lower_entry.ino_info.take(); - let entry = entries - .entry(name.clone()) - .or_insert_with(|| ResolvedEntry { - entry: lower_entry, - upper: false, - lower: Some(layer), - lower_directories: vec![false; self.lowers.len()], - }); - entry.lower.get_or_insert(layer); - if !entry.upper && entry.lower == Some(layer) { - // This layer owns the entry, so its node is the one callers see. - entry.entry.ino_info = lower_node.map(|node| { - self.map_node(&mut self.state.lock().ids, Some(layer), node) - }); - } - if blocked.contains(&name) { - continue; - } - if directory { - entry.lower_directories[layer] = true; - // Several layers describe one logical directory; the one already resolved - // above owns the identity, and this layer's node adopts it. - if let (Some(node), Some(id)) = (lower_node, entry.entry.ino_info) { - self.state - .lock() - .ids - .entry(layer_node(Some(layer), node)) - .or_insert(id); - } - } else { - blocked.insert(name); - } + for (layer, handle) in dir.lowers.iter().enumerate() { + let Some(handle) = handle else { + continue; + }; + let Some(mut lower_entry) = lookup(self.lowers[layer].as_ref(), handle, name)? else { + continue; + }; + let directory = lower_entry.file_type == FileType::Directory; + let lower_node = lower_entry.ino_info.take(); + let entry = merged.get_or_insert_with(|| ResolvedEntry { + entry: lower_entry, + upper: false, + lower: Some(layer), + lower_directories: vec![false; self.lowers.len()], + }); + entry.lower.get_or_insert(layer); + if !entry.upper && entry.lower == Some(layer) { + // This layer owns the entry, so its node is the one callers see. + entry.entry.ino_info = lower_node + .map(|node| self.map_node(&mut self.state.lock().ids, Some(layer), node)); + } + if blocked { + continue; + } + if directory { + entry.lower_directories[layer] = true; + // Several layers describe one logical directory; the one already resolved above + // owns the identity, and this layer's node adopts it. + if let (Some(node), Some(id)) = (lower_node, entry.entry.ino_info) { + self.state + .lock() + .ids + .entry(layer_node(Some(layer), node)) + .or_insert(id); } + } else { + blocked = true; } } + Ok(merged) + } - Ok(ResolvedDir { - upper, - lowers, - entries, - }) + /// Every overlay-visible entry of `dir`, sorted by name. + fn entries(&self, dir: &ResolvedDir) -> Result, OpenError> { + let mut names = BTreeSet::new(); + let mut add_names = |backend: &dyn Backend, dir: &DirHandle| { + let entries = backend + .list_dir_at(dir.clone()) + .map_err(|_| OpenError::Io)?; + names.extend(entries.into_iter().map(|entry| entry.name)); + Ok::<_, OpenError>(()) + }; + if let Some(upper) = &dir.upper { + add_names(self.upper.as_ref(), upper)?; + } + if !dir.opaque { + for (lower, handle) in self.lowers.iter().zip(&dir.lowers) { + if let Some(handle) = handle { + add_names(lower.as_ref(), handle)?; + } + } + } + let mut entries = Vec::with_capacity(names.len()); + for name in names { + if let Some(entry) = self.entry(dir, &name)? { + entries.push(entry.entry); + } + } + Ok(entries) } } @@ -745,9 +774,9 @@ impl Backend for Overlay { if !valid(name) { return Err(PathError::InvalidPathname.into()); } - let entry = current - .entries - .get(*name) + let entry = self + .entry(¤t, name) + .map_err(open_to_walk_error)? .ok_or(PathError::NoSuchFileOrDirectory)?; if entry.entry.file_type != FileType::Directory { return Ok(WalkOutcome { @@ -757,7 +786,7 @@ impl Backend for Overlay { }); } let (child, component) = self - .resolve_child_dir(¤t, name) + .resolve_child_dir(¤t, name, &entry) .map_err(open_to_walk_error)?; current = child; path.push(String::from(*name)); @@ -804,9 +833,8 @@ impl Backend for Overlay { let path = dir.into_typed::().path; let guard = self.namespace.lock(); let resolved = self.resolve_dir(&path)?; - let entry = resolved - .entries - .get(name) + let entry = self + .entry(&resolved, name)? .ok_or(OpenError::PathError(PathError::NoSuchFileOrDirectory))?; // The resolver only reaches `create_file_at` once a walk reported the name as missing, so // an existing entry means an exclusive create must fail here. @@ -890,13 +918,15 @@ impl Backend for Overlay { fn list_dir_at(&self, handle: DirHandle) -> Result, ReadDirError> { let path = handle.into_typed::().path; let resolved = self.resolve_dir(&path).map_err(|_| ReadDirError::Io)?; - let mut entries: Vec = resolved - .entries - .into_values() - .map(|entry| entry.entry) - .collect(); - entries.sort_by(|left, right| left.name.cmp(&right.name)); - Ok(entries) + self.entries(&resolved).map_err(|_| ReadDirError::Io) + } + + fn lookup_at(&self, dir: &DirHandle, name: &str) -> Result, ReadDirError> { + let resolved = self + .resolve_dir(&dir.get_typed::().path) + .map_err(|_| ReadDirError::Io)?; + let entry = self.entry(&resolved, name).map_err(|_| ReadDirError::Io)?; + Ok(entry.map(|entry| entry.entry)) } fn read(&self, h: &FileHandle, buf: &mut [u8], offset: usize) -> Result { @@ -982,7 +1012,7 @@ impl Backend for Overlay { } let path = dir.into_typed::().path; let locked = self.namespace.lock(); - if self.resolve_dir(&path)?.entries.contains_key(name) { + if self.entry(&self.resolve_dir(&path)?, name)?.is_some() { return Err(OpenError::AlreadyExists); } let upper = self.ensure_upper_dir(&locked, &path)?; @@ -1019,7 +1049,11 @@ impl Backend for Overlay { let mut path = dir.into_typed::().path; let locked = self.namespace.lock(); let resolved = self.resolve_dir(&path).map_err(open_to_mkdir_error)?; - if resolved.entries.contains_key(name) { + if self + .entry(&resolved, name) + .map_err(open_to_mkdir_error)? + .is_some() + { return Err(MkdirError::AlreadyExists); } let upper = self @@ -1068,9 +1102,9 @@ impl Backend for Overlay { let path = dir.into_typed::().path; let locked = self.namespace.lock(); let resolved = self.resolve_dir(&path).map_err(open_to_unlink_error)?; - let entry = resolved - .entries - .get(name) + let entry = self + .entry(&resolved, name) + .map_err(open_to_unlink_error)? .ok_or(PathError::NoSuchFileOrDirectory)?; if entry.entry.file_type == FileType::Directory { return Err(UnlinkError::IsADirectory); @@ -1109,17 +1143,21 @@ impl Backend for Overlay { let path = dir.into_typed::().path; let locked = self.namespace.lock(); let resolved = self.resolve_dir(&path).map_err(open_to_rmdir_error)?; - let entry = resolved - .entries - .get(name) + let entry = self + .entry(&resolved, name) + .map_err(open_to_rmdir_error)? .ok_or(PathError::NoSuchFileOrDirectory)?; if entry.entry.file_type != FileType::Directory { return Err(RmdirError::NotADirectory); } let (child, _) = self - .resolve_child_dir(&resolved, name) + .resolve_child_dir(&resolved, name, &entry) .map_err(open_to_rmdir_error)?; - if !child.entries.is_empty() { + if !self + .entries(&child) + .map_err(open_to_rmdir_error)? + .is_empty() + { return Err(RmdirError::NotEmpty); } diff --git a/litebox_broker_core/src/fs/tar_ro.rs b/litebox_broker_core/src/fs/tar_ro.rs index c4ba8ba38f..bd61ceee1e 100644 --- a/litebox_broker_core/src/fs/tar_ro.rs +++ b/litebox_broker_core/src/fs/tar_ro.rs @@ -24,6 +24,7 @@ //! Taro Milk Tea, Tapioca Bubbles, 50% Sugar, No Ice. //! ``` +use alloc::boxed::Box; use alloc::string::String; use alloc::vec::Vec; use core::ops::Range; @@ -50,15 +51,22 @@ pub struct TarRo { tar_index: TarIndex, } +/// The storage holding a [`TarRo`] archive. +type TarData = Box + Send + Sync>; + impl TarRo { - /// Construct a tar backend using a caller-provided inode allocator. + /// Construct a tar backend over the archive in `tar_data`, using a caller-provided inode + /// allocator. + /// + /// `tar_data` is any storage that can lend out the archive's bytes, such as a static slice, an + /// owned buffer, or a memory-mapped file; it must lend out the same bytes every time. #[must_use] pub fn new( - tar_data: alloc::borrow::Cow<'static, [u8]>, + tar_data: impl AsRef<[u8]> + Send + Sync + 'static, inode_allocator: InodeAllocator, ) -> Self { Self { - tar_index: TarIndex::new(tar_data, inode_allocator), + tar_index: TarIndex::new(Box::new(tar_data), inode_allocator), } } } @@ -181,24 +189,17 @@ impl super::backend::Backend for TarRo { Ok(self.tar_index.dirs[handle.idx] .children .iter() - .map(|(name, child)| { - let (file_type, node_info) = match *child { - IndexedChild::File(idx) => { - (FileType::RegularFile, self.tar_index.files[idx].node_info) - } - IndexedChild::Dir(idx) => { - (FileType::Directory, self.tar_index.dirs[idx].node_info) - } - }; - DirEntry { - name: name.clone(), - file_type, - ino_info: Some(node_info), - } - }) + .map(|(name, child)| self.tar_index.dir_entry(name, *child)) .collect()) } + fn lookup_at(&self, dir: &DirHandle, name: &str) -> Result, ReadDirError> { + Ok(self.tar_index.dirs[dir.get_typed::().idx] + .children + .get_key_value(name) + .map(|(name, child)| self.tar_index.dir_entry(name, *child))) + } + fn read(&self, h: &FileHandle, buf: &mut [u8], offset: usize) -> Result { let file = self.tar_index.file_data(h.get_typed::().idx); let start = offset.min(file.len()); @@ -325,15 +326,16 @@ enum IndexedChild { } struct TarIndex { - tar_data: alloc::borrow::Cow<'static, [u8]>, + tar_data: TarData, files: Vec, dirs: Vec, } impl TarIndex { - fn new(tar_data: alloc::borrow::Cow<'static, [u8]>, inode_allocator: InodeAllocator) -> Self { - let archive = tar_no_std::TarArchiveRef::new(tar_data.as_ref()).expect("invalid tar data"); - let base_ptr = tar_data.as_ptr() as usize; + fn new(tar_data: TarData, inode_allocator: InodeAllocator) -> Self { + let bytes = (*tar_data).as_ref(); + let archive = tar_no_std::TarArchiveRef::new(bytes).expect("invalid tar data"); + let base_ptr = bytes.as_ptr() as usize; let mut files = Vec::new(); let mut files_by_path: HashMap = HashMap::new(); @@ -422,7 +424,20 @@ impl TarIndex { fn file_data(&self, file_idx: usize) -> &[u8] { let range = self.files[file_idx].data_range.clone(); - &self.tar_data[range] + &(*self.tar_data).as_ref()[range] + } + + /// The directory entry for `child`, named `name` in its parent. + fn dir_entry(&self, name: &str, child: IndexedChild) -> DirEntry { + let (file_type, node_info) = match child { + IndexedChild::File(idx) => (FileType::RegularFile, self.files[idx].node_info), + IndexedChild::Dir(idx) => (FileType::Directory, self.dirs[idx].node_info), + }; + DirEntry { + name: name.into(), + file_type, + ino_info: Some(node_info), + } } } diff --git a/litebox_broker_core/src/fs/tests.rs b/litebox_broker_core/src/fs/tests.rs index 5efb791c41..456c2942ad 100644 --- a/litebox_broker_core/src/fs/tests.rs +++ b/litebox_broker_core/src/fs/tests.rs @@ -1795,6 +1795,79 @@ mod overlay { rx.recv_timeout(Duration::from_secs(2)) .expect("copy-up deadlocked"); } + + /// `lookup_at` finds exactly what `list_dir_at` lists, across merged, whited-out, and opaque + /// directories, and never exposes the overlay's own markers. + #[test] + fn lookup_agrees_with_listing() { + use crate::fs::backend::{Backend, CreationMetadata, DirHandle}; + use crate::fs::inode_allocator::InodeAllocator; + use crate::fs::tar_ro::TarRo; + + fn assert_lookups(overlay: &Overlay, dir: &DirHandle, absent: &[&str]) { + for entry in overlay.list_dir_at(dir.clone()).unwrap() { + assert_eq!(overlay.lookup_at(dir, &entry.name).unwrap(), Some(entry)); + } + for name in absent { + assert_eq!(overlay.lookup_at(dir, name).unwrap(), None, "{name}"); + } + } + + let overlay = Overlay::::new( + upper([ + ( + "/bar", + InitialNode::Directory { + mode: ALL_PERMS, + owner: ACTING_USER, + }, + ), + ( + "/bar/qux", + InitialNode::File { + mode: Mode::RWXU, + owner: ACTING_USER, + data: alloc::borrow::Cow::Borrowed(b""), + }, + ), + ]), + TarRo::new(TEST_TAR_FILE, InodeAllocator::standalone()), + InodeAllocator::standalone(), + ); + let root = overlay.owned_dir_at(overlay.root(), OFlags::PATH).unwrap(); + let bar = || { + let walked = overlay.walk_directories(overlay.root(), &["bar"]).unwrap(); + overlay.owned_dir_at(walked.last, OFlags::PATH).unwrap() + }; + let names = |dir: &DirHandle| -> Vec<_> { + let entries = overlay.list_dir_at(dir.clone()).unwrap(); + entries.into_iter().map(|entry| entry.name).collect() + }; + + // `bar` merges the upper and lower directories. + assert_eq!(names(&bar()), ["baz", "qux"]); + assert_lookups(&overlay, &root, &["missing"]); + assert_lookups(&overlay, &bar(), &["missing"]); + + // Unlinking a lower file leaves a whiteout behind, which hides it without being visible. + overlay.unlink_at(root.clone(), "foo").unwrap(); + assert_eq!(names(&root), ["bar"]); + assert_lookups(&overlay, &root, &["foo", ".litebox-overlay-whiteout-foo"]); + + // A directory recreated over a whiteout is opaque, hiding the lower directory's entries. + overlay.unlink_at(bar(), "baz").unwrap(); + overlay.unlink_at(bar(), "qux").unwrap(); + overlay.rmdir_at(root.clone(), "bar").unwrap(); + let metadata = CreationMetadata { + mode: ALL_PERMS, + owner: ACTING_USER, + }; + overlay.mkdir_at(root.clone(), "bar", metadata).unwrap(); + overlay.create_file_at(bar(), "new", metadata).unwrap(); + assert_eq!(names(&bar()), ["new"]); + assert_lookups(&overlay, &bar(), &["baz", ".litebox-overlay-opaque"]); + assert_lookups(&overlay, &root, &["foo"]); + } } mod devices { diff --git a/litebox_broker_platform_linux_userland/src/socket.rs b/litebox_broker_platform_linux_userland/src/socket.rs index e2e97e4779..e91229a657 100644 --- a/litebox_broker_platform_linux_userland/src/socket.rs +++ b/litebox_broker_platform_linux_userland/src/socket.rs @@ -171,8 +171,22 @@ pub struct LinuxSocketProvider { impl LinuxSocketProvider { /// Starts a provider with global and per-process socket limits. pub fn new(max_sockets: usize, max_sockets_per_process: usize) -> IoResult { + Self::with_wait_hook(max_sockets, max_sockets_per_process, || {}) + } + + /// Starts a provider that calls `before_wait` before a broker worker waits + /// for the reactor. + pub fn with_wait_hook( + max_sockets: usize, + max_sockets_per_process: usize, + before_wait: fn(), + ) -> IoResult { Ok(Self { - reactor: Arc::new(ReactorClient::start(max_sockets, max_sockets_per_process)?), + reactor: Arc::new(ReactorClient::start( + max_sockets, + max_sockets_per_process, + before_wait, + )?), }) } } @@ -415,10 +429,15 @@ struct ReactorClient { wake: Arc, next_socket_id: AtomicU64, thread: Mutex>>, + before_wait: fn(), } impl ReactorClient { - fn start(max_sockets: usize, max_sockets_per_process: usize) -> IoResult { + fn start( + max_sockets: usize, + max_sockets_per_process: usize, + before_wait: fn(), + ) -> IoResult { let epoll_fd = epoll::create(epoll::CreateFlags::CLOEXEC)?; let wake = Arc::new(eventfd(0, EventfdFlags::CLOEXEC | EventfdFlags::NONBLOCK)?); epoll::add( @@ -483,6 +502,7 @@ impl ReactorClient { wake, next_socket_id: AtomicU64::new(1), thread: Mutex::new(Some(reactor_thread)), + before_wait, }) } @@ -498,6 +518,7 @@ impl ReactorClient { &self, make_command: impl FnOnce(SyncSender>) -> ReactorCommand, ) -> BrokerResult { + (self.before_wait)(); let (response, receive) = sync_channel(1); let command = make_command(response); // Block until the reactor has queue space rather than surfacing a @@ -522,6 +543,7 @@ impl ReactorClient { address: SocketAddrV4, guest_source_lease: Option, ) -> core::result::Result { + (self.before_wait)(); let (response, receive) = sync_channel(1); let command = ReactorCommand::Connect { id, @@ -547,6 +569,7 @@ impl ReactorClient { } fn close_socket(&self, id: u64) { + (self.before_wait)(); let (response, receive) = sync_channel(1); if self .commands @@ -717,6 +740,7 @@ impl ReactorClient { } fn close_process(&self, process_authority: ProcessId) { + (self.before_wait)(); let (response, receive) = sync_channel(1); if self .commands diff --git a/litebox_broker_platform_linux_userland/src/socket/tests/mod.rs b/litebox_broker_platform_linux_userland/src/socket/tests/mod.rs index 4ce729a0c4..32cd65dbeb 100644 --- a/litebox_broker_platform_linux_userland/src/socket/tests/mod.rs +++ b/litebox_broker_platform_linux_userland/src/socket/tests/mod.rs @@ -143,6 +143,7 @@ fn failed_close_acknowledgement_waits_for_reactor_termination() { wake, next_socket_id: AtomicU64::new(1), thread: Mutex::new(Some(reactor_thread)), + before_wait: || {}, }); let closing_reactor = Arc::clone(&reactor); diff --git a/litebox_broker_transport/src/channel.rs b/litebox_broker_transport/src/channel.rs index b333aa2d11..b344898404 100644 --- a/litebox_broker_transport/src/channel.rs +++ b/litebox_broker_transport/src/channel.rs @@ -111,8 +111,9 @@ pub trait LocalNotificationChannel { /// Receives one broker notification. /// - /// Returns `Ok(None)` when the broker closed the channel cleanly before - /// starting another notification frame. + /// Returns `Ok(None)` when the channel closed cleanly: the broker closed it + /// before starting another notification frame, or the local side closed + /// the association. fn recv_notification(&mut self) -> Result, Self::Error>; } diff --git a/litebox_broker_transport/src/lib.rs b/litebox_broker_transport/src/lib.rs index a103f59a2d..eaecb9a560 100644 --- a/litebox_broker_transport/src/lib.rs +++ b/litebox_broker_transport/src/lib.rs @@ -30,6 +30,11 @@ extern crate std; pub mod channel; pub mod control_ring; +#[cfg(any( + all(target_arch = "x86_64", target_feature = "sse2"), + target_arch = "aarch64" +))] +pub mod peer_memory; pub mod pending_calls; pub mod setup_frame; pub mod shared_memory; diff --git a/litebox_broker_transport/src/peer_memory.rs b/litebox_broker_transport/src/peer_memory.rs new file mode 100644 index 0000000000..d58be9a7a8 --- /dev/null +++ b/litebox_broker_transport/src/peer_memory.rs @@ -0,0 +1,411 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Raw access to memory that a peer can modify concurrently. +//! +//! A peer process may write shared memory at any time, through an alias Rust +//! does not know about and with accesses of any width. Ordinary loads and +//! stores, and even Rust atomics, would race with those writes under Rust's +//! memory model, so these primitives perform each access in inline assembly, +//! which the compiler treats as an opaque hardware operation. No Rust reference +//! into peer-writable memory is ever formed: values are copied out as untrusted +//! snapshots for the caller to validate, and copied in from private buffers. +//! +//! Word operations are indivisible and ordered with respect to a peer that +//! uses matching atomic operations. A peer that does not can only corrupt the +//! values it shares, never the caller's memory safety. + +use core::arch::asm; + +/// Reads a `u32` with acquire semantics. +/// +/// # Safety +/// +/// `address` must be naturally aligned and readable for the whole call. +pub unsafe fn load_u32_acquire(address: *const u32) -> u32 { + let value: u32; + // SAFETY: Guaranteed by the caller. An aligned load is indivisible; on + // x86-64, total store order makes every load an acquire load. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "mov {value:e}, dword ptr [{address}]", + value = out(reg) value, + address = in(reg) address, + options(nostack, preserves_flags), + ); + } + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "ldar {value:w}, [{address}]", + value = out(reg) value, + address = in(reg) address, + options(nostack, preserves_flags), + ); + } + value +} + +/// Indivisibly increments a `u32` with release semantics, wrapping on +/// overflow. +/// +/// # Safety +/// +/// `address` must be naturally aligned and writable for the whole call. +pub unsafe fn increment_u32_release(address: *mut u32) { + // SAFETY: Guaranteed by the caller. A locked read-modify-write is + // indivisible and fully ordered on x86-64. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "lock inc dword ptr [{address}]", + address = in(reg) address, + options(nostack), + ); + } + // SAFETY: Guaranteed by the caller. The exclusive pair retries until the + // increment is indivisible, and the store-release orders earlier accesses + // before it. This avoids requiring the ARMv8.1 atomic instructions. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "2:", + "ldxr {value:w}, [{address}]", + "add {value:w}, {value:w}, #1", + "stlxr {failed:w}, {value:w}, [{address}]", + "cbnz {failed:w}, 2b", + address = in(reg) address, + value = out(reg) _, + failed = out(reg) _, + options(nostack, preserves_flags), + ); + } +} + +/// Reads a `u64` with acquire semantics. +/// +/// # Safety +/// +/// `address` must be naturally aligned and readable for the whole call. +pub unsafe fn load_u64_acquire(address: *const u64) -> u64 { + let value: u64; + // SAFETY: Guaranteed by the caller. An aligned load is indivisible; on + // x86-64, total store order makes every load an acquire load. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "mov {value}, qword ptr [{address}]", + value = out(reg) value, + address = in(reg) address, + options(nostack, preserves_flags), + ); + } + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "ldar {value}, [{address}]", + value = out(reg) value, + address = in(reg) address, + options(nostack, preserves_flags), + ); + } + value +} + +/// Writes a `u64` with release semantics. +/// +/// # Safety +/// +/// `address` must be naturally aligned and writable for the whole call. +pub unsafe fn store_u64_release(address: *mut u64, value: u64) { + // SAFETY: Guaranteed by the caller. `xchg` with memory is indivisible and + // fully ordered on x86-64. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "xchg qword ptr [{address}], {value}", + address = in(reg) address, + value = inout(reg) value => _, + options(nostack, preserves_flags), + ); + } + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "stlr {value}, [{address}]", + address = in(reg) address, + value = in(reg) value, + options(nostack, preserves_flags), + ); + } +} + +/// Copies `destination.len()` bytes from `source` into a private buffer. +/// +/// The bytes are an untrusted snapshot that a concurrent peer write may tear. +/// The copy is unordered: callers order it against word operations with +/// [`core::sync::atomic::fence`]. +/// +/// # Safety +/// +/// `source..source + destination.len()` must be readable for the whole call +/// and must not overlap `destination`. +pub unsafe fn copy_from_peer(source: *const u8, destination: &mut [u8]) { + // SAFETY: Guaranteed by the caller. + unsafe { copy(source, destination.as_mut_ptr(), destination.len()) }; +} + +/// Copies `source` into peer-visible memory at `destination`. +/// +/// The copy is unordered: callers publish it with a later release operation. +/// +/// # Safety +/// +/// `destination..destination + source.len()` must be writable for the whole +/// call and must not overlap `source`. +pub unsafe fn copy_to_peer(source: &[u8], destination: *mut u8) { + // SAFETY: Guaranteed by the caller. + unsafe { copy(source.as_ptr(), destination, source.len()) }; +} + +/// Copies `length` bytes between non-overlapping ranges. +/// +/// Each chunk is loaded completely before any of it is stored. A copy that +/// interleaves loads and stores, such as `rep movsb`, can stall on every access +/// when the destination's page offset sits just past the source's, because +/// the CPU mistakes the pair for a store-to-load dependency. Callers copy +/// between page-aligned shared buffers and arbitrarily aligned private ones, so +/// that case is common. +/// +/// # Safety +/// +/// `source` must be readable and `destination` writable for `length` bytes, +/// and the ranges must not overlap. +unsafe fn copy(source: *const u8, destination: *mut u8, length: usize) { + let mut offset = 0; + while length - offset >= CHUNK_SIZE { + let (from, to) = ( + source.wrapping_add(offset), + destination.wrapping_add(offset), + ); + // SAFETY: Guaranteed by the caller for the remaining range. + unsafe { copy_chunk(from, to) }; + offset += CHUNK_SIZE; + } + while length - offset >= size_of::() { + let (from, to) = ( + source.wrapping_add(offset), + destination.wrapping_add(offset), + ); + // SAFETY: Guaranteed by the caller for the remaining range. + unsafe { copy_word(from, to) }; + offset += size_of::(); + } + while offset < length { + let (from, to) = ( + source.wrapping_add(offset), + destination.wrapping_add(offset), + ); + // SAFETY: Guaranteed by the caller for the remaining range. + unsafe { copy_byte(from, to) }; + offset += 1; + } +} + +const CHUNK_SIZE: usize = 64; + +/// Copies [`CHUNK_SIZE`] bytes. +/// +/// # Safety +/// +/// As for [`copy`] with `length` equal to [`CHUNK_SIZE`]. +unsafe fn copy_chunk(source: *const u8, destination: *mut u8) { + // SAFETY: Guaranteed by the caller. SSE2 is part of the x86-64 baseline + // that this module requires. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "movdqu {a}, xmmword ptr [{source}]", + "movdqu {b}, xmmword ptr [{source} + 16]", + "movdqu {c}, xmmword ptr [{source} + 32]", + "movdqu {d}, xmmword ptr [{source} + 48]", + "movdqu xmmword ptr [{destination}], {a}", + "movdqu xmmword ptr [{destination} + 16], {b}", + "movdqu xmmword ptr [{destination} + 32], {c}", + "movdqu xmmword ptr [{destination} + 48], {d}", + source = in(reg) source, + destination = in(reg) destination, + a = out(xmm_reg) _, + b = out(xmm_reg) _, + c = out(xmm_reg) _, + d = out(xmm_reg) _, + options(nostack, preserves_flags), + ); + } + // SAFETY: Guaranteed by the caller. General-purpose registers keep this + // usable on targets without floating-point registers. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "ldp {a}, {b}, [{source}]", + "ldp {c}, {d}, [{source}, #16]", + "ldp {e}, {f}, [{source}, #32]", + "ldp {g}, {h}, [{source}, #48]", + "stp {a}, {b}, [{destination}]", + "stp {c}, {d}, [{destination}, #16]", + "stp {e}, {f}, [{destination}, #32]", + "stp {g}, {h}, [{destination}, #48]", + source = in(reg) source, + destination = in(reg) destination, + a = out(reg) _, + b = out(reg) _, + c = out(reg) _, + d = out(reg) _, + e = out(reg) _, + f = out(reg) _, + g = out(reg) _, + h = out(reg) _, + options(nostack, preserves_flags), + ); + } +} + +/// Copies eight bytes. +/// +/// # Safety +/// +/// As for [`copy`] with `length` equal to eight. +unsafe fn copy_word(source: *const u8, destination: *mut u8) { + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "mov {word}, qword ptr [{source}]", + "mov qword ptr [{destination}], {word}", + source = in(reg) source, + destination = in(reg) destination, + word = out(reg) _, + options(nostack, preserves_flags), + ); + } + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "ldr {word}, [{source}]", + "str {word}, [{destination}]", + source = in(reg) source, + destination = in(reg) destination, + word = out(reg) _, + options(nostack, preserves_flags), + ); + } +} + +/// Copies one byte. +/// +/// # Safety +/// +/// As for [`copy`] with `length` equal to one. +unsafe fn copy_byte(source: *const u8, destination: *mut u8) { + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "x86_64")] + unsafe { + asm!( + "mov {byte}, byte ptr [{source}]", + "mov byte ptr [{destination}], {byte}", + source = in(reg) source, + destination = in(reg) destination, + byte = out(reg_byte) _, + options(nostack, preserves_flags), + ); + } + // SAFETY: Guaranteed by the caller. + #[cfg(target_arch = "aarch64")] + unsafe { + asm!( + "ldrb {byte:w}, [{source}]", + "strb {byte:w}, [{destination}]", + source = in(reg) source, + destination = in(reg) destination, + byte = out(reg) _, + options(nostack, preserves_flags), + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn word_operations_round_trip() { + let mut words = [0u64; 2]; + let base = words.as_mut_ptr(); + // SAFETY: `words` is live, aligned, and only accessed through `base` + // during these calls. + unsafe { + store_u64_release(base, 0x0123_4567_89ab_cdef); + assert_eq!(load_u64_acquire(base), 0x0123_4567_89ab_cdef); + let epoch = base.add(1).cast::(); + increment_u32_release(epoch); + increment_u32_release(epoch); + assert_eq!(load_u32_acquire(epoch), 2); + epoch.write(u32::MAX); + increment_u32_release(epoch); + assert_eq!(load_u32_acquire(epoch), 0); + } + } + + #[test] + fn concurrent_increments_are_indivisible() { + const THREADS: u32 = 4; + const INCREMENTS: u32 = 10_000; + let mut counter = 0u32; + let address = std::ptr::from_mut(&mut counter).expose_provenance(); + std::thread::scope(|scope| { + for _ in 0..THREADS { + scope.spawn(move || { + for _ in 0..INCREMENTS { + let shared = std::ptr::with_exposed_provenance_mut(address); + // SAFETY: `counter` outlives the scope, is aligned, and + // is only accessed through these increments meanwhile. + unsafe { increment_u32_release(shared) }; + } + }); + } + }); + assert_eq!(counter, INCREMENTS * THREADS); + } + + #[test] + fn copies_cover_every_length_and_alignment() { + let source: alloc::vec::Vec = (0..=255u8).collect(); + for start in 0..8 { + for length in 0..(source.len() - start) { + let mut shared = [0xffu8; 256]; + // SAFETY: Both ranges are live, in bounds, and disjoint. + unsafe { + copy_to_peer( + &source[start..start + length], + shared.as_mut_ptr().add(start), + ); + }; + assert_eq!( + &shared[start..start + length], + &source[start..start + length] + ); + assert!(shared[..start].iter().all(|&byte| byte == 0xff)); + assert!(shared[start + length..].iter().all(|&byte| byte == 0xff)); + let mut private = alloc::vec![0u8; length]; + // SAFETY: Both ranges are live, in bounds, and disjoint. + unsafe { copy_from_peer(shared.as_ptr().add(start), &mut private) }; + assert_eq!(private, &source[start..start + length]); + } + } + } +} diff --git a/litebox_broker_transport/src/pending_calls.rs b/litebox_broker_transport/src/pending_calls.rs index c831404a9f..f54be5ba5a 100644 --- a/litebox_broker_transport/src/pending_calls.rs +++ b/litebox_broker_transport/src/pending_calls.rs @@ -76,49 +76,107 @@ pub enum PendingCallsError { } /// Concurrent registry of requests awaiting broker responses. +/// +/// Callers waiting for responses also read them: while any call is pending, +/// one waiting caller holds the reader role and completes responses for every +/// caller until its own arrives, then hands the role to another pending call. +/// A response therefore reaches a lone caller without a thread handoff. pub struct PendingCalls { state: Sync::Mutex>, capacity_available: Sync::Condvar>, } struct PendingCallsInner { - calls: BTreeMap>>, + calls: BTreeMap>>, failure: Option>, + /// The call whose caller holds the reader role, if any call is pending. + reader: Option>>, } -/// Completion state for one request awaiting a broker response. -pub struct PendingCall { - result: Sync::Mutex>>>, - result_ready: Sync::Condvar>>>, +/// Completion state shared between a pending call and its registry. +struct CallSlot { + state: Sync::Mutex>, + changed: Sync::Condvar>, } -impl PendingCall { - fn new() -> Self { +struct CallState { + result: Option>>, + /// Whether this call's caller holds the reader role. + reads_responses: bool, +} + +impl CallSlot { + fn new(reads_responses: bool) -> Self { Self { - result: Sync::mutex(None), - result_ready: Sync::condvar(), + state: Sync::mutex(CallState { + result: None, + reads_responses, + }), + changed: Sync::condvar(), + } + } + + fn resolve(&self, result: Result>, notify: bool) { + let mut state = self.state.lock(); + assert!( + state.result.is_none(), + "broker pending call already resolved" + ); + state.result = Some(result); + if notify { + self.changed.notify_one(); } } - fn resolve(&self, result: Result>) { - let mut stored = self.result.lock(); - assert!(stored.is_none(), "broker pending call already resolved"); - *stored = Some(result); - self.result_ready.notify_one(); + fn hand_reader_role(&self) { + self.state.lock().reads_responses = true; + self.changed.notify_one(); } +} + +/// One registered request awaiting its broker response. +/// +/// Dropping the call withdraws it from its registry and hands off the reader +/// role if it holds it, so an abandoned call never strands other callers. +pub struct PendingCall<'calls, Sync: PendingCallsSync, Error> { + pending_calls: &'calls PendingCalls, + request_id: RequestId, + slot: Arc>, +} +impl PendingCall<'_, Sync, Error> { /// Blocks until the broker responds or the association fails. - pub fn wait(&self) -> Result> { - let mut result = self.result.lock(); + /// + /// While this caller holds the reader role, it calls `read_responses` + /// repeatedly. Each call must block until it reads at least one response + /// and passes it to [`PendingCalls::complete`], or until the association + /// fails and the failure is recorded with [`PendingCalls::record_failure`]. + /// Only the caller holding the reader role calls `read_responses`. + pub fn wait(self, mut read_responses: impl FnMut()) -> Result> { loop { - if let Some(result) = result.take() { - return result; + { + let mut state = self.slot.state.lock(); + loop { + if let Some(result) = state.result.take() { + return result; + } + if state.reads_responses { + break; + } + state = self.slot.changed.wait(state); + } } - result = self.result_ready.wait(result); + read_responses(); } } } +impl Drop for PendingCall<'_, Sync, Error> { + fn drop(&mut self) { + self.pending_calls.withdraw(self.request_id, &self.slot); + } +} + impl PendingCalls { /// Creates an empty live pending-call registry. pub fn new() -> Self { @@ -126,55 +184,92 @@ impl PendingCalls { state: Sync::mutex(PendingCallsInner { calls: BTreeMap::new(), failure: None, + reader: None, }), capacity_available: Sync::condvar(), } } /// Registers a request, blocking while the pending-call limit is full. + /// + /// The first call registered while no call is pending takes the reader + /// role. pub fn register( &self, request_id: RequestId, - ) -> Result>, PendingCallsError> { - let pending_call = Arc::new(PendingCall::new()); + ) -> Result, PendingCallsError> { let mut state = self.state.lock(); while state.calls.len() >= MAX_PENDING_CALLS && state.failure.is_none() { state = self.capacity_available.wait(state); } + let state = &mut *state; if let Some(error) = state.failure.as_ref() { return Err(PendingCallsError::AssociationFailed(Arc::clone(error))); } - match state.calls.entry(request_id) { - Entry::Vacant(entry) => { - entry.insert(Arc::clone(&pending_call)); - } - Entry::Occupied(_) => return Err(PendingCallsError::DuplicateRequestId), + let Entry::Vacant(entry) = state.calls.entry(request_id) else { + return Err(PendingCallsError::DuplicateRequestId); + }; + let slot = Arc::new(CallSlot::new(state.reader.is_none())); + entry.insert(Arc::clone(&slot)); + if state.reader.is_none() { + state.reader = Some(Arc::clone(&slot)); } - Ok(pending_call) + Ok(PendingCall { + pending_calls: self, + request_id, + slot, + }) } /// Completes the pending call identified by `response`. - /// - /// # Panics - /// - /// Panics if the internal ordinary-call count is inconsistent with the - /// registered requests. pub fn complete(&self, response: BrokerResponse) -> Result<(), PendingCallsError> { - let pending_call = { + let (slot, notify) = { let mut state = self.state.lock(); if let Some(error) = state.failure.as_ref() { return Err(PendingCallsError::AssociationFailed(Arc::clone(error))); } - let Some(pending_call) = state.calls.remove(&response.request_id) else { + let Some(slot) = state.calls.remove(&response.request_id) else { return Err(PendingCallsError::UnknownResponseId); }; - self.capacity_available.notify_all(); - pending_call + // Registrations wait only while the registry is full. + if state.calls.len() + 1 == MAX_PENDING_CALLS { + self.capacity_available.notify_all(); + } + // The reader completes its own call while reading, not waiting. + let notify = !state + .reader + .as_ref() + .is_some_and(|reader| Arc::ptr_eq(reader, &slot)); + (slot, notify) }; - pending_call.resolve(Ok(response)); + slot.resolve(Ok(response), notify); Ok(()) } + fn withdraw(&self, request_id: RequestId, slot: &Arc>) { + let mut state = self.state.lock(); + if state + .calls + .get(&request_id) + .is_some_and(|registered| Arc::ptr_eq(registered, slot)) + { + state.calls.remove(&request_id); + if state.calls.len() + 1 == MAX_PENDING_CALLS { + self.capacity_available.notify_all(); + } + } + let holds_reader_role = state + .reader + .as_ref() + .is_some_and(|reader| Arc::ptr_eq(reader, slot)); + if holds_reader_role && state.failure.is_none() { + state.reader = state.calls.values().next().map(|next| { + next.hand_reader_role(); + Arc::clone(next) + }); + } + } + /// Records the first terminal failure and resolves every pending call. /// /// Returns whether this call recorded the first failure. @@ -190,7 +285,7 @@ impl PendingCalls { pending_calls }; for pending_call in pending_calls.into_values() { - pending_call.resolve(Err(Arc::clone(&error))); + pending_call.resolve(Err(Arc::clone(&error)), true); } true } diff --git a/litebox_broker_transport_linux_userland/src/memfd.rs b/litebox_broker_transport_linux_userland/src/memfd.rs index 881904ca0e..d2f5726afb 100644 --- a/litebox_broker_transport_linux_userland/src/memfd.rs +++ b/litebox_broker_transport_linux_userland/src/memfd.rs @@ -8,10 +8,10 @@ //! `litebox_broker_transport`, and its futex support lets portable control-ring //! endpoints block and wake without knowing anything about Linux. //! -//! Rust never dereferences the peer-writable mapping. Byte and word access uses -//! positional descriptor I/O into private buffers; the mapping exists only to -//! provide checked addresses to the kernel's futex operations and to process -//! image writes. +//! Rust never forms references into the peer-writable mapping. Byte and word +//! access goes through [`peer_memory`], which copies untrusted snapshots into +//! private buffers, and futex operations and process image writes hand checked +//! addresses to the kernel. use std::io::{Error, Result as IoResult}; use std::io::{ErrorKind, IoSlice, IoSliceMut}; @@ -25,7 +25,7 @@ use std::time::Instant; use rustix::fs::{ MemfdFlags, SealFlags, fcntl_add_seals, fcntl_get_seals, fstat, ftruncate, memfd_create, }; -use rustix::io::{Errno, pread, pwrite}; +use rustix::io::{Errno, pwrite}; use rustix::mm::{MapFlags, ProtFlags, mmap, munmap}; use rustix::net::{ RecvAncillaryBuffer, RecvAncillaryMessage, RecvFlags, ReturnFlags, SendAncillaryBuffer, @@ -39,6 +39,7 @@ use litebox_broker_transport::control_ring::{ use litebox_broker_transport::control_ring::{ memory_permits_byte_range, memory_permits_u32, memory_permits_u64, }; +use litebox_broker_transport::peer_memory; use litebox_broker_transport::shared_memory::{ControlRingMemory, SharedMemory, SharedMemoryError}; use crate::unix_io::{ @@ -63,18 +64,25 @@ struct MappedRegion { } // SAFETY: Moving or sharing this owner does not move or invalidate its OS -// mapping. Its metadata is immutable, and mapped contents are never -// dereferenced by Rust. +// mapping. Its metadata is immutable, and mapped contents are only accessed +// through `peer_memory` or the kernel, never through Rust references. unsafe impl Send for MappedRegion {} -// SAFETY: See the `Send` justification. Only checked raw futex addresses are -// derived from the mapping and passed to the kernel. +// SAFETY: See the `Send` justification. Every access derives a checked raw +// address from the mapping. unsafe impl Sync for MappedRegion {} -fn validate_u64_offset(memory: &MemfdSharedMemory, offset: usize) -> Result<(), SharedMemoryError> { +fn checked_u64_address( + memory: &MemfdSharedMemory, + offset: usize, +) -> Result<*mut u64, SharedMemoryError> { if !memory.policy.permits_u64(offset) { return Err(SharedMemoryError::InvalidRange); } - checked_range(&memory.mapping, offset, size_of::(), align_of::()) + let byte_address = + shared_address(&memory.mapping, offset, size_of::(), align_of::())?; + // The runtime check above establishes the required alignment. + #[allow(clippy::cast_ptr_alignment)] + Ok(byte_address.cast::()) } fn checked_u32_address( @@ -133,46 +141,6 @@ fn validate_nonoverlapping_word_ranges( Ok(()) } -fn read_exact_at( - memory: &MemfdSharedMemory, - offset: usize, - destination: &mut [u8], -) -> Result<(), SharedMemoryError> { - let mut completed = 0; - while completed < destination.len() { - let file_offset = - u64::try_from(offset + completed).map_err(|_| SharedMemoryError::InvalidRange)?; - match pread(&memory.fd, &mut destination[completed..], file_offset) { - Ok(0) => return Err(SharedMemoryError::AccessFailed), - Ok(read) => completed += read, - Err(Errno::INTR) => {} - Err(_) => return Err(SharedMemoryError::AccessFailed), - } - } - Ok(()) -} - -fn write_all_at( - memory: &MemfdSharedMemory, - offset: usize, - source: &[u8], -) -> Result<(), SharedMemoryError> { - let mut completed = 0; - while completed < source.len() { - let file_offset = - u64::try_from(offset + completed).map_err(|_| SharedMemoryError::InvalidRange)?; - match pwrite(&memory.fd, &source[completed..], file_offset) { - Ok(0) => return Err(SharedMemoryError::AccessFailed), - Ok(written) => completed += written, - Err(Errno::INTR) => {} - Err(_) => return Err(SharedMemoryError::AccessFailed), - } - } - Ok(()) -} - -const FUTEX_INCREMENT_OPERATION: libc::c_int = - (libc::FUTEX_OP_ADD << 28) | (libc::FUTEX_OP_CMP_EQ << 24) | (1 << 12); const FUTEX_WAIT_RECHECK_TIMEOUT: libc::timespec = libc::timespec { tv_sec: 0, tv_nsec: 100_000_000, @@ -180,28 +148,6 @@ const FUTEX_WAIT_RECHECK_TIMEOUT: libc::timespec = libc::timespec { // The rustix futex API requires `AtomicU32` references. Raw syscalls keep Rust // references out of memory that a peer can modify through an uncontrolled fd. -fn futex_increment(address: *mut u32) -> IoResult<()> { - // SAFETY: `address` is aligned and lies within the live shared mapping. - // FUTEX_WAKE_OP atomically increments it in the kernel, so Rust never forms - // an atomic reference that a peer could invalidate through another alias. - let result = unsafe { - libc::syscall( - libc::SYS_futex, - address, - libc::FUTEX_WAKE_OP, - 0, - 0, - address, - FUTEX_INCREMENT_OPERATION, - ) - }; - if result == -1 { - Err(Error::last_os_error()) - } else { - Ok(()) - } -} - fn futex_wait(address: *mut u32, expected: u32) -> IoResult<()> { let timeout = FUTEX_WAIT_RECHECK_TIMEOUT; // SAFETY: `address` is aligned and lies within the live shared mapping. The @@ -391,45 +337,53 @@ impl SharedMemory for MemfdSharedMemory { } fn read(&self, offset: usize, destination: &mut [u8]) -> Result<(), SharedMemoryError> { - checked_range(&self.mapping, offset, destination.len(), 1)?; + let source = shared_address(&self.mapping, offset, destination.len(), 1)?; if !self.policy.permits_byte_range(offset, destination.len()) { return Err(SharedMemoryError::InvalidRange); } - read_exact_at(self, offset, destination) + // SAFETY: The range lies within the live mapping, which cannot overlap + // the private destination. + unsafe { peer_memory::copy_from_peer(source, destination) }; + Ok(()) } fn write(&self, offset: usize, source: &[u8]) -> Result<(), SharedMemoryError> { - checked_range(&self.mapping, offset, source.len(), 1)?; + let destination = shared_address(&self.mapping, offset, source.len(), 1)?; if !self.policy.permits_byte_range(offset, source.len()) { return Err(SharedMemoryError::InvalidRange); } - write_all_at(self, offset, source) + // SAFETY: The range lies within the live mapping, which cannot overlap + // the private source. + unsafe { peer_memory::copy_to_peer(source, destination) }; + Ok(()) } } impl ControlRingMemory for MemfdSharedMemory { fn load_u32_acquire(&self, offset: usize) -> Result { - checked_u32_address(self, offset)?; - let mut bytes = [0; size_of::()]; - read_exact_at(self, offset, &mut bytes)?; - Ok(u32::from_ne_bytes(bytes)) + let address = checked_u32_address(self, offset)?; + // SAFETY: The address is aligned and lies within the live mapping. + Ok(unsafe { peer_memory::load_u32_acquire(address) }) } fn increment_u32_release(&self, offset: usize) -> Result<(), SharedMemoryError> { let address = checked_u32_address(self, offset)?; - futex_increment(address).map_err(|_| SharedMemoryError::AccessFailed) + // SAFETY: The address is aligned and lies within the live mapping. + unsafe { peer_memory::increment_u32_release(address) }; + Ok(()) } fn load_u64_acquire(&self, offset: usize) -> Result { - validate_u64_offset(self, offset)?; - let mut bytes = [0; size_of::()]; - read_exact_at(self, offset, &mut bytes)?; - Ok(u64::from_ne_bytes(bytes)) + let address = checked_u64_address(self, offset)?; + // SAFETY: The address is aligned and lies within the live mapping. + Ok(unsafe { peer_memory::load_u64_acquire(address) }) } fn store_u64_release(&self, offset: usize, value: u64) -> Result<(), SharedMemoryError> { - validate_u64_offset(self, offset)?; - write_all_at(self, offset, &value.to_ne_bytes()) + let address = checked_u64_address(self, offset)?; + // SAFETY: The address is aligned and lies within the live mapping. + unsafe { peer_memory::store_u64_release(address, value) }; + Ok(()) } fn store_u64_and_increment_u32_release( @@ -439,10 +393,14 @@ impl ControlRingMemory for MemfdSharedMemory { increment_offset: usize, ) -> Result<(), SharedMemoryError> { validate_nonoverlapping_word_ranges(store_offset, increment_offset)?; - validate_u64_offset(self, store_offset)?; + let store_address = checked_u64_address(self, store_offset)?; let increment_address = checked_u32_address(self, increment_offset)?; - write_all_at(self, store_offset, &value.to_ne_bytes())?; - futex_increment(increment_address).map_err(|_| SharedMemoryError::AccessFailed) + // SAFETY: Both addresses are aligned and lie within the live mapping. + unsafe { + peer_memory::store_u64_release(store_address, value); + peer_memory::increment_u32_release(increment_address); + } + Ok(()) } } @@ -1231,7 +1189,7 @@ mod tests { assert_eq!(fstat(&image.fd).unwrap().st_size, 0x1006); let mut bytes = [0; 6]; - assert_eq!(pread(&image.fd, &mut bytes, 0x1000).unwrap(), 6); + assert_eq!(rustix::io::pread(&image.fd, &mut bytes, 0x1000).unwrap(), 6); assert_eq!(&bytes, b"shared"); } diff --git a/litebox_broker_transport_linux_userland/src/pending_calls.rs b/litebox_broker_transport_linux_userland/src/pending_calls.rs index c430dc5820..e9e8665017 100644 --- a/litebox_broker_transport_linux_userland/src/pending_calls.rs +++ b/litebox_broker_transport_linux_userland/src/pending_calls.rs @@ -74,3 +74,125 @@ pub(crate) fn pending_calls_error(error: PendingCallsError) -> Error { } } } + +#[cfg(test)] +mod tests { + use std::io::Error; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::mpsc; + use std::time::Duration; + + use litebox_broker_protocol::RequestId; + use litebox_broker_protocol::message::{BrokerResponse, BrokerResult}; + use litebox_broker_transport::pending_calls::{MAX_PENDING_CALLS, PendingCallsError}; + + use super::PendingCalls; + + fn response(id: u64) -> BrokerResponse { + BrokerResponse { + request_id: RequestId(id), + result: BrokerResult::ObjectClosed, + } + } + + #[test] + fn a_full_registry_wakes_registration_when_a_call_completes() { + let pending_calls = PendingCalls::new(); + let mut calls: Vec<_> = (0..MAX_PENDING_CALLS as u64) + .map(|id| pending_calls.register(RequestId(id)).unwrap()) + .collect(); + let (registered, registration) = mpsc::channel(); + std::thread::scope(|scope| { + scope.spawn(|| { + let call = pending_calls.register(RequestId(MAX_PENDING_CALLS as u64)); + registered.send(call.is_ok()).unwrap(); + }); + assert!( + registration + .recv_timeout(Duration::from_millis(50)) + .is_err() + ); + pending_calls.complete(response(0)).unwrap(); + assert!(registration.recv_timeout(Duration::from_secs(10)).unwrap()); + }); + assert!(calls.remove(0).wait(|| unreachable!()).is_ok()); + } + + #[test] + fn the_first_caller_reads_responses_for_later_callers() { + let pending_calls = PendingCalls::new(); + let reader = pending_calls.register(RequestId(1)).unwrap(); + let follower = pending_calls.register(RequestId(2)).unwrap(); + std::thread::scope(|scope| { + let follower = scope.spawn(move || follower.wait(|| panic!("follower read"))); + let mut responses = [2, 1].into_iter(); + let result = reader.wait(|| { + pending_calls + .complete(response(responses.next().unwrap())) + .unwrap(); + }); + assert_eq!(result.unwrap().request_id, RequestId(1)); + assert_eq!(follower.join().unwrap().unwrap().request_id, RequestId(2)); + }); + } + + #[test] + fn a_finished_reader_hands_its_role_to_a_pending_call() { + let pending_calls = PendingCalls::new(); + let reader = pending_calls.register(RequestId(1)).unwrap(); + let follower = pending_calls.register(RequestId(2)).unwrap(); + let follower_read = AtomicBool::new(false); + std::thread::scope(|scope| { + let follower = scope.spawn(|| { + follower.wait(|| { + follower_read.store(true, Ordering::Relaxed); + pending_calls.complete(response(2)).unwrap(); + }) + }); + let result = reader.wait(|| pending_calls.complete(response(1)).unwrap()); + assert_eq!(result.unwrap().request_id, RequestId(1)); + assert_eq!(follower.join().unwrap().unwrap().request_id, RequestId(2)); + }); + assert!(follower_read.load(Ordering::Relaxed)); + } + + #[test] + fn an_abandoned_reader_withdraws_and_hands_off_its_role() { + let pending_calls = PendingCalls::new(); + let reader = pending_calls.register(RequestId(1)).unwrap(); + let follower = pending_calls.register(RequestId(2)).unwrap(); + drop(reader); + let result = follower.wait(|| pending_calls.complete(response(2)).unwrap()); + assert_eq!(result.unwrap().request_id, RequestId(2)); + assert!(matches!( + pending_calls.complete(response(1)), + Err(PendingCallsError::UnknownResponseId) + )); + } + + #[test] + fn a_call_registered_while_none_is_pending_reads_its_own_response() { + let pending_calls = PendingCalls::new(); + for id in 1..=2 { + let call = pending_calls.register(RequestId(id)).unwrap(); + let result = call.wait(|| pending_calls.complete(response(id)).unwrap()); + assert_eq!(result.unwrap().request_id, RequestId(id)); + } + } + + #[test] + fn a_reader_failure_resolves_every_pending_call() { + let pending_calls = PendingCalls::new(); + let reader = pending_calls.register(RequestId(1)).unwrap(); + let follower = pending_calls.register(RequestId(2)).unwrap(); + std::thread::scope(|scope| { + let follower = scope.spawn(move || follower.wait(|| panic!("follower read"))); + let result = reader.wait(|| { + pending_calls.record_failure(Arc::new(Error::other("failed"))); + }); + assert!(result.is_err()); + assert!(follower.join().unwrap().is_err()); + }); + } +} diff --git a/litebox_broker_transport_linux_userland/src/unix_socket/local.rs b/litebox_broker_transport_linux_userland/src/unix_socket/local.rs index 78caf1e076..486fa544c6 100644 --- a/litebox_broker_transport_linux_userland/src/unix_socket/local.rs +++ b/litebox_broker_transport_linux_userland/src/unix_socket/local.rs @@ -11,6 +11,7 @@ use std::io::{Error, ErrorKind, Read, Result as IoResult}; use std::os::fd::{AsFd, BorrowedFd, OwnedFd}; use std::os::unix::net::UnixStream; use std::path::Path; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::thread; use std::time::{Duration, Instant}; @@ -66,12 +67,17 @@ pub struct UnixControlRingLocalShutdown { } /// State shared by every activated local endpoint of one association: the -/// request producer, the setup socket used for liveness and teardown, pending -/// call tracking, and the wake handles of all three ring directions. +/// request producer, the response consumer read by whichever caller holds the +/// pending-call reader role, the setup socket used for liveness and teardown, +/// pending call tracking, and the wake handles of all three ring directions. struct LocalRingAssociation { request_producer: Mutex>, + response_consumer: Mutex>, control_stream: UnixStream, pending_calls: Arc, + /// Whether the call channel was dropped, which ends the association as a + /// local close rather than a failure. + closed: AtomicBool, on_failure: Arc, request_wake: ControlRingWakeHandle, response_wake: ControlRingWakeHandle, @@ -139,8 +145,8 @@ impl UnixStreamLocalSetupChannel { } /// Consumes a negotiated setup channel into independently usable active - /// call, notification, and shutdown handles, starting the response - /// dispatcher and liveness monitor. + /// call, notification, and shutdown handles, starting the liveness + /// monitor. /// /// The ring must be the validated control-ring memfd received during this /// setup exchange. @@ -186,20 +192,12 @@ impl UnixStreamLocalSetupChannel { request_producer: Mutex::new(request_producer), control_stream: shutdown_stream, pending_calls: Arc::clone(&pending_calls), + closed: AtomicBool::new(false), on_failure, response_wake: response_consumer.wake_handle(), + response_consumer: Mutex::new(response_consumer), notification_wake: notification_consumer.wake_handle(), }); - let response_association = Arc::clone(&association); - if let Err(error) = thread::Builder::new() - .name("litebox-broker-responses".to_owned()) - .spawn(move || { - dispatch_responses(response_consumer, response_association); - }) - { - let _ = association.fail(error); - return Err(Error::other("failed to start broker response dispatcher")); - } let monitor_association = Arc::clone(&association); if let Err(error) = thread::Builder::new() .name("litebox-broker-liveness".to_owned()) @@ -296,6 +294,7 @@ impl AsFd for UnixControlRingLocalShutdown { impl Drop for UnixControlRingLocalCallChannel { fn drop(&mut self) { + self.association.closed.store(true, Ordering::Release); let _ = self.association.fail(Error::new( ErrorKind::ConnectionAborted, "broker local call channel dropped", @@ -366,7 +365,9 @@ impl LocalCallChannel for UnixControlRingLocalCallChannel { let _ = association.fail(error); } - pending_call.wait().map_err(|error| copy_io_error(&error)) + pending_call + .wait(|| association.read_response()) + .map_err(|error| copy_io_error(&error)) } } @@ -376,7 +377,7 @@ impl LocalNotificationChannel for UnixControlRingLocalNotificationChannel { fn recv_notification(&mut self) -> IoResult> { loop { if let Some(error) = self.association.pending_calls.current_failure() { - return Err(copy_io_error(&error)); + return self.association.end_notifications(&error); } match self.consumer.try_read(decode_notification) { Ok(ControlRingReadStatus::Message(notification)) => { @@ -386,7 +387,7 @@ impl LocalNotificationChannel for UnixControlRingLocalNotificationChannel { } Ok(ControlRingReadStatus::Empty { wait_epoch }) => { if let Some(error) = self.association.pending_calls.current_failure() { - return Err(copy_io_error(&error)); + return self.association.end_notifications(&error); } if let Err(error) = self.consumer.wait_for_message(wait_epoch) { let result = Err(copy_io_error(&error)); @@ -412,6 +413,16 @@ impl LocalNotificationChannel for UnixControlRingLocalNotificationChannel { } impl LocalRingAssociation { + /// Ends notification receipt once the association has ended with + /// `failure`: cleanly if the local side closed it. + fn end_notifications(&self, failure: &Error) -> IoResult> { + if self.closed.load(Ordering::Acquire) { + Ok(None) + } else { + Err(copy_io_error(failure)) + } + } + fn acknowledge_notification( &self, consumer: &mut ControlRingConsumer, @@ -433,6 +444,51 @@ impl LocalRingAssociation { Ok(()) } + /// Reads one response for the caller holding the pending-call reader + /// role, blocking until one arrives or the association fails. + fn read_response(&self) { + let mut consumer = self + .response_consumer + .lock() + .expect("broker response reader mutex poisoned"); + loop { + match consumer.try_read(decode_response) { + Ok(ControlRingReadStatus::Message(response)) => { + if let Err(error) = consumer + .publish_head() + .map_err(ring_error) + .and_then(|()| consumer.wake_producer()) + .and_then(|()| { + self.pending_calls + .complete(response) + .map_err(pending_calls_error) + }) + { + let _ = self.fail(error); + } + return; + } + Ok(ControlRingReadStatus::Empty { wait_epoch }) => { + if self.pending_calls.current_failure().is_some() { + return; + } + if let Err(error) = consumer.wait_for_message(wait_epoch) { + let _ = self.fail(error); + return; + } + } + Err(ControlRingReadError::Ring(error)) => { + let _ = self.fail(ring_error(error)); + return; + } + Err(ControlRingReadError::Decode(error)) => { + let _ = self.fail(wire_error(error)); + return; + } + } + } + } + fn fail(&self, error: Error) -> IoResult<()> { let first_failure = self.pending_calls.record_failure(Arc::new(error)); let request_wake = self.request_wake.interrupt_wait(); @@ -473,49 +529,6 @@ fn wait_for_socket_termination(stream: &mut UnixStream, peer: &'static str) -> E } } -fn dispatch_responses( - mut consumer: ControlRingConsumer, - association: Arc, -) { - loop { - match consumer.try_read(decode_response) { - Ok(ControlRingReadStatus::Message(response)) => { - if let Err(error) = consumer - .publish_head() - .map_err(ring_error) - .and_then(|()| consumer.wake_producer()) - .and_then(|()| { - association - .pending_calls - .complete(response) - .map_err(pending_calls_error) - }) - { - let _ = association.fail(error); - return; - } - } - Ok(ControlRingReadStatus::Empty { wait_epoch }) => { - if association.pending_calls.current_failure().is_some() { - return; - } - if let Err(error) = consumer.wait_for_message(wait_epoch) { - let _ = association.fail(error); - return; - } - } - Err(ControlRingReadError::Ring(error)) => { - let _ = association.fail(ring_error(error)); - return; - } - Err(ControlRingReadError::Decode(error)) => { - let _ = association.fail(wire_error(error)); - return; - } - } - } -} - #[cfg(test)] mod control_ring_tests { use super::*; @@ -611,6 +624,7 @@ mod control_ring_tests { on_failure: impl Fn() + Send + Sync + 'static, ) -> ( UnixControlRingLocalCallChannel, + UnixControlRingLocalNotificationChannel, UnixControlRingLocalShutdown, Producer, Consumer, @@ -627,8 +641,7 @@ mod control_ring_tests { }); let (local_ring, broker_ring) = ring_pair(); let setup = negotiated_local(local_stream); - let (channel, _notifications, shutdown) = - setup.into_active(local_ring, on_failure).unwrap(); + let (channel, notifications, shutdown) = setup.into_active(local_ring, on_failure).unwrap(); acknowledgement.join().unwrap(); let litebox_broker_transport::control_ring::BrokerControlRingEndpoints { request_consumer, @@ -637,6 +650,7 @@ mod control_ring_tests { } = broker_ring.into_broker(); ( channel, + notifications, shutdown, response_producer, request_consumer, @@ -750,7 +764,8 @@ mod control_ring_tests { #[test] fn local_matches_out_of_order_ring_responses_without_socket_frames() { - let (channel, _shutdown, mut responses, mut requests, mut peer) = activate_local(|| {}); + let (channel, _notifications, _shutdown, mut responses, mut requests, mut peer) = + activate_local(|| {}); peer.set_read_timeout(Some(Duration::from_millis(100))) .unwrap(); let channel = Arc::new(channel); @@ -782,7 +797,7 @@ mod control_ring_tests { let failures = Arc::new(AtomicUsize::new(0)); let callback_failures = Arc::clone(&failures); let (failure_reported, wait_for_failure) = mpsc::channel(); - let (channel, _shutdown, mut responses, mut requests, _peer) = + let (channel, _notifications, _shutdown, mut responses, mut requests, _peer) = activate_local(move || { callback_failures.fetch_add(1, Ordering::SeqCst); failure_reported.send(()).unwrap(); @@ -818,7 +833,8 @@ mod control_ring_tests { #[test] fn local_socket_eof_and_shutdown_wake_pending_calls() { for close_peer in [false, true] { - let (channel, shutdown, _responses, mut requests, peer) = activate_local(|| {}); + let (channel, _notifications, shutdown, _responses, mut requests, peer) = + activate_local(|| {}); let caller = thread::spawn(move || channel.call(request(1))); read_request(&mut requests); if close_peer { @@ -830,6 +846,26 @@ mod control_ring_tests { } } + #[test] + fn notifications_end_cleanly_only_when_the_call_channel_is_dropped() { + for end in 0..3 { + let (channel, mut notifications, shutdown, _responses, _requests, peer) = + activate_local(|| {}); + let notification_receiver = thread::spawn(move || notifications.recv_notification()); + match end { + 0 => drop(channel), + 1 => drop(peer), + _ => shutdown.shutdown().unwrap(), + } + let received = notification_receiver.join().unwrap(); + if end == 0 { + assert!(received.unwrap().is_none()); + } else { + assert!(received.is_err()); + } + } + } + #[test] fn ready_ack_uses_absolute_setup_deadline() { let (local_stream, _peer) = UnixStream::pair().unwrap(); @@ -943,7 +979,10 @@ mod control_ring_tests { ErrorKind::ConnectionAborted, "test failure", ))); - assert_eq!(completed.wait().unwrap().request_id, RequestId(1)); + assert_eq!( + completed.wait(|| unreachable!()).unwrap().request_id, + RequestId(1) + ); let pending = PendingCalls::new(); let failed = pending.register(RequestId(2)).unwrap(); @@ -953,7 +992,7 @@ mod control_ring_tests { ))); assert!(pending.complete(response(RequestId(2))).is_err()); assert_eq!( - failed.wait().unwrap_err().kind(), + failed.wait(|| unreachable!()).unwrap_err().kind(), ErrorKind::ConnectionAborted ); } @@ -1006,7 +1045,7 @@ mod control_ring_tests { .unwrap(); failure.join().unwrap(); assert_eq!( - pending_call.wait().unwrap_err().kind(), + pending_call.wait(|| unreachable!()).unwrap_err().kind(), ErrorKind::ConnectionAborted ); } @@ -1020,6 +1059,9 @@ mod control_ring_tests { }; assert_eq!(error.kind(), ErrorKind::InvalidData); pending.complete(response(RequestId(1))).unwrap(); - assert_eq!(original.wait().unwrap().request_id, RequestId(1)); + assert_eq!( + original.wait(|| unreachable!()).unwrap().request_id, + RequestId(1) + ); } } diff --git a/litebox_broker_transport_windows_userland/src/local.rs b/litebox_broker_transport_windows_userland/src/local.rs index 96d263d6ae..6ce4737fdc 100644 --- a/litebox_broker_transport_windows_userland/src/local.rs +++ b/litebox_broker_transport_windows_userland/src/local.rs @@ -7,6 +7,7 @@ use std::ffi::OsStr; use std::fs::OpenOptions; use std::io::{Error, ErrorKind, Result as IoResult}; use std::os::windows::fs::OpenOptionsExt; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::thread; use std::time::{Duration, Instant}; @@ -60,7 +61,12 @@ pub struct WindowsControlRingLocalNotificationChannel { struct LocalRingAssociation { request_producer: Mutex>, + /// Read by whichever caller holds the pending-call reader role. + response_consumer: Mutex>, pending_calls: Arc, + /// Whether the call channel was dropped, which ends the association as a + /// local close rather than a failure. + closed: AtomicBool, liveness: PipeLiveness, request_wake: ControlRingWakeHandle, response_wake: ControlRingWakeHandle, @@ -202,13 +208,11 @@ fn activate_local( request_wake: request_producer.wake_handle(), request_producer: Mutex::new(request_producer), response_wake: response_consumer.wake_handle(), + response_consumer: Mutex::new(response_consumer), notification_wake: notification_consumer.wake_handle(), pending_calls, + closed: AtomicBool::new(false), }); - let response_association = Arc::clone(&association); - thread::Builder::new() - .name("litebox-broker-responses".to_owned()) - .spawn(move || dispatch_responses(response_consumer, response_association))?; let monitor_association = Arc::clone(&association); thread::Builder::new() .name("litebox-broker-liveness".to_owned()) @@ -261,12 +265,15 @@ impl LocalCallChannel for WindowsControlRingLocalCallChannel { if let Err(error) = write_result { let _ = association.fail(error); } - pending_call.wait().map_err(|error| copy_io_error(&error)) + pending_call + .wait(|| association.read_response()) + .map_err(|error| copy_io_error(&error)) } } impl Drop for WindowsControlRingLocalCallChannel { fn drop(&mut self) { + self.association.closed.store(true, Ordering::Release); let _ = self.association.fail(Error::new( ErrorKind::ConnectionAborted, "broker local call channel dropped", @@ -280,7 +287,7 @@ impl LocalNotificationChannel for WindowsControlRingLocalNotificationChannel { fn recv_notification(&mut self) -> IoResult> { loop { if let Some(error) = self.association.pending_calls.current_failure() { - return Err(copy_io_error(&error)); + return self.association.end_notifications(&error); } match self.consumer.try_read(decode_notification) { Ok(ControlRingReadStatus::Message(notification)) => { @@ -290,7 +297,7 @@ impl LocalNotificationChannel for WindowsControlRingLocalNotificationChannel { } Ok(ControlRingReadStatus::Empty { wait_epoch }) => { if let Some(error) = self.association.pending_calls.current_failure() { - return Err(copy_io_error(&error)); + return self.association.end_notifications(&error); } if let Err(error) = self.consumer.wait_for_message(wait_epoch) { let result = Err(copy_io_error(&error)); @@ -316,6 +323,16 @@ impl LocalNotificationChannel for WindowsControlRingLocalNotificationChannel { } impl LocalRingAssociation { + /// Ends notification receipt once the association has ended with + /// `failure`: cleanly if the local side closed it. + fn end_notifications(&self, failure: &Error) -> IoResult> { + if self.closed.load(Ordering::Acquire) { + Ok(None) + } else { + Err(copy_io_error(failure)) + } + } + fn acknowledge_notification( &self, consumer: &mut ControlRingConsumer, @@ -337,57 +354,59 @@ impl LocalRingAssociation { Ok(()) } - fn fail(&self, error: Error) -> IoResult<()> { - self.pending_calls.record_failure(Arc::new(error)); - self.request_wake - .interrupt_wait() - .and(self.response_wake.interrupt_wait()) - .and(self.notification_wake.interrupt_wait()) - .and(self.liveness.shutdown()) - } -} - -fn dispatch_responses( - mut consumer: ControlRingConsumer, - association: Arc, -) { - loop { - match consumer.try_read(decode_response) { - Ok(ControlRingReadStatus::Message(response)) => { - if let Err(error) = consumer - .publish_head() - .map_err(ring_error) - .and_then(|()| consumer.wake_producer()) - .and_then(|()| { - association - .pending_calls - .complete(response) - .map_err(pending_calls_error) - }) - { - let _ = association.fail(error); + /// Reads one response for the caller holding the pending-call reader + /// role, blocking until one arrives or the association fails. + fn read_response(&self) { + let Ok(mut consumer) = self.response_consumer.lock() else { + let _ = self.fail(Error::other("broker response reader mutex poisoned")); + return; + }; + loop { + match consumer.try_read(decode_response) { + Ok(ControlRingReadStatus::Message(response)) => { + if let Err(error) = consumer + .publish_head() + .map_err(ring_error) + .and_then(|()| consumer.wake_producer()) + .and_then(|()| { + self.pending_calls + .complete(response) + .map_err(pending_calls_error) + }) + { + let _ = self.fail(error); + } return; } - } - Ok(ControlRingReadStatus::Empty { wait_epoch }) => { - if association.pending_calls.current_failure().is_some() { + Ok(ControlRingReadStatus::Empty { wait_epoch }) => { + if self.pending_calls.current_failure().is_some() { + return; + } + if let Err(error) = consumer.wait_for_message(wait_epoch) { + let _ = self.fail(error); + return; + } + } + Err(ControlRingReadError::Ring(error)) => { + let _ = self.fail(ring_error(error)); return; } - if let Err(error) = consumer.wait_for_message(wait_epoch) { - let _ = association.fail(error); + Err(ControlRingReadError::Decode(error)) => { + let _ = self.fail(wire_error(error)); return; } } - Err(ControlRingReadError::Ring(error)) => { - let _ = association.fail(ring_error(error)); - return; - } - Err(ControlRingReadError::Decode(error)) => { - let _ = association.fail(wire_error(error)); - return; - } } } + + fn fail(&self, error: Error) -> IoResult<()> { + self.pending_calls.record_failure(Arc::new(error)); + self.request_wake + .interrupt_wait() + .and(self.response_wake.interrupt_wait()) + .and(self.notification_wake.interrupt_wait()) + .and(self.liveness.shutdown()) + } } fn monitor_local_pipe(stream: &WindowsNamedPipeStream, association: &LocalRingAssociation) { diff --git a/litebox_broker_transport_windows_userland/src/shared_memory.rs b/litebox_broker_transport_windows_userland/src/shared_memory.rs index 227b3b7e27..88d6fd0291 100644 --- a/litebox_broker_transport_windows_userland/src/shared_memory.rs +++ b/litebox_broker_transport_windows_userland/src/shared_memory.rs @@ -3,7 +3,6 @@ //! Windows file-mapping-backed broker shared memory. -use std::arch::asm; use std::io::{Error, ErrorKind, Result as IoResult}; use std::ptr::NonNull; use std::sync::Arc; @@ -11,6 +10,7 @@ use std::sync::Arc; use litebox_broker_transport::control_ring::{ CONTROL_RING_MEMORY_SIZE, ControlRingDirection, MemoryAccessPolicy, WaitableSharedMemory, }; +use litebox_broker_transport::peer_memory; use litebox_broker_transport::shared_memory::{ControlRingMemory, SharedMemory, SharedMemoryError}; use windows_sys::Win32::Foundation::{ CloseHandle, DUPLICATE_SAME_ACCESS, DuplicateHandle, HANDLE, INVALID_HANDLE_VALUE, WAIT_FAILED, @@ -36,60 +36,6 @@ const CONTROL_RING_WAKE_OFFSETS: [usize; 6] = [ ControlRingDirection::Notifications.consumer_epoch_offset(), ]; -unsafe fn atomic_load_u32(address: *const u32) -> u32 { - let value; - // SAFETY: The caller provides a readable, naturally aligned u32 address. An aligned x86-64 - // load is indivisible, and the compiler memory clobber plus x86-64 TSO provides acquire order. - unsafe { - asm!( - "mov {value:e}, dword ptr [{address}]", - value = out(reg) value, - address = in(reg) address, - options(nostack, preserves_flags), - ); - } - value -} - -unsafe fn atomic_increment_u32(address: *mut u32) { - // SAFETY: The caller provides a writable, naturally aligned u32 address. - unsafe { - asm!( - "lock inc dword ptr [{address}]", - address = in(reg) address, - options(nostack, preserves_flags), - ); - } -} - -unsafe fn atomic_load_u64(address: *const u64) -> u64 { - let value; - // SAFETY: The caller provides a readable, naturally aligned u64 address. Aligned x86-64 - // qword loads are indivisible. - unsafe { - asm!( - "mov {value}, qword ptr [{address}]", - value = out(reg) value, - address = in(reg) address, - options(nostack, preserves_flags), - ); - } - value -} - -unsafe fn atomic_store_u64(address: *mut u64, value: u64) { - // SAFETY: The caller provides a writable, naturally aligned u64 address. `xchg` with memory - // is indivisible and fully ordered on x86-64. - unsafe { - asm!( - "xchg qword ptr [{address}], {value}", - address = in(reg) address, - value = inout(reg) value => _, - options(nostack, preserves_flags), - ); - } -} - struct OwnedHandle(HANDLE); impl Drop for OwnedHandle { @@ -445,26 +391,26 @@ impl ControlRingMemory for WindowsSharedMemory { fn load_u32_acquire(&self, offset: usize) -> Result { let address = self.checked_u32(offset)?; // SAFETY: `checked_u32` validated the mapped range and alignment. - Ok(unsafe { atomic_load_u32(address) }) + Ok(unsafe { peer_memory::load_u32_acquire(address) }) } fn increment_u32_release(&self, offset: usize) -> Result<(), SharedMemoryError> { let address = self.checked_u32(offset)?; // SAFETY: `address` is checked and naturally aligned mapped memory. - unsafe { atomic_increment_u32(address) }; + unsafe { peer_memory::increment_u32_release(address) }; Ok(()) } fn load_u64_acquire(&self, offset: usize) -> Result { let address = self.checked_u64(offset)?; // SAFETY: `address` is checked and naturally aligned mapped memory. - Ok(unsafe { atomic_load_u64(address) }) + Ok(unsafe { peer_memory::load_u64_acquire(address) }) } fn store_u64_release(&self, offset: usize, value: u64) -> Result<(), SharedMemoryError> { let address = self.checked_u64(offset)?; // SAFETY: `address` is checked and naturally aligned mapped memory. - unsafe { atomic_store_u64(address, value) }; + unsafe { peer_memory::store_u64_release(address, value) }; Ok(()) } @@ -483,8 +429,8 @@ impl ControlRingMemory for WindowsSharedMemory { } // SAFETY: Both addresses are checked, aligned, and non-overlapping mapped words. unsafe { - atomic_store_u64(store_address, value); - atomic_increment_u32(increment_address); + peer_memory::store_u64_release(store_address, value); + peer_memory::increment_u32_release(increment_address); } Ok(()) } @@ -506,7 +452,7 @@ impl WaitableSharedMemory for WindowsSharedMemory { // Rechecking after resolving the event closes the publish-before-wait race. A signal that // arrives after this load remains pending on the auto-reset event until this waiter runs. // SAFETY: `checked_u32` validated the mapped range and alignment. - if unsafe { atomic_load_u32(address) } != expected { + if unsafe { peer_memory::load_u32_acquire(address) } != expected { return Ok(()); } // SAFETY: `event` is a live event handle owned by this memory object. diff --git a/litebox_broker_userland/Cargo.toml b/litebox_broker_userland/Cargo.toml index f197d68daa..47d4472353 100644 --- a/litebox_broker_userland/Cargo.toml +++ b/litebox_broker_userland/Cargo.toml @@ -24,6 +24,7 @@ tempfile = { version = "3", default-features = false } litebox_broker_platform_windows_userland = { path = "../litebox_broker_platform_windows_userland", version = "0.1.0" } litebox_broker_transport_windows_userland = { path = "../litebox_broker_transport_windows_userland", version = "0.1.0" } litebox_runner_windows_userland = { path = "../litebox_runner_windows_userland", version = "0.1.0" } +windows-sys = { version = "0.60.2", features = ["Win32_Foundation", "Win32_System_Memory"] } [features] lock_tracing = [ diff --git a/litebox_broker_userland/src/builder.rs b/litebox_broker_userland/src/builder.rs index 5bd46a2da7..d4c9c3ab3e 100644 --- a/litebox_broker_userland/src/builder.rs +++ b/litebox_broker_userland/src/builder.rs @@ -137,9 +137,10 @@ fn platform_socket_provider( limits: &BrokerCoreLimits, ) -> Result, BrokerBuildError> { Ok(Arc::new( - litebox_broker_platform_linux_userland::LinuxSocketProvider::new( + litebox_broker_platform_linux_userland::LinuxSocketProvider::with_wait_hook( limits.max_sockets, limits.max_sockets_per_process, + crate::runtime::before_blocking, )?, )) } diff --git a/litebox_broker_userland/src/lib.rs b/litebox_broker_userland/src/lib.rs index c666a8f2c9..c6183cce92 100644 --- a/litebox_broker_userland/src/lib.rs +++ b/litebox_broker_userland/src/lib.rs @@ -16,6 +16,8 @@ pub mod builder; #[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] +pub mod mapped_file; +#[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] mod process_launcher; pub mod readiness; #[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] @@ -26,4 +28,5 @@ pub mod random; pub mod stdio; mod timer; +/// The most request workers one association runs at a time. const WORKER_COUNT: usize = 8; diff --git a/litebox_broker_userland/src/linux.rs b/litebox_broker_userland/src/linux.rs index 770aedfcee..c44f2ecd87 100644 --- a/litebox_broker_userland/src/linux.rs +++ b/litebox_broker_userland/src/linux.rs @@ -22,9 +22,10 @@ use litebox_broker_transport_linux_userland::unix_socket::{ UnixStreamHostSetupChannel, validate_peer_process, }; use litebox_broker_userland::builder::BrokerCoreBuilder; +use litebox_broker_userland::runner::{accept_runner_channel, wait_for_connection}; use litebox_broker_userland::stdio::UserlandStdioProvider; -use super::{SETUP_TIMEOUT, configured_socket_policy}; +use super::{IN_PROCESS_RUNNER_POLL_INTERVAL, SETUP_TIMEOUT, configured_socket_policy}; const PROXY_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); @@ -214,11 +215,17 @@ fn serve_runner_in_process( runner: &JoinHandle, ) -> IoResult<()> { let setup_deadline = Instant::now() + SETUP_TIMEOUT; - let control_stream = crate::accept_runner_channel( + let control_stream = accept_runner_channel( setup_deadline, "control", || Ok(runner.is_finished().then(|| "thread stopped".to_owned())), || control_listener.accept().map(|(stream, _)| stream), + |remaining| { + wait_for_connection( + control_listener, + remaining.min(IN_PROCESS_RUNNER_POLL_INTERVAL), + ) + }, )?; validate_peer_process(&control_stream, std::process::id())?; serve_control_stream(broker, control_stream, setup_deadline) diff --git a/litebox_broker_userland/src/main.rs b/litebox_broker_userland/src/main.rs index a00b15ea7b..736e98e70c 100644 --- a/litebox_broker_userland/src/main.rs +++ b/litebox_broker_userland/src/main.rs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT license. -use std::borrow::Cow; use std::error::Error; use std::ffi::{OsStr, OsString}; use std::io::{Error as IoError, ErrorKind, Result as IoResult}; @@ -10,7 +9,7 @@ use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; use std::thread::JoinHandle; -use std::time::{Duration, Instant}; +use std::time::Duration; use clap::Parser; #[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] @@ -27,6 +26,7 @@ use litebox_broker_core::{ }; use litebox_broker_protocol::fs::{FileMode as Mode, FileUser as UserInfo}; use litebox_broker_protocol::socket::{Ipv4Address, Port}; +use litebox_broker_userland::mapped_file::MappedFile; use litebox_broker_userland::random::UserlandRandomProvider; use litebox_broker_userland::stdio::UserlandStdioProvider; use litebox_platform::sync::RawSyncPrimitivesProvider; @@ -37,7 +37,9 @@ mod linux; mod windows; const SETUP_TIMEOUT: Duration = Duration::from_secs(5); -const ACCEPT_RETRY_DELAY: Duration = Duration::from_millis(10); +/// How often an in-process runner thread, which cannot be waited on alongside +/// its connection, is rechecked for having stopped before connecting. +const IN_PROCESS_RUNNER_POLL_INTERVAL: Duration = Duration::from_millis(10); #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct AllowedDestination { @@ -129,6 +131,9 @@ struct CliArgs { )] runner: Option, /// Tar archive to mount as the broker-owned initial file system. + /// + /// The archive is mapped into memory rather than copied, so it must not be modified while the + /// broker runs. #[arg(long, value_name = "PATH", value_hint = clap::ValueHint::FilePath)] fs_initial_files: Option, /// Opaque arguments to pass to the local runner without interpretation. @@ -152,7 +157,7 @@ where ("/registry".to_owned(), writable_directory(UserInfo::ROOT)), ]; - let tar_data = match initial_files { + let tar_file = match initial_files { Some(path) => { if path.extension().and_then(|extension| extension.to_str()) != Some("tar") { return Err(IoError::new( @@ -160,18 +165,20 @@ where format!("expected a .tar file, found {}", path.display()), )); } - Cow::Owned(std::fs::read(path)?) + // SAFETY: As documented on `--fs-initial-files`, the archive must not be modified + // while the broker runs. + Some(unsafe { MappedFile::open(path) }?) } - None => Cow::Borrowed(EMPTY_TAR_FILE), + None => None, }; let in_mem = InMem::::new_initialized(entries); let backend = Composer::builder() .mount_nestable("/", |allocators| { - Overlay::::new( - in_mem, - TarRo::new(tar_data, allocators.next()), - allocators.next(), - ) + let lower = match tar_file { + Some(tar_file) => TarRo::new(tar_file, allocators.next()), + None => TarRo::new(EMPTY_TAR_FILE, allocators.next()), + }; + Overlay::::new(in_mem, lower, allocators.next()) }) .mount("/dev", |allocator| { litebox_broker_core::fs::devices::Devices::new( @@ -302,35 +309,6 @@ fn destination_rules(allowed_destinations: &[AllowedDestination]) -> Vec( - deadline: Instant, - channel_name: &'static str, - mut runner_status: impl FnMut() -> IoResult>, - mut try_accept: impl FnMut() -> IoResult, -) -> IoResult { - loop { - let remaining = deadline.saturating_duration_since(Instant::now()); - if remaining.is_zero() { - return Err(IoError::new( - ErrorKind::TimedOut, - format!("timed out waiting for runner {channel_name} channel"), - )); - } - if let Some(status) = runner_status()? { - return Err(IoError::new( - ErrorKind::BrokenPipe, - format!("runner {status} before connecting its {channel_name} channel"), - )); - } - match try_accept() { - Ok(channel) => return Ok(channel), - Err(error) if error.kind() == ErrorKind::WouldBlock => {} - Err(error) => return Err(error), - } - std::thread::sleep(remaining.min(ACCEPT_RETRY_DELAY)); - } -} - #[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] fn main() -> Result<(), Box> { let args = CliArgs::parse(); diff --git a/litebox_broker_userland/src/mapped_file.rs b/litebox_broker_userland/src/mapped_file.rs new file mode 100644 index 0000000000..daa2d16b02 --- /dev/null +++ b/litebox_broker_userland/src/mapped_file.rs @@ -0,0 +1,185 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +//! Read-only file mappings for userland brokers. + +use std::fs::File; +use std::io::{Error as IoError, ErrorKind, Result as IoResult}; +use std::path::Path; +use std::ptr::NonNull; + +/// The contents of a file, mapped read-only into the broker's address space. +/// +/// This lets the broker serve a large read-only input, such as the initial file system archive, +/// without first copying all of it into memory: pages come from the host's page cache as they are +/// read. +pub struct MappedFile { + /// Start of the mapping, or dangling if `len` is zero. + address: NonNull, + len: usize, +} + +// SAFETY: The mapping is read-only and owned by this value alone, so it can be sent and shared +// across threads just like an owned `Box<[u8]>`. +unsafe impl Send for MappedFile {} +// SAFETY: See `Send` above. +unsafe impl Sync for MappedFile {} + +impl MappedFile { + /// Map the whole file at `path`. + /// + /// # Safety + /// + /// The file must not be modified, by this or any other process, while the mapping is alive. + /// The mapping can observe such changes, which would break the immutability of the bytes it + /// lends out, and reading past the end of a file truncated in the meantime faults. + pub unsafe fn open(path: &Path) -> IoResult { + let file = File::open(path)?; + let len = usize::try_from(file.metadata()?.len()) + .map_err(|_| IoError::new(ErrorKind::InvalidInput, "file is too large to map"))?; + if len == 0 { + // Hosts refuse to map empty files, and there are no bytes to lend out anyway. + return Ok(Self { + address: NonNull::dangling(), + len, + }); + } + let address = host::map(&file, len)?; + Ok(Self { address, len }) + } +} + +impl AsRef<[u8]> for MappedFile { + fn as_ref(&self) -> &[u8] { + // SAFETY: `address` is either dangling with a zero `len`, or the start of a live, readable + // mapping of `len` bytes owned by `self`, whose contents the contract of `open` keeps + // unchanged. + unsafe { std::slice::from_raw_parts(self.address.as_ptr(), self.len) } + } +} + +impl Drop for MappedFile { + fn drop(&mut self) { + if self.len != 0 { + // SAFETY: A non-empty `self` owns the live mapping at `address`, and no borrow of its + // bytes can outlive `self`. + unsafe { host::unmap(self.address, self.len) }; + } + } +} + +#[cfg(target_os = "linux")] +mod host { + use std::fs::File; + use std::io::{Error as IoError, Result as IoResult}; + use std::os::fd::AsRawFd; + use std::ptr::NonNull; + + /// Map the first `len` bytes of `file` read-only. + pub(super) fn map(file: &File, len: usize) -> IoResult> { + // SAFETY: A new private mapping at an address chosen by the kernel aliases no existing + // memory, and `file` is a live descriptor. + let address = unsafe { + libc::mmap( + std::ptr::null_mut(), + len, + libc::PROT_READ, + libc::MAP_PRIVATE, + file.as_raw_fd(), + 0, + ) + }; + if address == libc::MAP_FAILED { + return Err(IoError::last_os_error()); + } + Ok(NonNull::new(address.cast()).expect("a successful mmap is not at address zero")) + } + + /// Unmap a mapping returned by [`map`]. + /// + /// # Safety + /// + /// `address` and `len` must describe a live mapping returned by [`map`] that is no longer + /// borrowed. + pub(super) unsafe fn unmap(address: NonNull, len: usize) { + // SAFETY: Guaranteed by the caller. + unsafe { libc::munmap(address.as_ptr().cast(), len) }; + } +} + +#[cfg(all(windows, target_arch = "x86_64"))] +mod host { + use std::fs::File; + use std::io::{Error as IoError, Result as IoResult}; + use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; + use std::ptr::NonNull; + + use windows_sys::Win32::System::Memory::{ + CreateFileMappingW, FILE_MAP_READ, MEMORY_MAPPED_VIEW_ADDRESS, MapViewOfFile, + PAGE_READONLY, UnmapViewOfFile, + }; + + /// Map the first `len` bytes of `file` read-only. + pub(super) fn map(file: &File, len: usize) -> IoResult> { + // SAFETY: `file` is a live handle, and the unnamed read-only mapping spans the whole file. + let mapping = unsafe { + CreateFileMappingW( + file.as_raw_handle(), + std::ptr::null(), + PAGE_READONLY, + 0, + 0, + std::ptr::null(), + ) + }; + if mapping.is_null() { + return Err(IoError::last_os_error()); + } + // SAFETY: `mapping` is a fresh handle owned by nothing else. Closing it once the view is + // mapped is fine, since the view keeps the mapping alive. + let mapping = unsafe { OwnedHandle::from_raw_handle(mapping) }; + // SAFETY: `mapping` is a live read-only file mapping spanning `len` bytes, and a view at an + // address chosen by the system aliases no existing memory. + let view = unsafe { MapViewOfFile(mapping.as_raw_handle(), FILE_MAP_READ, 0, 0, len) }; + NonNull::new(view.Value.cast()).ok_or_else(IoError::last_os_error) + } + + /// Unmap a mapping returned by [`map`]. + /// + /// # Safety + /// + /// `address` must be a live mapping returned by [`map`] that is no longer borrowed. + pub(super) unsafe fn unmap(address: NonNull, _len: usize) { + // SAFETY: Guaranteed by the caller. + unsafe { + UnmapViewOfFile(MEMORY_MAPPED_VIEW_ADDRESS { + Value: address.as_ptr().cast(), + }) + }; + } +} + +#[cfg(all(test, target_os = "linux"))] +mod tests { + use std::io::Write; + + use super::MappedFile; + + fn mapped(contents: &[u8]) -> MappedFile { + let mut file = tempfile::NamedTempFile::new().unwrap(); + file.write_all(contents).unwrap(); + // SAFETY: Nothing modifies the temporary file while it is mapped. + unsafe { MappedFile::open(file.path()) }.unwrap() + } + + #[test] + fn maps_file_contents() { + let contents: Vec = (0..10_000u32).flat_map(u32::to_le_bytes).collect(); + assert_eq!(mapped(&contents).as_ref(), contents); + } + + #[test] + fn maps_empty_file() { + assert!(mapped(b"").as_ref().is_empty()); + } +} diff --git a/litebox_broker_userland/src/process_launcher.rs b/litebox_broker_userland/src/process_launcher.rs index 2e5401e5e7..0da0c9200b 100644 --- a/litebox_broker_userland/src/process_launcher.rs +++ b/litebox_broker_userland/src/process_launcher.rs @@ -160,6 +160,8 @@ impl UserlandProcessLauncher { config: RunnerConfig, completion_sender: Option>>, ) -> Result<(), BrokerError> { + // Starting a runner waits for its setup to finish. + crate::runtime::before_blocking(); let process = Arc::clone(&association.process); let process_id = process.id(); let Ok(instance) = RunnerInstance::start(config) else { diff --git a/litebox_broker_userland/src/runner.rs b/litebox_broker_userland/src/runner.rs index 9f05c1d485..17bddae1de 100644 --- a/litebox_broker_userland/src/runner.rs +++ b/litebox_broker_userland/src/runner.rs @@ -415,7 +415,12 @@ const _: () = { assert!(runner_exit_code_is_crash(Some(breakpoint))); }; -fn accept_runner_channel( +/// Accepts a runner's `channel_name` channel with `try_accept`, calling `wait` +/// with the remaining setup time whenever no connection is pending. +/// +/// Fails once `runner_status` reports that the runner stopped or `deadline` +/// passes. +pub fn accept_runner_channel( deadline: Instant, channel_name: &'static str, mut runner_status: impl FnMut() -> IoResult>, @@ -445,6 +450,17 @@ fn accept_runner_channel( } } +/// Waits until `listener` has a pending connection or `timeout` elapses. +/// Callers recheck their own condition, so early returns are harmless. +#[cfg(target_os = "linux")] +pub fn wait_for_connection( + listener: &std::os::unix::net::UnixListener, + timeout: Duration, +) -> IoResult<()> { + use std::os::fd::AsFd; + linux::poll_readable([Some(listener.as_fd()), None], Some(timeout)) +} + fn runner_has_exited(runner: &Arc>) -> IoResult { // A pre-authentication caller stops accepting before acting on `true`; // post-authentication callers no longer rely on PID-based authentication. diff --git a/litebox_broker_userland/src/runner/linux.rs b/litebox_broker_userland/src/runner/linux.rs index 6e651c9f38..17823d8fbb 100644 --- a/litebox_broker_userland/src/runner/linux.rs +++ b/litebox_broker_userland/src/runner/linux.rs @@ -5,7 +5,7 @@ use std::any::Any; use std::ffi::OsStr; use std::io::{Error as IoError, ErrorKind, Result as IoResult}; use std::ops::Range; -use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd}; +use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd}; use std::os::unix::net::UnixListener; use std::os::unix::process::CommandExt; use std::path::PathBuf; @@ -236,16 +236,23 @@ pub(super) fn wait_for_runner_event( Some(_) => timeout, None => Some(timeout.map_or(RUNNER_POLL_INTERVAL, |t| t.min(RUNNER_POLL_INTERVAL))), }; + poll_readable( + [pidfd.as_ref().map(AsFd::as_fd), listener.map(AsFd::as_fd)], + timeout, + ) +} + +/// Waits until any of `fds` becomes readable or `timeout` elapses. +pub(super) fn poll_readable( + fds: [Option>; 2], + timeout: Option, +) -> IoResult<()> { let timeout_ms = timeout.map_or(-1, |timeout| { i32::try_from(timeout.as_nanos().div_ceil(1_000_000)).unwrap_or(i32::MAX) }); // `poll` ignores entries with negative descriptors. - let mut fds = [ - pidfd.as_ref().map_or(-1, AsRawFd::as_raw_fd), - listener.map_or(-1, AsRawFd::as_raw_fd), - ] - .map(|fd| libc::pollfd { - fd, + let mut fds = fds.map(|fd| libc::pollfd { + fd: fd.map_or(-1, |fd| fd.as_raw_fd()), events: libc::POLLIN, revents: 0, }); diff --git a/litebox_broker_userland/src/runtime.rs b/litebox_broker_userland/src/runtime.rs index e719ed7b16..c6898dc397 100644 --- a/litebox_broker_userland/src/runtime.rs +++ b/litebox_broker_userland/src/runtime.rs @@ -15,15 +15,15 @@ //! //! Concurrency and worker sizing are deliberately not part of the public //! surface. Both association entry points delegate to one internal runtime -//! that owns worker counts and request queue capacity. +//! that owns the worker count. +use std::cell::Cell; use std::io::{Error as IoError, ErrorKind, Result as IoResult}; -use std::sync::mpsc::{Receiver, SyncSender, TrySendError, sync_channel}; use std::sync::{ - Arc, Mutex, + Arc, Mutex, MutexGuard, atomic::{AtomicBool, Ordering}, }; -use std::time::{Duration, Instant}; +use std::thread::{JoinHandle, Thread}; use litebox_broker_core::BrokerCore; use litebox_broker_host::{ @@ -43,10 +43,6 @@ use litebox_broker_transport::shared_memory::{ControlRingMemory, SharedBufferPoo use crate::process_launcher::{PendingRunnerAssociation, UserlandProcessLauncher}; use crate::readiness::ReadinessPublisherRuntime; -const REQUEST_QUEUE_CAPACITY: usize = 64; -const REQUEST_QUEUE_RETRY_DELAY: Duration = Duration::from_millis(1); -const REQUEST_QUEUE_STALL_TIMEOUT: Duration = Duration::from_secs(5); - pub(crate) struct AssociationOutcome { pub(crate) result: IoResult<()>, pub(crate) abnormal: bool, @@ -96,8 +92,8 @@ pub fn serve_in_process_runner_association< where Memory: ControlRingMemory, SetupChannel: HostSetupChannel, - RequestSource: HostRequestSource, - ResponseSink: HostResponseSink + Clone + Send, + RequestSource: HostRequestSource + Send + 'static, + ResponseSink: HostResponseSink + Send + Sync + 'static, NotificationChannel: HostNotificationChannel + Send, Shutdown: HostAssociationShutdown + Send + Sync + 'static, { @@ -142,8 +138,8 @@ pub(crate) fn serve_out_of_process_runner_association< where Memory: ControlRingMemory, SetupChannel: HostSetupChannel, - RequestSource: HostRequestSource, - ResponseSink: HostResponseSink + Clone + Send, + RequestSource: HostRequestSource + Send + 'static, + ResponseSink: HostResponseSink + Send + Sync + 'static, NotificationChannel: HostNotificationChannel + Send, Shutdown: HostAssociationShutdown + Send + Sync + 'static, { @@ -196,8 +192,8 @@ fn serve_association_inner< where Memory: ControlRingMemory, SetupChannel: HostSetupChannel, - RequestSource: HostRequestSource, - ResponseSink: HostResponseSink + Clone + Send, + RequestSource: HostRequestSource + Send + 'static, + ResponseSink: HostResponseSink + Send + Sync + 'static, NotificationChannel: HostNotificationChannel + Send, Shutdown: HostAssociationShutdown + Send + Sync + 'static, { @@ -446,7 +442,7 @@ impl Drop for AssociationCancellationGuard<'_, Memory> { fn dispatch_requests( association: BrokerHostAssociation, readiness: Arc, - mut request_source: RequestSource, + request_source: RequestSource, response_sink: ResponseSink, mut notification_channel: NotificationChannel, shutdown: Shutdown, @@ -455,8 +451,8 @@ fn dispatch_requests AssociationOutcome where Memory: SharedMemory, - RequestSource: HostRequestSource, - ResponseSink: HostResponseSink + Clone + Send, + RequestSource: HostRequestSource + Send + 'static, + ResponseSink: HostResponseSink + Send + Sync + 'static, NotificationChannel: HostNotificationChannel + Send, Shutdown: HostAssociationShutdown + Send + Sync + 'static, { @@ -470,8 +466,14 @@ where abnormal: true, }; } - let (request_sender, request_receiver) = sync_channel(REQUEST_QUEUE_CAPACITY); - let request_receiver = Arc::new(Mutex::new(request_receiver)); + let workers = Arc::new(Workers { + association: Arc::clone(&association), + requests: Mutex::new(request_source), + response_sink, + failure_coordinator: Arc::clone(&failure_coordinator), + launcher, + role: ReceivingRole::new(), + }); std::thread::scope(|scope| { let publisher_readiness = Arc::clone(&readiness); @@ -482,9 +484,9 @@ where let _panicking = PublisherPanicGuard { failure_coordinator: &publisher_failure_coordinator, }; - // The request reader owns association termination. A failing + // The request workers own association termination. A failing // notification transport must fail the association before - // returning an error, so a reader still running observes and + // returning an error, so a worker still receiving observes and // reports the same failure. Reporting here would instead turn // a clean peer close into an error when its transport teardown // releases a blocked notification send. @@ -510,36 +512,10 @@ where association: &association, }; - let mut workers = Vec::with_capacity(crate::WORKER_COUNT); - for worker_id in 0..crate::WORKER_COUNT { - let association = Arc::clone(&association); - let request_receiver = Arc::clone(&request_receiver); - let response_sink = response_sink.clone(); - let worker_failure_coordinator = Arc::clone(&failure_coordinator); - let launcher_for_worker = launcher.clone(); - match std::thread::Builder::new() - .name(format!("litebox-broker-worker-{worker_id}")) - .spawn_scoped(scope, move || { - run_worker( - &association, - &request_receiver, - &response_sink, - &worker_failure_coordinator, - launcher_for_worker.as_ref(), - ); - }) { - Ok(worker) => workers.push(worker), - Err(error) => { - failure_coordinator.report(error); - break; - } - } - } - - read_requests(&mut request_source, request_sender, &failure_coordinator); + // This thread is the first worker; others start when requests block. + workers.run(false); drop(cancellation); - association.association_ending(); - for worker in workers { + for worker in workers.role.take_started() { if worker.join().is_err() { failure_coordinator.report_panic(IoError::other("broker request worker panicked")); } @@ -560,6 +536,7 @@ where } }); + drop(workers); let result = match failure_coordinator.take_error() { Some(error) => Err(error), None => Ok(()), @@ -577,122 +554,258 @@ where AssociationOutcome { result, abnormal } } -fn read_requests( - request_source: &mut RequestSource, - request_sender: SyncSender, - failure_coordinator: &HostAssociationFailureCoordinator, -) where - RequestSource: HostRequestSource, - Shutdown: HostAssociationShutdown, -{ - loop { - if failure_coordinator.failed() { - break; +/// The workers serving one association's requests. +/// +/// The association's own thread is the first worker; [`ReceivingRole::pass`] +/// starts others only when a request is about to block while no worker is +/// idle. +struct Workers { + association: Arc>, + /// Locked only by the worker holding `role`. + requests: Mutex, + response_sink: ResponseSink, + failure_coordinator: Arc>, + launcher: Option>, + role: ReceivingRole, +} + +/// The right to receive an association's next request. +/// +/// One worker at a time receives, and it keeps the role while it executes the +/// request it received, so the requests of a mostly sequential peer keep going +/// to one worker with a warm cache instead of rotating through ones that have +/// gone cold. A worker passes the role on only when execution is about to wait +/// (see [`before_blocking`]), so a slow request does not hold up later ones. +struct ReceivingRole { + state: Mutex, +} + +struct Receivers { + /// Whether the receiving worker has observed the end of the requests. + ended: bool, + /// Whether a worker holds the receiving role. + held: bool, + /// Workers waiting for the receiving role, most recently idle last. + /// + /// A worker waits only while another one holds the role, so this is empty + /// whenever `held` is false. + idle: Vec, + /// Workers started by [`ReceivingRole::pass`]. + started: Vec>, +} + +/// A receiving role held by the worker executing a request. +trait HeldReceivingRole { + fn pass(self: Arc); +} + +std::thread_local! { + /// The receiving role this thread holds while executing a request. + static HELD_RECEIVING_ROLE: Cell>> = + const { Cell::new(None) }; +} + +/// Lets another worker receive requests while the calling thread waits. +/// +/// Request execution calls this before it waits for anything other than the +/// CPU, such as another thread or process. It does nothing on a thread that is +/// not executing a request while holding its association's receiving role. +pub(crate) fn before_blocking() { + if let Some(role) = HELD_RECEIVING_ROLE.take() { + role.pass(); + } +} + +impl ReceivingRole { + fn new() -> Self { + Self { + state: Mutex::new(Receivers { + ended: false, + held: false, + idle: Vec::with_capacity(crate::WORKER_COUNT), + started: Vec::with_capacity(crate::WORKER_COUNT), + }), } - match request_source.recv_request() { - Ok(HostReceive::Message(request)) => { - if !enqueue_request( - &request_sender, - request, - failure_coordinator, - REQUEST_QUEUE_STALL_TIMEOUT, - ) { - break; - } - } - Ok(HostReceive::ProtocolViolation) => { - failure_coordinator.report(IoError::new( - ErrorKind::InvalidData, - "runner sent a request for the wrong protocol phase", - )); - break; + } + + fn state(&self) -> MutexGuard<'_, Receivers> { + self.state + .lock() + .expect("broker receiving role mutex poisoned") + } + + /// Waits for the role, or returns `false` once requests have ended. + fn take(&self) -> bool { + let mut state = self.state(); + if state.ended { + return false; + } + if !state.held { + state.held = true; + return true; + } + let worker = std::thread::current(); + let id = worker.id(); + state.idle.push(worker); + loop { + if state.ended { + return false; } - Ok(HostReceive::PeerClosed) => break, - Err(error) => { - failure_coordinator.report(error); - break; + // Only a worker passing the role removes a waiting one. + if !state.idle.iter().any(|idle| idle.id() == id) { + return true; } + drop(state); + std::thread::park(); + state = self.state(); } } + + /// Passes the role to the most recently idle worker or, while there are + /// fewer than [`crate::WORKER_COUNT`] workers, to one that `start` starts + /// already holding it given its index. + /// + /// Without either, the role is left for the next worker to finish. + fn pass(&self, start: impl FnOnce(usize) -> IoResult>) { + let mut state = self.state(); + if let Some(next) = state.idle.pop() { + drop(state); + next.unpark(); + return; + } + let index = state.started.len() + 1; + // Starting under the lock records the new worker before it can end + // the requests, so the association thread joins it. + if index < crate::WORKER_COUNT + && let Ok(worker) = start(index) + { + state.started.push(worker); + } else { + state.held = false; + } + } + + /// Ends receiving and returns every idle worker. + fn end(&self) { + let idle = { + let mut state = self.state(); + state.ended = true; + std::mem::take(&mut state.idle) + }; + for worker in idle { + worker.unpark(); + } + } + + /// Takes the workers [`Self::pass`] started. + /// + /// Once requests have ended, no worker holds the role to start another. + fn take_started(&self) -> Vec> { + let mut state = self.state(); + debug_assert!(state.ended, "broker requests must end before joining"); + std::mem::take(&mut state.started) + } } -fn enqueue_request( - request_sender: &SyncSender, - mut request: BrokerRequest, - failure_coordinator: &HostAssociationFailureCoordinator, - stall_timeout: Duration, -) -> bool +impl HeldReceivingRole + for Workers where - Shutdown: HostAssociationShutdown, + Memory: SharedMemory, + RequestSource: HostRequestSource + Send + 'static, + ResponseSink: HostResponseSink + Send + Sync + 'static, + Shutdown: HostAssociationShutdown + Send + Sync + 'static, { - let started = Instant::now(); - loop { - match request_sender.try_send(request) { - Ok(()) => return true, - Err(TrySendError::Disconnected(_)) => { - failure_coordinator.report(IoError::new( - ErrorKind::BrokenPipe, - "broker request workers stopped", - )); - return false; + fn pass(self: Arc) { + self.role.pass(|index| { + let workers = Arc::clone(&self); + std::thread::Builder::new() + .name(format!("litebox-broker-worker-{index}")) + .spawn(move || workers.run(true)) + }); + } +} + +impl + Workers +where + Memory: SharedMemory, + RequestSource: HostRequestSource + Send + 'static, + ResponseSink: HostResponseSink + Send + Sync + 'static, + Shutdown: HostAssociationShutdown + Send + Sync + 'static, +{ + /// Serves requests until the association's requests end, starting with + /// the receiving role if `holds_role`. + /// + /// Every worker takes its turn receiving and executes each request it + /// receives. Once every worker is busy, requests wait in the transport + /// until one finishes. + fn run(self: &Arc, mut holds_role: bool) { + loop { + if !holds_role && !self.role.take() { + return; } - Err(TrySendError::Full(pending)) => { - request = pending; - if failure_coordinator.failed() { - return false; - } - if started.elapsed() >= stall_timeout { - failure_coordinator.report(IoError::new( - ErrorKind::TimedOut, - "broker request queue remained full", - )); - return false; + let Some(request) = self.next_request() else { + return; + }; + HELD_RECEIVING_ROLE.set(Some(Arc::clone(self) as Arc)); + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.association.execute_request_with( + request, + |process, operation, shared_buffers| { + self.launcher.as_ref().and_then(|launcher| { + handle_process_operation(launcher, process, operation, shared_buffers) + }) + }, + |response| self.response_sink.send_response(response), + ) + })) { + Ok(Ok(()) | Err(BrokerHostError::AssociationFailed)) => {} + Ok(Err(error)) => self.failure_coordinator.report(map_host_error(error)), + Err(_) => { + self.failure_coordinator + .report_panic(IoError::other("broker request worker panicked")); } - std::thread::sleep(REQUEST_QUEUE_RETRY_DELAY); } + holds_role = HELD_RECEIVING_ROLE.take().is_some(); } } -} -fn run_worker( - association: &BrokerHostAssociation, - request_receiver: &Mutex>, - response_sink: &ResponseSink, - failure_coordinator: &HostAssociationFailureCoordinator, - launcher: Option<&Arc>, -) where - Memory: SharedMemory, - ResponseSink: HostResponseSink, - Shutdown: HostAssociationShutdown, -{ - loop { - let request = request_receiver - .lock() - .expect("broker request receiver mutex poisoned") - .recv(); - let Ok(request) = request else { - break; - }; - if failure_coordinator.failed() { - continue; + /// Receives the next request while holding the receiving role, or returns + /// `None` once the association's requests have ended. + /// + /// The worker that observes the end reports why and starts association + /// teardown, so workers still executing requests see cancellation + /// promptly. + fn next_request(&self) -> Option { + let received = (!self.failure_coordinator.failed()).then(|| { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.requests + .lock() + .expect("broker request source mutex poisoned") + .recv_request() + })) + }); + if let Some(Ok(Ok(HostReceive::Message(request)))) = received { + return Some(request); } - match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - association.execute_request_with( - request, - |process, operation, shared_buffers| { - launcher.and_then(|launcher| { - handle_process_operation(launcher, process, operation, shared_buffers) - }) - }, - |response| response_sink.send_response(response), - ) - })) { - Ok(Ok(()) | Err(BrokerHostError::AssociationFailed)) => {} - Ok(Err(error)) => failure_coordinator.report(map_host_error(error)), - Err(_) => { - failure_coordinator.report_panic(IoError::other("broker request worker panicked")); + self.role.end(); + match received { + Some(Ok(Ok(HostReceive::ProtocolViolation))) => { + self.failure_coordinator.report(IoError::new( + ErrorKind::InvalidData, + "runner sent a request for the wrong protocol phase", + )); + } + Some(Ok(Err(error))) => self.failure_coordinator.report(error), + Some(Err(_)) => { + self.failure_coordinator + .report_panic(IoError::other("broker request reader panicked")); } + Some(Ok(Ok(HostReceive::PeerClosed | HostReceive::Message(_)))) | None => {} } + self.association.request_cancellation(); + self.association.association_ending(); + None } } @@ -709,13 +822,8 @@ mod tests { use litebox_broker_core::{ObjectRights, PolicyEngine}; use litebox_broker_host::setup_connection; use litebox_broker_protocol::BROKER_PROTOCOL_VERSION; - use litebox_broker_protocol::RequestId; - use litebox_broker_protocol::message::{ - BrokerHandshakeResponse, BrokerNotification, BrokerOperation, - }; - use litebox_broker_protocol::shared_buffer::{ - SHARED_BUFFER_LAYOUT, SHARED_BUFFER_POOL_SIZE, SharedBufferSequence, SharedBufferSlotIndex, - }; + use litebox_broker_protocol::message::{BrokerHandshakeResponse, BrokerNotification}; + use litebox_broker_protocol::shared_buffer::{SHARED_BUFFER_LAYOUT, SHARED_BUFFER_POOL_SIZE}; use litebox_broker_transport::channel::{ HostNotificationChannel, HostReceive, HostSetupChannel, LocalSetupChannel, }; @@ -934,17 +1042,6 @@ mod tests { (local, notifications, shutdown, outcome, host) } - struct RecordingShutdown(Arc); - - impl HostAssociationShutdown for RecordingShutdown { - type Error = IoError; - - fn shutdown(&self) -> IoResult<()> { - self.0.store(true, Ordering::Release); - Ok(()) - } - } - #[test] fn publication_guard_ends_a_parked_publisher() { let association = live_association(); @@ -1100,34 +1197,6 @@ mod tests { ); } - #[test] - fn a_stalled_request_queue_fails_after_its_deadline() { - let request = |request_id| BrokerRequest { - request_id: RequestId(request_id), - operation: BrokerOperation::FillRandom( - SharedBufferSequence::new(&[SharedBufferSlotIndex(0)], 1).unwrap(), - ), - }; - let shutdown_called = Arc::new(AtomicBool::new(false)); - let failure_coordinator = - HostAssociationFailureCoordinator::new(RecordingShutdown(Arc::clone(&shutdown_called))); - let (request_sender, _request_receiver) = sync_channel(1); - request_sender.send(request(1)).unwrap(); - - assert!(!enqueue_request( - &request_sender, - request(2), - &failure_coordinator, - Duration::ZERO, - )); - - assert!(shutdown_called.load(Ordering::Acquire)); - assert_eq!( - failure_coordinator.take_error().unwrap().kind(), - ErrorKind::TimedOut - ); - } - #[test] fn dispatching_requests_publishes_readiness_until_the_association_ends() { use litebox_broker_protocol::ObjectHandle; @@ -1194,4 +1263,107 @@ mod tests { drop(local); host.join().unwrap(); } + + fn wait_for_idle_workers(role: &ReceivingRole, count: usize) { + let deadline = Instant::now() + TEST_SETUP_TIMEOUT; + while role.state().idle.len() < count { + assert!( + Instant::now() < deadline, + "workers did not wait for the receiving role" + ); + std::thread::yield_now(); + } + } + + fn start_no_worker(_index: usize) -> IoResult> { + Err(IoError::other("this test starts no workers")) + } + + #[test] + fn the_receiving_role_passes_to_the_most_recently_idle_worker() { + let role = ReceivingRole::new(); + assert!(role.take()); + let (taken, order) = sync_channel(2); + std::thread::scope(|scope| { + for worker in 1..=2 { + let taken = taken.clone(); + let role = &role; + scope.spawn(move || { + assert!(role.take()); + taken.send(worker).unwrap(); + }); + wait_for_idle_workers(role, worker); + } + let start = |_| panic!("an idle worker must take the role"); + role.pass(start); + assert_eq!(order.recv_timeout(TEST_SETUP_TIMEOUT).unwrap(), 2); + role.pass(start); + assert_eq!(order.recv_timeout(TEST_SETUP_TIMEOUT).unwrap(), 1); + }); + // Without a worker to take it, the role is left for the next taker. + role.pass(start_no_worker); + assert!(role.take()); + } + + #[test] + fn passing_the_receiving_role_starts_workers_up_to_the_limit() { + let role = ReceivingRole::new(); + assert!(role.take()); + for expected in 1..crate::WORKER_COUNT { + role.pass(|index| { + assert_eq!(index, expected); + Ok(std::thread::spawn(|| {})) + }); + assert!(role.state().held, "a started worker holds the role"); + } + role.pass(|_| panic!("no worker may start beyond the limit")); + assert!(!role.state().held); + + role.end(); + let started = role.take_started(); + assert_eq!(started.len(), crate::WORKER_COUNT - 1); + for worker in started { + worker.join().unwrap(); + } + } + + #[test] + fn ending_the_receiving_role_returns_every_idle_worker() { + let role = ReceivingRole::new(); + assert!(role.take()); + std::thread::scope(|scope| { + let workers: Vec<_> = (1..=2) + .map(|count| { + let worker = scope.spawn(|| role.take()); + wait_for_idle_workers(&role, count); + worker + }) + .collect(); + role.end(); + for worker in workers { + assert!(!worker.join().unwrap()); + } + }); + assert!(!role.take()); + } + + #[test] + fn blocking_passes_only_a_held_receiving_role() { + struct Held(ReceivingRole); + impl HeldReceivingRole for Held { + fn pass(self: Arc) { + self.0.pass(start_no_worker); + } + } + + let held = Arc::new(Held(ReceivingRole::new())); + assert!(held.0.take()); + before_blocking(); + assert!(held.0.state().held); + + HELD_RECEIVING_ROLE.set(Some(Arc::clone(&held) as Arc)); + before_blocking(); + assert!(!held.0.state().held); + assert!(HELD_RECEIVING_ROLE.take().is_none()); + } } diff --git a/litebox_broker_userland/src/windows.rs b/litebox_broker_userland/src/windows.rs index a9c84a301e..9a0ed9af70 100644 --- a/litebox_broker_userland/src/windows.rs +++ b/litebox_broker_userland/src/windows.rs @@ -20,9 +20,10 @@ use litebox_broker_transport_windows_userland::named_pipe::{ }; use litebox_broker_transport_windows_userland::shared_memory::WindowsSharedMemory; use litebox_broker_userland::builder::BrokerCoreBuilder; +use litebox_broker_userland::runner::accept_runner_channel; use litebox_broker_userland::stdio::UserlandStdioProvider; -use super::{SETUP_TIMEOUT, configured_socket_policy}; +use super::{IN_PROCESS_RUNNER_POLL_INTERVAL, SETUP_TIMEOUT, configured_socket_policy}; pub(super) fn run( args: super::CliArgs, @@ -74,11 +75,15 @@ fn serve_runner_in_process( runner: &JoinHandle, ) -> IoResult<()> { let setup_deadline = Instant::now() + SETUP_TIMEOUT; - let control_stream = crate::accept_runner_channel( + let control_stream = accept_runner_channel( setup_deadline, "control", || Ok(runner.is_finished().then(|| "thread stopped".to_owned())), || control_listener.try_accept(), + |remaining| { + std::thread::sleep(remaining.min(IN_PROCESS_RUNNER_POLL_INTERVAL)); + Ok(()) + }, )?; validate_client_process(&control_stream, std::process::id())?; serve_control_stream( From cfce8cdd08ce35792cb0159077c375d052e5b2c4 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 3 Oct 2026 23:18:19 -0700 Subject: [PATCH 2/4] Cut page-table costs of copying large files - Add a defaulted `PageManagementProvider::advise_fill` hint, which the Linux shim gives before copying file contents into a new mapping. The Linux userland platform backs such ranges with transparent huge pages, so the copy takes one fault and clear per 2 MiB instead of per page, while pages past the end of the file stay lazily allocated. - Read tar file contents through `TarStorage::read_exact_at`, which the userland broker's mapped archive serves with positional reads, so the broker no longer maps every page it serves and tears them down at exit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b5ed1835-8e32-4a35-aae1-8d8a8a276d09 --- litebox/src/platform/page_mgmt.rs | 10 +++ litebox_broker_core/src/fs/tar_ro.rs | 50 ++++++++----- litebox_broker_userland/src/mapped_file.rs | 70 +++++++++++++++---- .../src/page_mgmt.rs | 15 ++++ litebox_shim_linux/src/syscalls/mm.rs | 5 ++ 5 files changed, 118 insertions(+), 32 deletions(-) diff --git a/litebox/src/platform/page_mgmt.rs b/litebox/src/platform/page_mgmt.rs index c08acf7d7a..c55365ab22 100644 --- a/litebox/src/platform/page_mgmt.rs +++ b/litebox/src/platform/page_mgmt.rs @@ -220,6 +220,16 @@ pub trait PageManagementProvider: RawPointerProvider { /// Note that the returned ranges should be `ALIGN`-aligned. fn reserved_pages(&self) -> impl Iterator>; + /// Hints that the caller is about to fill the allocated pages in `range` with data from their + /// start, possibly stopping short of the end (e.g., at the end of a file copied into them). + /// + /// A platform may, for example, back these pages with larger host pages, which take fewer + /// faults to fill but can leave up to one larger page past the filled part resident. + /// + /// The default implementation does nothing. + #[expect(unused_variables, reason = "default body")] + fn advise_fill(&self, range: Range) {} + /// Attempt to allocate pages with copy-on-write semantics backed by static data. /// /// This method allows platforms that support it to create CoW mappings instead of performing diff --git a/litebox_broker_core/src/fs/tar_ro.rs b/litebox_broker_core/src/fs/tar_ro.rs index bd61ceee1e..f1a8e5df82 100644 --- a/litebox_broker_core/src/fs/tar_ro.rs +++ b/litebox_broker_core/src/fs/tar_ro.rs @@ -52,19 +52,34 @@ pub struct TarRo { } /// The storage holding a [`TarRo`] archive. -type TarData = Box + Send + Sync>; +/// +/// Anything that can lend out the archive's bytes, such as a static slice or an owned buffer, is +/// such storage. +pub trait TarStorage: Send + Sync + 'static { + /// Lends out the whole archive, which must be the same bytes every time. + fn bytes(&self) -> &[u8]; + + /// Fills `buf` with the archive's bytes starting at `offset`, which lie within [`Self::bytes`]. + /// + /// [`TarRo`] indexes the archive through [`Self::bytes`] but reads file contents through this, + /// so storage that can read its bytes more cheaply than it lends them out can override it. + fn read_exact_at(&self, buf: &mut [u8], offset: usize) -> Result<(), ReadError> { + buf.copy_from_slice(&self.bytes()[offset..][..buf.len()]); + Ok(()) + } +} + +impl + Send + Sync + 'static> TarStorage for T { + fn bytes(&self) -> &[u8] { + self.as_ref() + } +} impl TarRo { /// Construct a tar backend over the archive in `tar_data`, using a caller-provided inode /// allocator. - /// - /// `tar_data` is any storage that can lend out the archive's bytes, such as a static slice, an - /// owned buffer, or a memory-mapped file; it must lend out the same bytes every time. #[must_use] - pub fn new( - tar_data: impl AsRef<[u8]> + Send + Sync + 'static, - inode_allocator: InodeAllocator, - ) -> Self { + pub fn new(tar_data: impl TarStorage, inode_allocator: InodeAllocator) -> Self { Self { tar_index: TarIndex::new(Box::new(tar_data), inode_allocator), } @@ -201,12 +216,16 @@ impl super::backend::Backend for TarRo { } fn read(&self, h: &FileHandle, buf: &mut [u8], offset: usize) -> Result { - let file = self.tar_index.file_data(h.get_typed::().idx); + let file = self.tar_index.files[h.get_typed::().idx] + .data_range + .clone(); let start = offset.min(file.len()); let end = offset.checked_add(buf.len()).unwrap().min(file.len()); debug_assert!(start <= end); let len = end - start; - buf[..len].copy_from_slice(&file[start..end]); + self.tar_index + .tar_data + .read_exact_at(&mut buf[..len], file.start + start)?; Ok(len) } @@ -326,14 +345,14 @@ enum IndexedChild { } struct TarIndex { - tar_data: TarData, + tar_data: Box, files: Vec, dirs: Vec, } impl TarIndex { - fn new(tar_data: TarData, inode_allocator: InodeAllocator) -> Self { - let bytes = (*tar_data).as_ref(); + fn new(tar_data: Box, inode_allocator: InodeAllocator) -> Self { + let bytes = tar_data.bytes(); let archive = tar_no_std::TarArchiveRef::new(bytes).expect("invalid tar data"); let base_ptr = bytes.as_ptr() as usize; @@ -422,11 +441,6 @@ impl TarIndex { } } - fn file_data(&self, file_idx: usize) -> &[u8] { - let range = self.files[file_idx].data_range.clone(); - &(*self.tar_data).as_ref()[range] - } - /// The directory entry for `child`, named `name` in its parent. fn dir_entry(&self, name: &str, child: IndexedChild) -> DirEntry { let (file_type, node_info) = match child { diff --git a/litebox_broker_userland/src/mapped_file.rs b/litebox_broker_userland/src/mapped_file.rs index daa2d16b02..07f8085ce3 100644 --- a/litebox_broker_userland/src/mapped_file.rs +++ b/litebox_broker_userland/src/mapped_file.rs @@ -8,12 +8,17 @@ use std::io::{Error as IoError, ErrorKind, Result as IoResult}; use std::path::Path; use std::ptr::NonNull; +use litebox_broker_core::fs::errors::ReadError; +use litebox_broker_core::fs::tar_ro::TarStorage; + /// The contents of a file, mapped read-only into the broker's address space. /// /// This lets the broker serve a large read-only input, such as the initial file system archive, /// without first copying all of it into memory: pages come from the host's page cache as they are -/// read. +/// read. Reads at an offset copy from the page cache directly, sparing the broker the cost of +/// mapping those pages in and tearing them down again. pub struct MappedFile { + file: File, /// Start of the mapping, or dangling if `len` is zero. address: NonNull, len: usize, @@ -37,25 +42,27 @@ impl MappedFile { let file = File::open(path)?; let len = usize::try_from(file.metadata()?.len()) .map_err(|_| IoError::new(ErrorKind::InvalidInput, "file is too large to map"))?; - if len == 0 { + let address = if len == 0 { // Hosts refuse to map empty files, and there are no bytes to lend out anyway. - return Ok(Self { - address: NonNull::dangling(), - len, - }); - } - let address = host::map(&file, len)?; - Ok(Self { address, len }) + NonNull::dangling() + } else { + host::map(&file, len)? + }; + Ok(Self { file, address, len }) } } -impl AsRef<[u8]> for MappedFile { - fn as_ref(&self) -> &[u8] { +impl TarStorage for MappedFile { + fn bytes(&self) -> &[u8] { // SAFETY: `address` is either dangling with a zero `len`, or the start of a live, readable // mapping of `len` bytes owned by `self`, whose contents the contract of `open` keeps // unchanged. unsafe { std::slice::from_raw_parts(self.address.as_ptr(), self.len) } } + + fn read_exact_at(&self, buf: &mut [u8], offset: usize) -> Result<(), ReadError> { + host::read_exact_at(&self.file, buf, offset as u64).map_err(|_| ReadError::Io) + } } impl Drop for MappedFile { @@ -73,8 +80,14 @@ mod host { use std::fs::File; use std::io::{Error as IoError, Result as IoResult}; use std::os::fd::AsRawFd; + use std::os::unix::fs::FileExt as _; use std::ptr::NonNull; + /// Fill `buf` with the bytes of `file` at `offset`. + pub(super) fn read_exact_at(file: &File, buf: &mut [u8], offset: u64) -> IoResult<()> { + file.read_exact_at(buf, offset) + } + /// Map the first `len` bytes of `file` read-only. pub(super) fn map(file: &File, len: usize) -> IoResult> { // SAFETY: A new private mapping at an address chosen by the kernel aliases no existing @@ -110,10 +123,27 @@ mod host { #[cfg(all(windows, target_arch = "x86_64"))] mod host { use std::fs::File; - use std::io::{Error as IoError, Result as IoResult}; + use std::io::{Error as IoError, ErrorKind, Result as IoResult}; + use std::os::windows::fs::FileExt as _; use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; use std::ptr::NonNull; + /// Fill `buf` with the bytes of `file` at `offset`. + pub(super) fn read_exact_at(file: &File, mut buf: &mut [u8], mut offset: u64) -> IoResult<()> { + while !buf.is_empty() { + match file.seek_read(buf, offset) { + Ok(0) => return Err(ErrorKind::UnexpectedEof.into()), + Ok(read) => { + buf = &mut buf[read..]; + offset += read as u64; + } + Err(error) if error.kind() == ErrorKind::Interrupted => {} + Err(error) => return Err(error), + } + } + Ok(()) + } + use windows_sys::Win32::System::Memory::{ CreateFileMappingW, FILE_MAP_READ, MEMORY_MAPPED_VIEW_ADDRESS, MapViewOfFile, PAGE_READONLY, UnmapViewOfFile, @@ -163,6 +193,8 @@ mod host { mod tests { use std::io::Write; + use litebox_broker_core::fs::tar_ro::TarStorage as _; + use super::MappedFile; fn mapped(contents: &[u8]) -> MappedFile { @@ -175,11 +207,21 @@ mod tests { #[test] fn maps_file_contents() { let contents: Vec = (0..10_000u32).flat_map(u32::to_le_bytes).collect(); - assert_eq!(mapped(&contents).as_ref(), contents); + assert_eq!(mapped(&contents).bytes(), contents); + } + + #[test] + fn reads_file_contents_at_an_offset() { + let contents: Vec = (0..10_000u32).flat_map(u32::to_le_bytes).collect(); + let mut buf = [0; 100]; + mapped(&contents).read_exact_at(&mut buf, 5000).unwrap(); + assert_eq!(buf, contents[5000..5100]); } #[test] fn maps_empty_file() { - assert!(mapped(b"").as_ref().is_empty()); + let file = mapped(b""); + assert!(file.bytes().is_empty()); + file.read_exact_at(&mut [], 0).unwrap(); } } diff --git a/litebox_platform_linux_userland/src/page_mgmt.rs b/litebox_platform_linux_userland/src/page_mgmt.rs index f07afaa7c5..6fc9ec9b47 100644 --- a/litebox_platform_linux_userland/src/page_mgmt.rs +++ b/litebox_platform_linux_userland/src/page_mgmt.rs @@ -272,6 +272,21 @@ impl litebox::platform::PageManagementProvider for Li self.reserved_pages.iter() } + /// Asks the host to back `range` with transparent huge pages, each of which takes one + /// fault and one clear instead of one per base page. This is best-effort: the host may + /// have them disabled, or another thread may have unmapped the range. + fn advise_fill(&self, range: core::ops::Range) { + // SAFETY: `MADV_HUGEPAGE` changes how the host backs `range`, never what it contains. + let _ = unsafe { + syscalls::syscall3( + syscalls::Sysno::madvise, + range.start, + range.len(), + libc::MADV_HUGEPAGE.reinterpret_as_unsigned() as usize, + ) + }; + } + fn try_allocate_cow_pages( &self, suggested_start: usize, diff --git a/litebox_shim_linux/src/syscalls/mm.rs b/litebox_shim_linux/src/syscalls/mm.rs index cdc9bf8142..d5c3f4b1d5 100644 --- a/litebox_shim_linux/src/syscalls/mm.rs +++ b/litebox_shim_linux/src/syscalls/mm.rs @@ -378,6 +378,11 @@ impl Task { offset: usize, ) -> Result, MappingError> { let op = |ptr: UserPtrMut| -> Result { + // The copy below fills the mapping from its start up to the end of the file. + >::advise_fill( + self.global.platform, + ptr.as_usize()..ptr.as_usize() + len, + ); // Note a malicious user may unmap ptr while we are reading. // `sys_read` does not handle page faults, so we need to use a // temporary buffer to read the data from fs (without worrying page From 715e51ad7ce51a2c3b757b8803a2f44700ebfa13 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 3 Oct 2026 23:35:12 -0700 Subject: [PATCH 3/4] Keep overlay reads consistent and local closes quiet Read each mutable upper directory in a single listing so a lookup's entry, whiteout, and opaque marker come from one snapshot, list each layer once when enumerating a directory, and end notification receipt cleanly when an acknowledgement fails because the local call channel was closed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b5ed1835-8e32-4a35-aae1-8d8a8a276d09 --- litebox_broker_core/src/fs/overlay.rs | 194 ++++++++++-------- .../src/unix_socket/local.rs | 8 +- .../src/local.rs | 8 +- 3 files changed, 118 insertions(+), 92 deletions(-) diff --git a/litebox_broker_core/src/fs/overlay.rs b/litebox_broker_core/src/fs/overlay.rs index 80d587b2bf..068db0860e 100644 --- a/litebox_broker_core/src/fs/overlay.rs +++ b/litebox_broker_core/src/fs/overlay.rs @@ -14,7 +14,7 @@ //! changes do not compromise memory safety or the structural integrity of its internal state. use alloc::boxed::Box; -use alloc::collections::BTreeSet; +use alloc::collections::{BTreeMap, BTreeSet}; use alloc::string::String; use alloc::sync::Arc; use alloc::vec; @@ -112,13 +112,11 @@ enum OverlayFileLayer { /// A logical directory, resolved to the per-layer directories that make it up. /// /// Its entries are looked up by name on demand (see [`Overlay::entry`]), so resolving a path never -/// reads whole directories. +/// reads whole lower directories. struct ResolvedDir { upper: Option, /// Per lower backend, in precedence order, its directory at this path if any. lowers: Vec>, - /// Whether the upper directory hides every lower entry. - opaque: bool, } /// An overlay-visible directory entry, plus which layers contribute to it. @@ -132,6 +130,10 @@ struct ResolvedEntry { lower_directories: Vec, } +/// One name's entries in the layers of a directory: its upper entry, and its lower entries (by +/// lower backend, in precedence order) that the upper directory does not hide. +type LayerEntries = (Option, Vec<(usize, DirEntry)>); + impl Overlay { /// Construct an overlay over a single `lower`, using `allocator` for overlay-visible inodes. pub fn new(upper: impl Backend, lower: impl Backend, allocator: InodeAllocator) -> Self { @@ -171,25 +173,9 @@ impl Overlay { .iter() .map(|lower| lower.owned_dir_at(lower.root(), OFlags::PATH).map(Some)) .collect::, _>>()?; - self.resolved_dir(Some(upper), lowers) - } - - /// The logical directory made up of the given per-layer directories. - fn resolved_dir( - &self, - upper: Option, - lowers: Vec>, - ) -> Result { - let opaque = match &upper { - Some(upper) => self - .marker_present(upper, OPAQUE_MARKER) - .map_err(|_| OpenError::Io)?, - None => false, - }; Ok(ResolvedDir { - upper, + upper: Some(upper), lowers, - opaque, }) } @@ -255,7 +241,7 @@ impl Overlay { let component = owner_component.expect("a merged directory is owned by upper or by a lower directory"); - Ok((self.resolved_dir(upper, lowers)?, component)) + Ok((ResolvedDir { upper, lowers }, component)) } /// Resolve a logical `path` (relative to the overlay root) to its per-layer directories. @@ -597,52 +583,113 @@ impl Overlay { /// The overlay-visible entry `name` of `dir`, if any, merged across the layers that hold it. fn entry(&self, dir: &ResolvedDir, name: &str) -> Result, OpenError> { - fn lookup( - backend: &dyn Backend, - dir: &DirHandle, - name: &str, - ) -> Result, OpenError> { - backend.lookup_at(dir, name).map_err(|_| OpenError::Io) - } if !valid(name) { return Ok(None); } + let mut upper = None; + let mut lowers = Vec::new(); + if let Some(handle) = &dir.upper { + // Read paths do not take the namespace lock, so read the mutable upper directory in a + // single listing: separate lookups of `name` and of the markers hiding it could + // straddle a concurrent writer and expose deleted lower entries. + let whiteout = whiteout(name); + let mut hidden = false; + for entry in self + .upper + .list_dir_at(handle.clone()) + .map_err(|_| OpenError::Io)? + { + if entry.name == name { + upper = Some(entry); + } else if entry.name == whiteout || entry.name == OPAQUE_MARKER { + // Markers held by the upper directory say what it hides from the lowers. + hidden = true; + } + } + if hidden { + return Ok(self.merge_entry(upper, lowers)); + } + } + // Lower backends are immutable, so they can be looked up by name. + for (layer, (backend, handle)) in self.lowers.iter().zip(&dir.lowers).enumerate() { + if let Some(handle) = handle + && let Some(entry) = backend.lookup_at(handle, name).map_err(|_| OpenError::Io)? + { + lowers.push((layer, entry)); + } + } + Ok(self.merge_entry(upper, lowers)) + } + + /// Every overlay-visible entry of `dir`, sorted by name. + fn entries(&self, dir: &ResolvedDir) -> Result, OpenError> { + fn list(backend: &dyn Backend, dir: &DirHandle) -> Result, OpenError> { + backend.list_dir_at(dir.clone()).map_err(|_| OpenError::Io) + } + + let mut names: BTreeMap = BTreeMap::new(); + // Markers held by the upper directory say what it hides from the lowers. + let mut markers = BTreeSet::new(); + if let Some(handle) = &dir.upper { + for entry in list(self.upper.as_ref(), handle)? { + if valid(&entry.name) { + let name = entry.name.clone(); + names.entry(name).or_default().0 = Some(entry); + } else { + markers.insert(entry.name); + } + } + } + if !markers.contains(OPAQUE_MARKER) { + for (layer, (backend, handle)) in self.lowers.iter().zip(&dir.lowers).enumerate() { + let Some(handle) = handle else { + continue; + }; + for entry in list(backend.as_ref(), handle)? { + if valid(&entry.name) && !markers.contains(&whiteout(&entry.name)) { + names + .entry(entry.name.clone()) + .or_default() + .1 + .push((layer, entry)); + } + } + } + } + Ok(names + .into_values() + .filter_map(|(upper, lowers)| self.merge_entry(upper, lowers)) + .map(|merged| merged.entry) + .collect()) + } + + /// Merge the entries of one name in one logical directory: `upper` from its upper directory, + /// and `lowers` (in precedence order) from the lower directories it does not hide. + fn merge_entry( + &self, + upper: Option, + lowers: Vec<(usize, DirEntry)>, + ) -> Option { let mut merged: Option = None; // Whether lower entries can no longer be merged in: the entry is not a directory in every // layer that contributed to it. let mut blocked = false; - if let Some(upper) = &dir.upper { - if let Some(mut entry) = lookup(self.upper.as_ref(), upper, name)? { - blocked = entry.file_type != FileType::Directory; - entry.ino_info = entry - .ino_info - .take() - .map(|node| self.map_node(&mut self.state.lock().ids, None, node)); - merged = Some(ResolvedEntry { - entry, - upper: true, - lower: None, - lower_directories: vec![false; self.lowers.len()], - }); - } - // Markers held by the upper directory say what it hides from the lowers. - if dir.opaque - || self - .marker_present(upper, &whiteout(name)) - .map_err(|_| OpenError::Io)? - { - return Ok(merged); - } + if let Some(mut entry) = upper { + blocked = entry.file_type != FileType::Directory; + entry.ino_info = entry + .ino_info + .take() + .map(|node| self.map_node(&mut self.state.lock().ids, None, node)); + merged = Some(ResolvedEntry { + entry, + upper: true, + lower: None, + lower_directories: vec![false; self.lowers.len()], + }); } - for (layer, handle) in dir.lowers.iter().enumerate() { - let Some(handle) = handle else { - continue; - }; - let Some(mut lower_entry) = lookup(self.lowers[layer].as_ref(), handle, name)? else { - continue; - }; + for (layer, mut lower_entry) in lowers { let directory = lower_entry.file_type == FileType::Directory; let lower_node = lower_entry.ino_info.take(); let entry = merged.get_or_insert_with(|| ResolvedEntry { @@ -675,36 +722,7 @@ impl Overlay { blocked = true; } } - Ok(merged) - } - - /// Every overlay-visible entry of `dir`, sorted by name. - fn entries(&self, dir: &ResolvedDir) -> Result, OpenError> { - let mut names = BTreeSet::new(); - let mut add_names = |backend: &dyn Backend, dir: &DirHandle| { - let entries = backend - .list_dir_at(dir.clone()) - .map_err(|_| OpenError::Io)?; - names.extend(entries.into_iter().map(|entry| entry.name)); - Ok::<_, OpenError>(()) - }; - if let Some(upper) = &dir.upper { - add_names(self.upper.as_ref(), upper)?; - } - if !dir.opaque { - for (lower, handle) in self.lowers.iter().zip(&dir.lowers) { - if let Some(handle) = handle { - add_names(lower.as_ref(), handle)?; - } - } - } - let mut entries = Vec::with_capacity(names.len()); - for name in names { - if let Some(entry) = self.entry(dir, &name)? { - entries.push(entry.entry); - } - } - Ok(entries) + merged } } diff --git a/litebox_broker_transport_linux_userland/src/unix_socket/local.rs b/litebox_broker_transport_linux_userland/src/unix_socket/local.rs index 486fa544c6..a26f898b44 100644 --- a/litebox_broker_transport_linux_userland/src/unix_socket/local.rs +++ b/litebox_broker_transport_linux_userland/src/unix_socket/local.rs @@ -381,8 +381,12 @@ impl LocalNotificationChannel for UnixControlRingLocalNotificationChannel { } match self.consumer.try_read(decode_notification) { Ok(ControlRingReadStatus::Message(notification)) => { - self.association - .acknowledge_notification(&mut self.consumer)?; + if let Err(error) = self + .association + .acknowledge_notification(&mut self.consumer) + { + return self.association.end_notifications(&error); + } return Ok(Some(notification)); } Ok(ControlRingReadStatus::Empty { wait_epoch }) => { diff --git a/litebox_broker_transport_windows_userland/src/local.rs b/litebox_broker_transport_windows_userland/src/local.rs index 6ce4737fdc..a1fca4b6b1 100644 --- a/litebox_broker_transport_windows_userland/src/local.rs +++ b/litebox_broker_transport_windows_userland/src/local.rs @@ -291,8 +291,12 @@ impl LocalNotificationChannel for WindowsControlRingLocalNotificationChannel { } match self.consumer.try_read(decode_notification) { Ok(ControlRingReadStatus::Message(notification)) => { - self.association - .acknowledge_notification(&mut self.consumer)?; + if let Err(error) = self + .association + .acknowledge_notification(&mut self.consumer) + { + return self.association.end_notifications(&error); + } return Ok(Some(notification)); } Ok(ControlRingReadStatus::Empty { wait_epoch }) => { From 84ab6f98d07d144ae5cf451087468e4c9ef046ab Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sun, 4 Oct 2026 06:12:01 -0700 Subject: [PATCH 4/4] Drop in-process runner accept polling The in-process runner is for development and testing only, so its startup does not need optimizing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b5ed1835-8e32-4a35-aae1-8d8a8a276d09 --- litebox_broker_userland/src/linux.rs | 11 ++----- litebox_broker_userland/src/main.rs | 35 ++++++++++++++++++--- litebox_broker_userland/src/runner.rs | 18 +---------- litebox_broker_userland/src/runner/linux.rs | 21 +++++-------- litebox_broker_userland/src/windows.rs | 9 ++---- 5 files changed, 43 insertions(+), 51 deletions(-) diff --git a/litebox_broker_userland/src/linux.rs b/litebox_broker_userland/src/linux.rs index c44f2ecd87..770aedfcee 100644 --- a/litebox_broker_userland/src/linux.rs +++ b/litebox_broker_userland/src/linux.rs @@ -22,10 +22,9 @@ use litebox_broker_transport_linux_userland::unix_socket::{ UnixStreamHostSetupChannel, validate_peer_process, }; use litebox_broker_userland::builder::BrokerCoreBuilder; -use litebox_broker_userland::runner::{accept_runner_channel, wait_for_connection}; use litebox_broker_userland::stdio::UserlandStdioProvider; -use super::{IN_PROCESS_RUNNER_POLL_INTERVAL, SETUP_TIMEOUT, configured_socket_policy}; +use super::{SETUP_TIMEOUT, configured_socket_policy}; const PROXY_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); @@ -215,17 +214,11 @@ fn serve_runner_in_process( runner: &JoinHandle, ) -> IoResult<()> { let setup_deadline = Instant::now() + SETUP_TIMEOUT; - let control_stream = accept_runner_channel( + let control_stream = crate::accept_runner_channel( setup_deadline, "control", || Ok(runner.is_finished().then(|| "thread stopped".to_owned())), || control_listener.accept().map(|(stream, _)| stream), - |remaining| { - wait_for_connection( - control_listener, - remaining.min(IN_PROCESS_RUNNER_POLL_INTERVAL), - ) - }, )?; validate_peer_process(&control_stream, std::process::id())?; serve_control_stream(broker, control_stream, setup_deadline) diff --git a/litebox_broker_userland/src/main.rs b/litebox_broker_userland/src/main.rs index 736e98e70c..c66f3ef66a 100644 --- a/litebox_broker_userland/src/main.rs +++ b/litebox_broker_userland/src/main.rs @@ -9,7 +9,7 @@ use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; use std::thread::JoinHandle; -use std::time::Duration; +use std::time::{Duration, Instant}; use clap::Parser; #[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] @@ -37,9 +37,7 @@ mod linux; mod windows; const SETUP_TIMEOUT: Duration = Duration::from_secs(5); -/// How often an in-process runner thread, which cannot be waited on alongside -/// its connection, is rechecked for having stopped before connecting. -const IN_PROCESS_RUNNER_POLL_INTERVAL: Duration = Duration::from_millis(10); +const ACCEPT_RETRY_DELAY: Duration = Duration::from_millis(10); #[derive(Clone, Copy, Debug, PartialEq, Eq)] struct AllowedDestination { @@ -309,6 +307,35 @@ fn destination_rules(allowed_destinations: &[AllowedDestination]) -> Vec( + deadline: Instant, + channel_name: &'static str, + mut runner_status: impl FnMut() -> IoResult>, + mut try_accept: impl FnMut() -> IoResult, +) -> IoResult { + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(IoError::new( + ErrorKind::TimedOut, + format!("timed out waiting for runner {channel_name} channel"), + )); + } + if let Some(status) = runner_status()? { + return Err(IoError::new( + ErrorKind::BrokenPipe, + format!("runner {status} before connecting its {channel_name} channel"), + )); + } + match try_accept() { + Ok(channel) => return Ok(channel), + Err(error) if error.kind() == ErrorKind::WouldBlock => {} + Err(error) => return Err(error), + } + std::thread::sleep(remaining.min(ACCEPT_RETRY_DELAY)); + } +} + #[cfg(any(target_os = "linux", all(windows, target_arch = "x86_64")))] fn main() -> Result<(), Box> { let args = CliArgs::parse(); diff --git a/litebox_broker_userland/src/runner.rs b/litebox_broker_userland/src/runner.rs index 17bddae1de..9f05c1d485 100644 --- a/litebox_broker_userland/src/runner.rs +++ b/litebox_broker_userland/src/runner.rs @@ -415,12 +415,7 @@ const _: () = { assert!(runner_exit_code_is_crash(Some(breakpoint))); }; -/// Accepts a runner's `channel_name` channel with `try_accept`, calling `wait` -/// with the remaining setup time whenever no connection is pending. -/// -/// Fails once `runner_status` reports that the runner stopped or `deadline` -/// passes. -pub fn accept_runner_channel( +fn accept_runner_channel( deadline: Instant, channel_name: &'static str, mut runner_status: impl FnMut() -> IoResult>, @@ -450,17 +445,6 @@ pub fn accept_runner_channel( } } -/// Waits until `listener` has a pending connection or `timeout` elapses. -/// Callers recheck their own condition, so early returns are harmless. -#[cfg(target_os = "linux")] -pub fn wait_for_connection( - listener: &std::os::unix::net::UnixListener, - timeout: Duration, -) -> IoResult<()> { - use std::os::fd::AsFd; - linux::poll_readable([Some(listener.as_fd()), None], Some(timeout)) -} - fn runner_has_exited(runner: &Arc>) -> IoResult { // A pre-authentication caller stops accepting before acting on `true`; // post-authentication callers no longer rely on PID-based authentication. diff --git a/litebox_broker_userland/src/runner/linux.rs b/litebox_broker_userland/src/runner/linux.rs index 17823d8fbb..6e651c9f38 100644 --- a/litebox_broker_userland/src/runner/linux.rs +++ b/litebox_broker_userland/src/runner/linux.rs @@ -5,7 +5,7 @@ use std::any::Any; use std::ffi::OsStr; use std::io::{Error as IoError, ErrorKind, Result as IoResult}; use std::ops::Range; -use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd}; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd}; use std::os::unix::net::UnixListener; use std::os::unix::process::CommandExt; use std::path::PathBuf; @@ -236,23 +236,16 @@ pub(super) fn wait_for_runner_event( Some(_) => timeout, None => Some(timeout.map_or(RUNNER_POLL_INTERVAL, |t| t.min(RUNNER_POLL_INTERVAL))), }; - poll_readable( - [pidfd.as_ref().map(AsFd::as_fd), listener.map(AsFd::as_fd)], - timeout, - ) -} - -/// Waits until any of `fds` becomes readable or `timeout` elapses. -pub(super) fn poll_readable( - fds: [Option>; 2], - timeout: Option, -) -> IoResult<()> { let timeout_ms = timeout.map_or(-1, |timeout| { i32::try_from(timeout.as_nanos().div_ceil(1_000_000)).unwrap_or(i32::MAX) }); // `poll` ignores entries with negative descriptors. - let mut fds = fds.map(|fd| libc::pollfd { - fd: fd.map_or(-1, |fd| fd.as_raw_fd()), + let mut fds = [ + pidfd.as_ref().map_or(-1, AsRawFd::as_raw_fd), + listener.map_or(-1, AsRawFd::as_raw_fd), + ] + .map(|fd| libc::pollfd { + fd, events: libc::POLLIN, revents: 0, }); diff --git a/litebox_broker_userland/src/windows.rs b/litebox_broker_userland/src/windows.rs index 9a0ed9af70..a9c84a301e 100644 --- a/litebox_broker_userland/src/windows.rs +++ b/litebox_broker_userland/src/windows.rs @@ -20,10 +20,9 @@ use litebox_broker_transport_windows_userland::named_pipe::{ }; use litebox_broker_transport_windows_userland::shared_memory::WindowsSharedMemory; use litebox_broker_userland::builder::BrokerCoreBuilder; -use litebox_broker_userland::runner::accept_runner_channel; use litebox_broker_userland::stdio::UserlandStdioProvider; -use super::{IN_PROCESS_RUNNER_POLL_INTERVAL, SETUP_TIMEOUT, configured_socket_policy}; +use super::{SETUP_TIMEOUT, configured_socket_policy}; pub(super) fn run( args: super::CliArgs, @@ -75,15 +74,11 @@ fn serve_runner_in_process( runner: &JoinHandle, ) -> IoResult<()> { let setup_deadline = Instant::now() + SETUP_TIMEOUT; - let control_stream = accept_runner_channel( + let control_stream = crate::accept_runner_channel( setup_deadline, "control", || Ok(runner.is_finished().then(|| "thread stopped".to_owned())), || control_listener.try_accept(), - |remaining| { - std::thread::sleep(remaining.min(IN_PROCESS_RUNNER_POLL_INTERVAL)); - Ok(()) - }, )?; validate_client_process(&control_stream, std::process::id())?; serve_control_stream(