From 875e63e2d19b254d5557a44b468ab49764bd26ba Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 18 Jul 2026 06:06:45 -0700 Subject: [PATCH 1/4] Fail broker objects closed on association loss Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ac86dd2d-0280-4d7b-87a8-5654ad1503a6 --- litebox/src/broker/mod.rs | 84 +++++++++++++-- litebox/src/event/counter.rs | 108 +++++++++++++++++++- litebox/src/litebox.rs | 31 ++++-- litebox_broker_transport/src/unix_socket.rs | 70 +++++++++++++ litebox_runner_linux_userland/src/broker.rs | 32 ++++-- litebox_runner_linux_userland/src/lib.rs | 4 +- 6 files changed, 303 insertions(+), 26 deletions(-) diff --git a/litebox/src/broker/mod.rs b/litebox/src/broker/mod.rs index 5334ae9587..144db5a449 100644 --- a/litebox/src/broker/mod.rs +++ b/litebox/src/broker/mod.rs @@ -5,6 +5,7 @@ use alloc::{ sync::{Arc, Weak}, vec::Vec, }; +use core::sync::atomic::{AtomicBool, Ordering}; use hashbrown::HashMap; use litebox_broker_local::BrokerLocal; @@ -53,6 +54,8 @@ pub(crate) trait BrokerControl: Send + Sync { ) -> core::result::Result; fn close_object(&self, handle: ObjectHandle) -> core::result::Result<(), BrokerControlError>; + + fn fail_connection(&self); } pub(crate) struct BrokerHandleRegistry { @@ -112,6 +115,21 @@ impl BrokerHandleRegistry { pollee.notify_observers(events); } } + + fn notify_all(&self, events: Events) { + let pollables = { + let mut pollables = Vec::new(); + self.handles.lock().retain(|_, entry| { + entry.prune_stale_pollables(); + pollables.extend(entry.pollables.iter().filter_map(Weak::upgrade)); + !entry.is_empty() + }); + pollables + }; + for pollee in pollables { + pollee.notify_observers(events); + } + } } struct BrokerHandleEntry { @@ -146,7 +164,9 @@ pub(crate) struct BrokerLocalControl< Platform: RawSyncPrimitivesProvider, Channel: LocalControlChannel + Send, > { - local: Mutex>, + local: Mutex>>, + handles: Arc>, + failed: AtomicBool, } impl BrokerLocalControl @@ -154,9 +174,53 @@ where Platform: RawSyncPrimitivesProvider, Channel: LocalControlChannel + Send, { - pub(crate) const fn new(local: BrokerLocal) -> Self { + pub(crate) fn new( + local: BrokerLocal, + handles: Arc>, + ) -> Self { Self { - local: Mutex::new(local), + local: Mutex::new(Some(local)), + handles, + failed: AtomicBool::new(false), + } + } + + fn request( + &self, + request: impl FnOnce( + &mut BrokerLocal, + ) -> litebox_broker_local::Result, + ) -> core::result::Result { + if self.failed.load(Ordering::Acquire) { + return Err(BrokerControlError::Transport); + } + let (result, notify_failure) = { + let mut local = self.local.lock(); + if self.failed.load(Ordering::Acquire) { + return Err(BrokerControlError::Transport); + } + let result = request( + local + .as_mut() + .expect("active broker connection must retain its control channel"), + ) + .map_err(BrokerControlError::from); + let notify_failure = matches!(result.as_ref(), Err(BrokerControlError::Transport)) + && !self.failed.swap(true, Ordering::AcqRel); + if matches!(result.as_ref(), Err(BrokerControlError::Transport)) { + local.take(); + } + (result, notify_failure) + }; + if notify_failure { + self.handles.notify_all(Events::ERR); + } + result + } + + fn mark_failed(&self) { + if !self.failed.swap(true, Ordering::AcqRel) { + self.handles.notify_all(Events::ERR); } } } @@ -170,14 +234,14 @@ where &self, initial_count: u64, ) -> core::result::Result { - Ok(self.local.lock().create_event_with_count(initial_count)?) + self.request(|local| local.create_event_with_count(initial_count)) } fn wait_event( &self, handle: ObjectHandle, ) -> core::result::Result { - Ok(self.local.lock().wait_event(handle)?) + self.request(|local| local.wait_event(handle)) } fn add_event( @@ -185,7 +249,7 @@ where handle: ObjectHandle, value: u64, ) -> core::result::Result { - Ok(self.local.lock().add_event(handle, value)?) + self.request(|local| local.add_event(handle, value)) } fn consume_event( @@ -193,11 +257,15 @@ where handle: ObjectHandle, mode: EventConsumeMode, ) -> core::result::Result { - Ok(self.local.lock().consume_event(handle, mode)?) + self.request(|local| local.consume_event(handle, mode)) } fn close_object(&self, handle: ObjectHandle) -> core::result::Result<(), BrokerControlError> { - Ok(self.local.lock().close_object(handle)?) + self.request(|local| local.close_object(handle)) + } + + fn fail_connection(&self) { + self.mark_failed(); } } diff --git a/litebox/src/event/counter.rs b/litebox/src/event/counter.rs index b346e1d36d..cb7a206380 100644 --- a/litebox/src/event/counter.rs +++ b/litebox/src/event/counter.rs @@ -200,10 +200,13 @@ mod tests { let handle = ObjectHandle(7); let consume_attempts = Arc::new(AtomicUsize::new(0)); let read_ready = Arc::new(AtomicBool::new(false)); + let request_count = Arc::new(AtomicUsize::new(0)); let local = BrokerLocal::negotiate(FakeLocalControlChannel { handle, consume_attempts: consume_attempts.clone(), read_ready: read_ready.clone(), + request_count, + fail_requests: Arc::new(AtomicBool::new(false)), last_request: None, }) .unwrap(); @@ -246,15 +249,113 @@ mod tests { ); } + #[test] + fn broker_association_failure_wakes_blocked_read() { + use std::time::{Duration, Instant}; + + let platform = MockPlatform::new(); + let handle = ObjectHandle(7); + let consume_attempts = Arc::new(AtomicUsize::new(0)); + let request_count = Arc::new(AtomicUsize::new(0)); + let local = BrokerLocal::negotiate(FakeLocalControlChannel { + handle, + consume_attempts: Arc::clone(&consume_attempts), + read_ready: Arc::new(AtomicBool::new(false)), + request_count: Arc::clone(&request_count), + fail_requests: Arc::new(AtomicBool::new(false)), + last_request: None, + }) + .unwrap(); + let litebox = Arc::new(LiteBox::new_with_broker_local(platform, local)); + let counter = Arc::new(EventCounter::new(&litebox, 0).unwrap()); + + let (result_sender, result_receiver) = std::sync::mpsc::sync_channel(1); + let read_counter = Arc::clone(&counter); + let reader = std::thread::spawn(move || { + result_sender + .send(read_counter.read( + &WaitState::new(platform).context(), + false, + EventCounterReadMode::One, + )) + .unwrap(); + }); + let deadline = Instant::now() + Duration::from_secs(1); + while consume_attempts.load(Ordering::SeqCst) < 2 { + assert!(Instant::now() < deadline); + std::thread::yield_now(); + } + + litebox.broker_failure_dispatcher()(); + + assert!(matches!( + result_receiver + .recv_timeout(Duration::from_secs(1)) + .unwrap(), + Err(TryOpError::Other(EventCounterError::Io)) + )); + reader.join().unwrap(); + assert_eq!(request_count.load(Ordering::SeqCst), 3); + assert_eq!(counter.check_io_events(), Events::ERR); + assert_eq!(request_count.load(Ordering::SeqCst), 3); + } + + #[test] + fn control_transport_failure_notifies_all_event_counters() { + let platform = MockPlatform::new(); + let handle = ObjectHandle(7); + let request_count = Arc::new(AtomicUsize::new(0)); + let fail_requests = Arc::new(AtomicBool::new(false)); + let local = BrokerLocal::negotiate(FakeLocalControlChannel { + handle, + consume_attempts: Arc::new(AtomicUsize::new(0)), + read_ready: Arc::new(AtomicBool::new(false)), + request_count: Arc::clone(&request_count), + fail_requests: Arc::clone(&fail_requests), + last_request: None, + }) + .unwrap(); + let litebox = LiteBox::new_with_broker_local(platform, local); + let first = EventCounter::new(&litebox, 0).unwrap(); + let second = EventCounter::new(&litebox, 0).unwrap(); + let first_observer = Arc::new(ErrorObserver(AtomicBool::new(false))); + let first_observer_dyn: Arc> = first_observer.clone(); + first.register_observer(Arc::downgrade(&first_observer_dyn), Events::ERR); + let second_observer = Arc::new(ErrorObserver(AtomicBool::new(false))); + let second_observer_dyn: Arc> = second_observer.clone(); + second.register_observer(Arc::downgrade(&second_observer_dyn), Events::ERR); + + fail_requests.store(true, Ordering::SeqCst); + + assert_eq!(first.check_io_events(), Events::ERR); + assert!(first_observer.0.load(Ordering::SeqCst)); + assert!(second_observer.0.load(Ordering::SeqCst)); + assert_eq!(request_count.load(Ordering::SeqCst), 3); + assert_eq!(second.check_io_events(), Events::ERR); + assert_eq!(request_count.load(Ordering::SeqCst), 3); + } + + struct ErrorObserver(AtomicBool); + + impl Observer for ErrorObserver { + fn on_events(&self, events: &Events) { + if events.contains(Events::ERR) { + self.0.store(true, Ordering::SeqCst); + } + } + } + struct FakeLocalControlChannel { handle: ObjectHandle, consume_attempts: Arc, read_ready: Arc, + request_count: Arc, + fail_requests: Arc, last_request: Option, } impl LocalControlChannel for FakeLocalControlChannel { - type Error = core::convert::Infallible; + type Error = (); fn send_handshake_request( &mut self, @@ -276,10 +377,15 @@ mod tests { request: &BrokerRequest, ) -> core::result::Result<(), Self::Error> { self.last_request = Some(request.clone()); + self.request_count.fetch_add(1, Ordering::SeqCst); Ok(()) } fn recv_response(&mut self) -> core::result::Result, Self::Error> { + if self.fail_requests.load(Ordering::SeqCst) { + self.last_request.take(); + return Err(()); + } let response = match self.last_request.take().unwrap() { BrokerRequest::Event(EventRequest::Create(_)) => { BrokerResponse::Event(EventResponse::Create(CreateEventResponse { diff --git a/litebox/src/litebox.rs b/litebox/src/litebox.rs index 637c6301b4..936b66809a 100644 --- a/litebox/src/litebox.rs +++ b/litebox/src/litebox.rs @@ -36,7 +36,11 @@ impl LiteBox { /// If the `enforce_singleton_litebox_instance` compilation feature has been enabled, and more /// than one instance is made, will panic. pub fn new(platform: &'static Platform) -> Self { - Self::new_inner(platform, None) + Self::new_inner( + platform, + None, + Arc::new(broker::BrokerHandleRegistry::new()), + ) } /// Create a new [`LiteBox`] instance with a negotiated broker-local control adapter installed. @@ -47,17 +51,18 @@ impl LiteBox { where Channel: LocalControlChannel + Send + 'static, { - Self::new_inner( - platform, - Some(Arc::new( - broker::BrokerLocalControl::::new(broker_local), - )), - ) + let broker_handles = Arc::new(broker::BrokerHandleRegistry::new()); + let broker_control = Arc::new(broker::BrokerLocalControl::::new( + broker_local, + Arc::clone(&broker_handles), + )); + Self::new_inner(platform, Some(broker_control), broker_handles) } fn new_inner( platform: &'static Platform, broker_control: Option>, + broker_handles: Arc>, ) -> Self { // This check ensures that there is exactly one `LiteBox` instance in the process. // @@ -102,7 +107,7 @@ impl LiteBox { platform, descriptors, broker: broker_control, - broker_handles: Arc::new(broker::BrokerHandleRegistry::new()), + broker_handles, }), } } @@ -167,6 +172,16 @@ impl LiteBox { litebox.dispatch_broker_notification(notification); } } + + /// Returns a dispatcher that fails all broker-backed objects when the association closes. + pub fn broker_failure_dispatcher(&self) -> impl Fn() + Send + 'static { + let litebox = self.clone(); + move || { + if let Some(broker) = &litebox.x.broker { + broker.fail_connection(); + } + } + } } /// The actual body of [`LiteBox`], containing any components that might be shared. diff --git a/litebox_broker_transport/src/unix_socket.rs b/litebox_broker_transport/src/unix_socket.rs index 1b5f956a39..2f3c5ee250 100644 --- a/litebox_broker_transport/src/unix_socket.rs +++ b/litebox_broker_transport/src/unix_socket.rs @@ -8,6 +8,7 @@ //! no_std protocol, local, core, and host crates. use std::io::{Error, ErrorKind, Read, Result as IoResult, Write}; +use std::net::Shutdown; use std::os::unix::net::UnixStream; use std::path::Path; use std::time::{Duration, Instant}; @@ -34,6 +35,11 @@ pub struct UnixStreamLocalControlChannel { setup_deadline: Option, } +/// Independently owned handle for interrupting local control-channel I/O. +pub struct UnixStreamLocalControlCancellation { + stream: UnixStream, +} + impl UnixStreamLocalControlChannel { /// Creates a local control channel from an already-connected Unix stream. pub const fn from_connected(stream: UnixStream) -> Self { @@ -62,6 +68,29 @@ impl UnixStreamLocalControlChannel { setup_deadline: Some(deadline), }) } + + /// Creates a handle that can interrupt pending control-channel I/O. + pub fn cancellation_handle(&self) -> IoResult { + self.stream + .try_clone() + .map(|stream| UnixStreamLocalControlCancellation { stream }) + } +} + +impl UnixStreamLocalControlCancellation { + /// Shuts down the control stream, unblocking pending reads or writes. + pub fn cancel(&self) -> IoResult<()> { + match self.stream.shutdown(Shutdown::Both) { + Err(error) if error.kind() == ErrorKind::NotConnected => Ok(()), + result => result, + } + } +} + +impl Drop for UnixStreamLocalControlChannel { + fn drop(&mut self) { + let _ = self.stream.shutdown(Shutdown::Both); + } } /// Host-side Unix-domain-socket control channel for the hosted userland POC. @@ -403,6 +432,47 @@ mod tests { ); } + #[test] + fn local_control_cancellation_unblocks_response_read() { + let (local_stream, _host_stream) = UnixStream::pair().unwrap(); + let mut channel = UnixStreamLocalControlChannel::from_connected(local_stream); + let cancellation = channel.cancellation_handle().unwrap(); + let completed = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let (started_sender, started_receiver) = std::sync::mpsc::sync_channel(1); + let (result_sender, result_receiver) = std::sync::mpsc::sync_channel(1); + let reader_completed = completed.clone(); + let reader = std::thread::spawn(move || { + started_sender.send(()).unwrap(); + result_sender.send(channel.recv_response()).unwrap(); + reader_completed.store(true, std::sync::atomic::Ordering::Release); + }); + + started_receiver + .recv_timeout(Duration::from_secs(1)) + .unwrap(); + std::thread::sleep(Duration::from_millis(50)); + assert!(!completed.load(std::sync::atomic::Ordering::Acquire)); + cancellation.cancel().unwrap(); + + assert!(result_receiver.recv_timeout(Duration::from_secs(1)).is_ok()); + reader.join().unwrap(); + } + + #[test] + fn dropping_local_control_closes_connection_with_cancellation_clone() { + let (local_stream, mut host_stream) = UnixStream::pair().unwrap(); + let channel = UnixStreamLocalControlChannel::from_connected(local_stream); + let _cancellation = channel.cancellation_handle().unwrap(); + host_stream + .set_read_timeout(Some(Duration::from_secs(1))) + .unwrap(); + + drop(channel); + + let mut byte = [0]; + assert_eq!(host_stream.read(&mut byte).unwrap(), 0); + } + #[test] fn host_reports_wrong_phase_request_frames_as_protocol_violations() { let (mut peer_stream, host_stream) = UnixStream::pair().unwrap(); diff --git a/litebox_runner_linux_userland/src/broker.rs b/litebox_runner_linux_userland/src/broker.rs index d6bbc471d6..c17b75d2c8 100644 --- a/litebox_runner_linux_userland/src/broker.rs +++ b/litebox_runner_linux_userland/src/broker.rs @@ -10,7 +10,8 @@ use anyhow::{Context as _, Result}; use litebox_broker_local::{BrokerLocal, BrokerNotifications}; use litebox_broker_protocol::message::BrokerNotification; use litebox_broker_transport::unix_socket::{ - UnixStreamLocalControlChannel, UnixStreamLocalNotificationChannel, + UnixStreamLocalControlCancellation, UnixStreamLocalControlChannel, + UnixStreamLocalNotificationChannel, }; const SETUP_TIMEOUT: Duration = Duration::from_secs(5); @@ -22,6 +23,7 @@ pub(crate) fn connect( ) -> Result<( BrokerLocal, BrokerNotifications, + UnixStreamLocalControlCancellation, )> { let setup_deadline = Instant::now() + SETUP_TIMEOUT; let control_channel = connect_with_retry( @@ -48,26 +50,40 @@ pub(crate) fn connect( notification_socket_path.display() ) })?; + let control_cancellation = control_channel + .cancellation_handle() + .context("failed to create broker control cancellation handle")?; let local = BrokerLocal::negotiate(control_channel).context("broker negotiation failed")?; - Ok((local, BrokerNotifications::new(notification_channel))) + Ok(( + local, + BrokerNotifications::new(notification_channel), + control_cancellation, + )) } pub(crate) fn start_notification_receiver( mut notifications: BrokerNotifications, + control_cancellation: UnixStreamLocalControlCancellation, dispatch_notification: impl Fn(BrokerNotification) + Send + 'static, + dispatch_failure: impl Fn() + Send + 'static, ) -> Result<()> { std::thread::Builder::new() .name("litebox-broker-notifications".to_owned()) .spawn(move || { - loop { + let receive_error = loop { match notifications.recv_notification() { Ok(Some(notification)) => dispatch_notification(notification), - Ok(None) => break, - Err(error) => { - eprintln!("failed to receive broker notification: {error}"); - break; - } + Ok(None) => break None, + Err(error) => break Some(error), } + }; + let cancellation_error = control_cancellation.cancel().err(); + dispatch_failure(); + if let Some(error) = receive_error { + eprintln!("failed to receive broker notification: {error}"); + } + if let Some(error) = cancellation_error { + eprintln!("failed to cancel broker control channel: {error}"); } }) .context("failed to start broker notification receiver")?; diff --git a/litebox_runner_linux_userland/src/lib.rs b/litebox_runner_linux_userland/src/lib.rs index 9b5495b81f..0fc0eee0a7 100644 --- a/litebox_runner_linux_userland/src/lib.rs +++ b/litebox_runner_linux_userland/src/lib.rs @@ -241,14 +241,16 @@ pub fn run(cli_args: CliArgs) -> Result<()> { }; let shim_builder = if let Some(broker_connection) = broker_connection { - let (broker_local, broker_notifications) = broker_connection; + let (broker_local, broker_notifications, broker_control_cancellation) = broker_connection; let litebox = litebox::LiteBox::new_with_broker_local( litebox_platform_multiplex::platform(), broker_local, ); broker::start_notification_receiver( broker_notifications, + broker_control_cancellation, litebox.broker_notification_dispatcher(), + litebox.broker_failure_dispatcher(), )?; litebox_shim_linux::LinuxShimBuilder::new_with_litebox(litebox) } else { From 880117d837ae7353f813f6b26d40459088d8d1cf Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 18 Jul 2026 12:56:28 -0700 Subject: [PATCH 2/4] Simplify broker failure state Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ac86dd2d-0280-4d7b-87a8-5654ad1503a6 --- litebox/src/broker/mod.rs | 58 ++++++++++++++++-------------------- litebox/src/event/counter.rs | 46 ++++++++++++++++++++++++++++ litebox/src/litebox.rs | 16 +++++++--- 3 files changed, 83 insertions(+), 37 deletions(-) diff --git a/litebox/src/broker/mod.rs b/litebox/src/broker/mod.rs index 144db5a449..e97806cd03 100644 --- a/litebox/src/broker/mod.rs +++ b/litebox/src/broker/mod.rs @@ -5,7 +5,6 @@ use alloc::{ sync::{Arc, Weak}, vec::Vec, }; -use core::sync::atomic::{AtomicBool, Ordering}; use hashbrown::HashMap; use litebox_broker_local::BrokerLocal; @@ -116,13 +115,16 @@ impl BrokerHandleRegistry { } } - fn notify_all(&self, events: Events) { + fn notify_all(&self, events: Events) + where + Platform: TimeProvider, + { let pollables = { let mut pollables = Vec::new(); self.handles.lock().retain(|_, entry| { entry.prune_stale_pollables(); pollables.extend(entry.pollables.iter().filter_map(Weak::upgrade)); - !entry.is_empty() + !entry.pollables.is_empty() }); pollables }; @@ -166,12 +168,11 @@ pub(crate) struct BrokerLocalControl< > { local: Mutex>>, handles: Arc>, - failed: AtomicBool, } impl BrokerLocalControl where - Platform: RawSyncPrimitivesProvider, + Platform: RawSyncPrimitivesProvider + TimeProvider, Channel: LocalControlChannel + Send, { pub(crate) fn new( @@ -181,7 +182,6 @@ where Self { local: Mutex::new(Some(local)), handles, - failed: AtomicBool::new(false), } } @@ -191,43 +191,31 @@ where &mut BrokerLocal, ) -> litebox_broker_local::Result, ) -> core::result::Result { - if self.failed.load(Ordering::Acquire) { - return Err(BrokerControlError::Transport); - } - let (result, notify_failure) = { + let (result, failed_connection) = { let mut local = self.local.lock(); - if self.failed.load(Ordering::Acquire) { + let Some(connection) = local.as_mut() else { return Err(BrokerControlError::Transport); - } - let result = request( - local - .as_mut() - .expect("active broker connection must retain its control channel"), - ) - .map_err(BrokerControlError::from); - let notify_failure = matches!(result.as_ref(), Err(BrokerControlError::Transport)) - && !self.failed.swap(true, Ordering::AcqRel); - if matches!(result.as_ref(), Err(BrokerControlError::Transport)) { - local.take(); - } - (result, notify_failure) + }; + let result = request(connection).map_err(BrokerControlError::from); + let failed_connection = if matches!(result.as_ref(), Err(BrokerControlError::Transport)) + { + local.take() + } else { + None + }; + (result, failed_connection) }; - if notify_failure { + if let Some(connection) = failed_connection { + drop(connection); self.handles.notify_all(Events::ERR); } result } - - fn mark_failed(&self) { - if !self.failed.swap(true, Ordering::AcqRel) { - self.handles.notify_all(Events::ERR); - } - } } impl BrokerControl for BrokerLocalControl where - Platform: RawSyncPrimitivesProvider, + Platform: RawSyncPrimitivesProvider + TimeProvider, Channel: LocalControlChannel + Send, { fn create_event_with_count( @@ -265,7 +253,11 @@ where } fn fail_connection(&self) { - self.mark_failed(); + let connection = self.local.lock().take(); + if let Some(connection) = connection { + drop(connection); + self.handles.notify_all(Events::ERR); + } } } diff --git a/litebox/src/event/counter.rs b/litebox/src/event/counter.rs index cb7a206380..82759f619b 100644 --- a/litebox/src/event/counter.rs +++ b/litebox/src/event/counter.rs @@ -335,6 +335,42 @@ mod tests { assert_eq!(request_count.load(Ordering::SeqCst), 3); } + #[test] + fn broker_dispatchers_follow_objects_that_outlive_litebox() { + let platform = MockPlatform::new(); + let handle = ObjectHandle(7); + let request_count = Arc::new(AtomicUsize::new(0)); + let local = BrokerLocal::negotiate(FakeLocalControlChannel { + handle, + consume_attempts: Arc::new(AtomicUsize::new(0)), + read_ready: Arc::new(AtomicBool::new(false)), + request_count: Arc::clone(&request_count), + fail_requests: Arc::new(AtomicBool::new(false)), + last_request: None, + }) + .unwrap(); + let litebox = LiteBox::new_with_broker_local(platform, local); + let counter = EventCounter::new(&litebox, 0).unwrap(); + let read_observer = Arc::new(ReadObserver(AtomicBool::new(false))); + let read_observer_dyn: Arc> = read_observer.clone(); + counter.register_observer(Arc::downgrade(&read_observer_dyn), Events::IN); + let litebox_weak = Arc::downgrade(&litebox.x); + let dispatch_notification = litebox.broker_notification_dispatcher(); + let dispatch_failure = litebox.broker_failure_dispatcher(); + + drop(litebox); + + assert!(litebox_weak.upgrade().is_none()); + dispatch_notification(BrokerNotification::Readiness(ReadinessNotification { + handle, + readiness: ReadinessFlags::READ, + })); + assert!(read_observer.0.load(Ordering::SeqCst)); + dispatch_failure(); + assert_eq!(counter.check_io_events(), Events::ERR); + assert_eq!(request_count.load(Ordering::SeqCst), 1); + } + struct ErrorObserver(AtomicBool); impl Observer for ErrorObserver { @@ -345,6 +381,16 @@ mod tests { } } + struct ReadObserver(AtomicBool); + + impl Observer for ReadObserver { + fn on_events(&self, events: &Events) { + if events.contains(Events::IN) { + self.0.store(true, Ordering::SeqCst); + } + } + } + struct FakeLocalControlChannel { handle: ObjectHandle, consume_attempts: Arc, diff --git a/litebox/src/litebox.rs b/litebox/src/litebox.rs index 936b66809a..4d0c1d1737 100644 --- a/litebox/src/litebox.rs +++ b/litebox/src/litebox.rs @@ -49,6 +49,7 @@ impl LiteBox { broker_local: BrokerLocal, ) -> Self where + Platform: TimeProvider, Channel: LocalControlChannel + Send + 'static, { let broker_handles = Arc::new(broker::BrokerHandleRegistry::new()); @@ -167,17 +168,24 @@ impl LiteBox { where Platform: TimeProvider + 'static, { - let litebox = self.clone(); + let broker_handles = Arc::downgrade(&self.x.broker_handles); move |notification| { - litebox.dispatch_broker_notification(notification); + if let Some(broker_handles) = broker_handles.upgrade() { + match notification { + BrokerNotification::Readiness(notification) => { + broker_handles + .notify_readiness(notification.handle, notification.readiness); + } + } + } } } /// Returns a dispatcher that fails all broker-backed objects when the association closes. pub fn broker_failure_dispatcher(&self) -> impl Fn() + Send + 'static { - let litebox = self.clone(); + let broker = self.x.broker.as_ref().map(Arc::downgrade); move || { - if let Some(broker) = &litebox.x.broker { + if let Some(broker) = broker.as_ref().and_then(alloc::sync::Weak::upgrade) { broker.fail_connection(); } } From 1ab51aa5330279311f749558c94cd5a9ea76f772 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 18 Jul 2026 13:05:56 -0700 Subject: [PATCH 3/4] Clarify broker pollable registry naming Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ac86dd2d-0280-4d7b-87a8-5654ad1503a6 --- litebox/src/broker/mod.rs | 22 +++++++++++----------- litebox/src/event/counter.rs | 13 +++++++------ litebox/src/litebox.rs | 26 +++++++++++++------------- 3 files changed, 31 insertions(+), 30 deletions(-) diff --git a/litebox/src/broker/mod.rs b/litebox/src/broker/mod.rs index e97806cd03..e67c3e8c8c 100644 --- a/litebox/src/broker/mod.rs +++ b/litebox/src/broker/mod.rs @@ -57,11 +57,11 @@ pub(crate) trait BrokerControl: Send + Sync { fn fail_connection(&self); } -pub(crate) struct BrokerHandleRegistry { - handles: Mutex>>, +pub(crate) struct BrokerPollableRegistry { + handles: Mutex>>, } -impl BrokerHandleRegistry { +impl BrokerPollableRegistry { pub(crate) fn new() -> Self { Self { handles: Mutex::new(HashMap::new()), @@ -72,7 +72,7 @@ impl BrokerHandleRegistry { self.handles .lock() .entry(handle) - .or_insert_with(BrokerHandleEntry::new) + .or_insert_with(BrokerPollableEntry::new) .register_pollable(pollee); } @@ -134,11 +134,11 @@ impl BrokerHandleRegistry { } } -struct BrokerHandleEntry { +struct BrokerPollableEntry { pollables: Vec>>, } -impl BrokerHandleEntry { +impl BrokerPollableEntry { fn new() -> Self { Self { pollables: Vec::new(), @@ -167,7 +167,7 @@ pub(crate) struct BrokerLocalControl< Channel: LocalControlChannel + Send, > { local: Mutex>>, - handles: Arc>, + pollable_registry: Arc>, } impl BrokerLocalControl @@ -177,11 +177,11 @@ where { pub(crate) fn new( local: BrokerLocal, - handles: Arc>, + pollable_registry: Arc>, ) -> Self { Self { local: Mutex::new(Some(local)), - handles, + pollable_registry, } } @@ -207,7 +207,7 @@ where }; if let Some(connection) = failed_connection { drop(connection); - self.handles.notify_all(Events::ERR); + self.pollable_registry.notify_all(Events::ERR); } result } @@ -256,7 +256,7 @@ where let connection = self.local.lock().take(); if let Some(connection) = connection { drop(connection); - self.handles.notify_all(Events::ERR); + self.pollable_registry.notify_all(Events::ERR); } } } diff --git a/litebox/src/event/counter.rs b/litebox/src/event/counter.rs index 82759f619b..e6272481cb 100644 --- a/litebox/src/event/counter.rs +++ b/litebox/src/event/counter.rs @@ -12,7 +12,7 @@ use thiserror::Error; use crate::{ LiteBox, broker::{ - BrokerControl, BrokerHandleRegistry, + BrokerControl, BrokerPollableRegistry, error::{BrokerControlError, BrokerObjectError}, readiness_events, }, @@ -46,7 +46,7 @@ pub enum EventCounterError { pub struct EventCounter { broker: Arc, handle: ObjectHandle, - registry: Arc>, + pollable_registry: Arc>, pollee: Arc>, } @@ -68,13 +68,13 @@ where .create_event_with_count(initial_count) .map_err(BrokerObjectError::from) .map_err(EventCounterError::from)?; - let registry = litebox.broker_handle_registry(); + let pollable_registry = litebox.broker_pollable_registry(); let pollee = Arc::new(Pollee::new()); - registry.register_pollable(handle, &pollee); + pollable_registry.register_pollable(handle, &pollee); Ok(Self { broker, handle, - registry, + pollable_registry, pollee, }) } @@ -143,7 +143,8 @@ where Platform: RawSyncPrimitivesProvider + TimeProvider, { fn drop(&mut self) { - self.registry.unregister_pollable(self.handle, &self.pollee); + self.pollable_registry + .unregister_pollable(self.handle, &self.pollee); let _ = self.broker.close_object(self.handle); } } diff --git a/litebox/src/litebox.rs b/litebox/src/litebox.rs index 4d0c1d1737..36838a8db7 100644 --- a/litebox/src/litebox.rs +++ b/litebox/src/litebox.rs @@ -39,7 +39,7 @@ impl LiteBox { Self::new_inner( platform, None, - Arc::new(broker::BrokerHandleRegistry::new()), + Arc::new(broker::BrokerPollableRegistry::new()), ) } @@ -52,18 +52,18 @@ impl LiteBox { Platform: TimeProvider, Channel: LocalControlChannel + Send + 'static, { - let broker_handles = Arc::new(broker::BrokerHandleRegistry::new()); + let broker_pollables = Arc::new(broker::BrokerPollableRegistry::new()); let broker_control = Arc::new(broker::BrokerLocalControl::::new( broker_local, - Arc::clone(&broker_handles), + Arc::clone(&broker_pollables), )); - Self::new_inner(platform, Some(broker_control), broker_handles) + Self::new_inner(platform, Some(broker_control), broker_pollables) } fn new_inner( platform: &'static Platform, broker_control: Option>, - broker_handles: Arc>, + broker_pollables: Arc>, ) -> Self { // This check ensures that there is exactly one `LiteBox` instance in the process. // @@ -108,7 +108,7 @@ impl LiteBox { platform, descriptors, broker: broker_control, - broker_handles, + broker_pollables, }), } } @@ -146,8 +146,8 @@ impl LiteBox { self.x.broker.clone() } - pub(crate) fn broker_handle_registry(&self) -> Arc> { - Arc::clone(&self.x.broker_handles) + pub(crate) fn broker_pollable_registry(&self) -> Arc> { + Arc::clone(&self.x.broker_pollables) } /// Dispatches one broker notification to the matching local-core object. @@ -158,7 +158,7 @@ impl LiteBox { match notification { BrokerNotification::Readiness(notification) => self .x - .broker_handles + .broker_pollables .notify_readiness(notification.handle, notification.readiness), } } @@ -168,12 +168,12 @@ impl LiteBox { where Platform: TimeProvider + 'static, { - let broker_handles = Arc::downgrade(&self.x.broker_handles); + let broker_pollables = Arc::downgrade(&self.x.broker_pollables); move |notification| { - if let Some(broker_handles) = broker_handles.upgrade() { + if let Some(broker_pollables) = broker_pollables.upgrade() { match notification { BrokerNotification::Readiness(notification) => { - broker_handles + broker_pollables .notify_readiness(notification.handle, notification.readiness); } } @@ -197,5 +197,5 @@ pub(crate) struct LiteBoxX { pub(crate) platform: &'static Platform, descriptors: RwLock>, broker: Option>, - broker_handles: Arc>, + broker_pollables: Arc>, } From 2002108a94f4e0e849e844746f1dd85ddc662200 Mon Sep 17 00:00:00 2001 From: Weidong Cui Date: Sat, 18 Jul 2026 13:12:45 -0700 Subject: [PATCH 4/4] Simplify broker pollable registry Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ac86dd2d-0280-4d7b-87a8-5654ad1503a6 --- litebox/src/broker/mod.rs | 84 ++++++++++-------------------------- litebox/src/event/counter.rs | 29 ++++++------- 2 files changed, 35 insertions(+), 78 deletions(-) diff --git a/litebox/src/broker/mod.rs b/litebox/src/broker/mod.rs index e67c3e8c8c..42d040bad7 100644 --- a/litebox/src/broker/mod.rs +++ b/litebox/src/broker/mod.rs @@ -58,32 +58,26 @@ pub(crate) trait BrokerControl: Send + Sync { } pub(crate) struct BrokerPollableRegistry { - handles: Mutex>>, + pollables: Mutex>>>, } impl BrokerPollableRegistry { pub(crate) fn new() -> Self { Self { - handles: Mutex::new(HashMap::new()), + pollables: Mutex::new(HashMap::new()), } } pub(crate) fn register_pollable(&self, handle: ObjectHandle, pollee: &Arc>) { - self.handles - .lock() - .entry(handle) - .or_insert_with(BrokerPollableEntry::new) - .register_pollable(pollee); + let previous = self.pollables.lock().insert(handle, Arc::downgrade(pollee)); + assert!( + previous.is_none(), + "broker handle already has a registered pollable" + ); } - pub(crate) fn unregister_pollable(&self, handle: ObjectHandle, pollee: &Arc>) { - let mut handles = self.handles.lock(); - if let Some(entry) = handles.get_mut(&handle) { - entry.unregister_pollable(pollee); - if entry.pollables.is_empty() { - handles.remove(&handle); - } - } + pub(crate) fn unregister_pollable(&self, handle: ObjectHandle) { + self.pollables.lock().remove(&handle); } pub(crate) fn notify_readiness(&self, handle: ObjectHandle, readiness: ReadinessFlags) @@ -94,23 +88,15 @@ impl BrokerPollableRegistry { if events.is_empty() { return; } - let pollables = { - let mut handles = self.handles.lock(); - let Some(entry) = handles.get_mut(&handle) else { - return; - }; - entry.prune_stale_pollables(); - let pollables = entry - .pollables - .iter() - .filter_map(Weak::upgrade) - .collect::>(); - if entry.pollables.is_empty() { - handles.remove(&handle); + let pollee = { + let mut pollables = self.pollables.lock(); + let pollee = pollables.get(&handle).and_then(Weak::upgrade); + if pollee.is_none() { + pollables.remove(&handle); } - pollables + pollee }; - for pollee in pollables { + if let Some(pollee) = pollee { pollee.notify_observers(events); } } @@ -121,10 +107,12 @@ impl BrokerPollableRegistry { { let pollables = { let mut pollables = Vec::new(); - self.handles.lock().retain(|_, entry| { - entry.prune_stale_pollables(); - pollables.extend(entry.pollables.iter().filter_map(Weak::upgrade)); - !entry.pollables.is_empty() + self.pollables.lock().retain(|_, registered| { + let Some(pollee) = registered.upgrade() else { + return false; + }; + pollables.push(pollee); + true }); pollables }; @@ -134,34 +122,6 @@ impl BrokerPollableRegistry { } } -struct BrokerPollableEntry { - pollables: Vec>>, -} - -impl BrokerPollableEntry { - fn new() -> Self { - Self { - pollables: Vec::new(), - } - } - - fn register_pollable(&mut self, pollee: &Arc>) { - self.pollables.push(Arc::downgrade(pollee)); - } - - fn unregister_pollable(&mut self, pollee: &Arc>) { - self.pollables.retain(|registered| { - registered - .upgrade() - .is_some_and(|registered| !Arc::ptr_eq(®istered, pollee)) - }); - } - - fn prune_stale_pollables(&mut self) { - self.pollables - .retain(|registered| registered.strong_count() > 0); - } -} pub(crate) struct BrokerLocalControl< Platform: RawSyncPrimitivesProvider, Channel: LocalControlChannel + Send, diff --git a/litebox/src/event/counter.rs b/litebox/src/event/counter.rs index e6272481cb..4ebd50e604 100644 --- a/litebox/src/event/counter.rs +++ b/litebox/src/event/counter.rs @@ -143,8 +143,7 @@ where Platform: RawSyncPrimitivesProvider + TimeProvider, { fn drop(&mut self) { - self.pollable_registry - .unregister_pollable(self.handle, &self.pollee); + self.pollable_registry.unregister_pollable(self.handle); let _ = self.broker.close_object(self.handle); } } @@ -203,7 +202,7 @@ mod tests { let read_ready = Arc::new(AtomicBool::new(false)); let request_count = Arc::new(AtomicUsize::new(0)); let local = BrokerLocal::negotiate(FakeLocalControlChannel { - handle, + next_handle: handle.0, consume_attempts: consume_attempts.clone(), read_ready: read_ready.clone(), request_count, @@ -259,7 +258,7 @@ mod tests { let consume_attempts = Arc::new(AtomicUsize::new(0)); let request_count = Arc::new(AtomicUsize::new(0)); let local = BrokerLocal::negotiate(FakeLocalControlChannel { - handle, + next_handle: handle.0, consume_attempts: Arc::clone(&consume_attempts), read_ready: Arc::new(AtomicBool::new(false)), request_count: Arc::clone(&request_count), @@ -308,7 +307,7 @@ mod tests { let request_count = Arc::new(AtomicUsize::new(0)); let fail_requests = Arc::new(AtomicBool::new(false)); let local = BrokerLocal::negotiate(FakeLocalControlChannel { - handle, + next_handle: handle.0, consume_attempts: Arc::new(AtomicUsize::new(0)), read_ready: Arc::new(AtomicBool::new(false)), request_count: Arc::clone(&request_count), @@ -342,7 +341,7 @@ mod tests { let handle = ObjectHandle(7); let request_count = Arc::new(AtomicUsize::new(0)); let local = BrokerLocal::negotiate(FakeLocalControlChannel { - handle, + next_handle: handle.0, consume_attempts: Arc::new(AtomicUsize::new(0)), read_ready: Arc::new(AtomicBool::new(false)), request_count: Arc::clone(&request_count), @@ -393,7 +392,7 @@ mod tests { } struct FakeLocalControlChannel { - handle: ObjectHandle, + next_handle: u64, consume_attempts: Arc, read_ready: Arc, request_count: Arc, @@ -435,13 +434,11 @@ mod tests { } let response = match self.last_request.take().unwrap() { BrokerRequest::Event(EventRequest::Create(_)) => { - BrokerResponse::Event(EventResponse::Create(CreateEventResponse { - handle: self.handle, - })) + let handle = ObjectHandle(self.next_handle); + self.next_handle += 1; + BrokerResponse::Event(EventResponse::Create(CreateEventResponse { handle })) } - BrokerRequest::Event(EventRequest::Consume(request)) - if request.handle == self.handle => - { + BrokerRequest::Event(EventRequest::Consume(_)) => { self.consume_attempts.fetch_add(1, Ordering::SeqCst); if self.read_ready.swap(false, Ordering::SeqCst) { BrokerResponse::Event(EventResponse::Consume(EventConsumption { @@ -452,10 +449,10 @@ mod tests { BrokerResponse::Error(ErrorCode::WouldBlock) } } - BrokerRequest::CloseObject(handle) if handle == self.handle => { - BrokerResponse::ObjectClosed + BrokerRequest::CloseObject(_) => BrokerResponse::ObjectClosed, + request @ BrokerRequest::Event(_) => { + panic!("unexpected broker request: {request:?}") } - request => panic!("unexpected broker request: {request:?}"), }; Ok(Some(response)) }