diff --git a/native/rust/dimos-module/src/lcm.rs b/native/rust/dimos-module/src/lcm.rs index 6fe3838253..6a834976e4 100644 --- a/native/rust/dimos-module/src/lcm.rs +++ b/native/rust/dimos-module/src/lcm.rs @@ -21,7 +21,7 @@ use std::time::Duration; use dimos_lcm::{Lcm, LcmOptions}; use url::Url; -use crate::transport::{Dispatch, Transport}; +use crate::transport::{Dispatch, TopicDispatch, Transport}; /// LCM UDP multicast transport. Wraps `dimos_lcm::Lcm`. /// @@ -30,6 +30,7 @@ use crate::transport::{Dispatch, Transport}; pub struct LcmTransport { inner: Arc, routes: Arc>>>, + all_routes: Arc>>, listening: AtomicBool, /// The runtime the transport was opened on. In a baked host each module has /// its own runtime, so the one shared recv loop must not land on whichever @@ -96,6 +97,7 @@ impl LcmTransport { Self { inner: Arc::new(inner), routes: Arc::new(Mutex::new(HashMap::new())), + all_routes: Arc::new(Mutex::new(Vec::new())), listening: AtomicBool::new(false), runtime: tokio::runtime::Handle::current(), } @@ -104,6 +106,7 @@ impl LcmTransport { fn spawn_recv_loop(&self) { let inner = Arc::clone(&self.inner); let routes = Arc::clone(&self.routes); + let all_routes = Arc::clone(&self.all_routes); self.runtime.spawn(async move { loop { match inner.recv().await { @@ -114,6 +117,7 @@ impl LcmTransport { cb(&msg.data); } } + dispatch_all(&all_routes, &msg.channel, &msg.data); } Err(e) => { crate::error_throttled!( @@ -128,6 +132,15 @@ impl LcmTransport { } } +fn dispatch_all(all_routes: &Mutex>, channel: &str, data: &[u8]) { + if channel != "LCM_SELF_TEST" { + let callbacks = all_routes.lock().unwrap().clone(); + for callback in &callbacks { + callback(data, channel); + } + } +} + impl Transport for LcmTransport { async fn publish(&self, channel: &str, data: Vec) -> io::Result<()> { self.inner.publish(channel, &data).await @@ -146,6 +159,14 @@ impl Transport for LcmTransport { Ok(()) } + async fn subscribe_all(&self, on_msg: TopicDispatch) -> io::Result<()> { + self.all_routes.lock().unwrap().push(on_msg); + if !self.listening.swap(true, Ordering::SeqCst) { + self.spawn_recv_loop(); + } + Ok(()) + } + /// LCM has no per-topic publisher settings and no notion of a session-local /// publisher, so a baked host cannot hide an internal hop on this transport. fn set_publisher_qos(&self, qos: &serde_json::Value) { @@ -169,8 +190,10 @@ impl Transport for LcmTransport { #[cfg(test)] mod tests { - use super::options_from_url; + use super::{dispatch_all, options_from_url}; + use crate::transport::TopicDispatch; use std::net::Ipv4Addr; + use std::sync::{Arc, Mutex}; #[test] fn reads_group_port_and_ttl() { @@ -202,4 +225,24 @@ mod tests { assert_eq!(options.ttl, defaults.ttl, "{url}"); } } + + #[test] + fn subscribe_all_receives_the_channel_and_payload() { + let all = Arc::new(Mutex::new(Vec::new())); + let all_sink = Arc::clone(&all); + let all_routes = Mutex::new(vec![Arc::new(move |data: &[u8], topic: &str| { + all_sink + .lock() + .unwrap() + .push((topic.to_string(), data.to_vec())); + }) as TopicDispatch]); + + dispatch_all(&all_routes, "camera", b"frame"); + dispatch_all(&all_routes, "LCM_SELF_TEST", b"ignore"); + + assert_eq!( + *all.lock().unwrap(), + [("camera".to_string(), b"frame".to_vec())] + ); + } } diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 23fe965ba3..6038b9053b 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -27,7 +27,7 @@ use serde::de::DeserializeOwned; use serde::{Deserialize, Serialize}; use validator::Validate; -use crate::transport::{Dispatch, Transport}; +use crate::transport::{Dispatch, TopicDispatch, Transport}; /// Marker trait for a config checked by `#[native_config]`: every field required, /// no Rust-side defaults, no unknown fields. Implemented only by the macro. @@ -309,6 +309,7 @@ pub struct Builder { // Every port the module asked for a topic, matched against topics after build. requested: BTreeSet, routes: HashMap>>, + all_routes: Vec, // One publish queue per output channel, drained by its own worker. outputs: Vec<(String, mpsc::Receiver>)>, tf: Option, @@ -320,6 +321,7 @@ impl Builder { topics, requested: BTreeSet::new(), routes: HashMap::new(), + all_routes: Vec::new(), outputs: Vec::new(), tf: None, } @@ -386,6 +388,14 @@ impl Builder { Input { topic, receiver } } + /// Register a raw callback for every DimOS topic. The callback runs on the + /// transport delivery path, so it should hand expensive work to another task. + /// It receives the transport-native channel name and remains registered for + /// the transport's lifetime. + pub fn subscribe_all(&mut self, on_msg: impl Fn(&[u8], &str) + Send + Sync + 'static) { + self.all_routes.push(Arc::new(on_msg)); + } + pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { let topic = self.topic_for(port); let sender = self.add_publisher(&topic); @@ -456,6 +466,25 @@ pub(crate) async fn subscribe_routes( Ok(()) } +pub(crate) async fn subscribe_all_routes( + transport: &T, + routes: Vec, +) -> io::Result<()> { + if routes.is_empty() { + return Ok(()); + } + let dispatch: TopicDispatch = Arc::new(move |bytes, topic| { + for route in &routes { + let dispatched = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| route(bytes, topic))); + if dispatched.is_err() { + error!(topic, "dispatch handler panicked; message dropped"); + } + } + }); + transport.subscribe_all(dispatch).await +} + /// Spawn one worker per output channel so they don't block each other pub(crate) fn spawn_publish_tasks( transport: Arc, @@ -527,6 +556,7 @@ where builder.enforce_topics_match_ports()?; subscribe_routes(transport.as_ref(), builder.routes).await?; + subscribe_all_routes(transport.as_ref(), builder.all_routes).await?; // Kept alive until teardown so the subscriptions stay live. let mut pub_tasks = spawn_publish_tasks(Arc::clone(&transport), builder.outputs); @@ -639,6 +669,7 @@ mod tests { inbound: Arc, inbound_notify: Arc, subscriptions: Arc>>>, + all_subscriptions: Arc>>, listening: Arc, publish_delay_ms: Arc, publish_entered: Arc, @@ -653,6 +684,7 @@ mod tests { inbound: Arc::new(InboundQueue::new(VecDeque::new())), inbound_notify: Arc::new(Notify::new()), subscriptions: Arc::new(Mutex::new(HashMap::new())), + all_subscriptions: Arc::new(Mutex::new(Vec::new())), listening: Arc::new(AtomicBool::new(false)), publish_delay_ms: Arc::new(AtomicU64::new(0)), publish_entered: Arc::new(Notify::new()), @@ -666,6 +698,7 @@ mod tests { let inbound = Arc::clone(&self.inbound); let inbound_notify = Arc::clone(&self.inbound_notify); let subscriptions = Arc::clone(&self.subscriptions); + let all_subscriptions = Arc::clone(&self.all_subscriptions); let dispatch_entered = Arc::clone(&self.dispatch_entered); let dispatch_log = Arc::clone(&self.dispatch_log); tokio::spawn(async move { @@ -679,6 +712,10 @@ mod tests { cb(&data); } } + let callbacks = all_subscriptions.lock().unwrap().clone(); + for callback in &callbacks { + callback(&data, &channel); + } dispatch_log.lock().unwrap().push(Instant::now()); } else { inbound_notify.notified().await; @@ -711,6 +748,14 @@ mod tests { } Ok(()) } + + async fn subscribe_all(&self, on_msg: TopicDispatch) -> io::Result<()> { + self.all_subscriptions.lock().unwrap().push(on_msg); + if !self.listening.swap(true, Ordering::SeqCst) { + self.spawn_delivery_loop(); + } + Ok(()) + } } fn inject_inbound(inbound: &InboundQueue, notify: &Notify, channel: &str, data: Vec) { @@ -939,6 +984,38 @@ mod tests { assert_eq!(input.topic, "/data"); } + #[tokio::test] + async fn subscribe_all_routes_delivers_new_topics() { + let transport = ControllableMockTransport::new(); + let inbound = Arc::clone(&transport.inbound); + let inbound_notify = Arc::clone(&transport.inbound_notify); + let seen = Arc::new(Mutex::new(Vec::new())); + let sink = Arc::clone(&seen); + let mut builder = builder_with_topics(&[]); + builder.subscribe_all(move |data, topic| { + sink.lock() + .unwrap() + .push((topic.to_string(), data.to_vec())); + }); + builder + .enforce_topics_match_ports() + .expect("subscribe-all is not a fixed port"); + subscribe_all_routes(&transport, builder.all_routes) + .await + .expect("subscribe to all topics"); + + inject_inbound(&inbound, &inbound_notify, "/late", b"message".to_vec()); + wait_for("subscribe-all callback", || { + !seen.lock().unwrap().is_empty() + }) + .await; + + assert_eq!( + *seen.lock().unwrap(), + [("/late".to_string(), b"message".to_vec())] + ); + } + #[test] fn output_uses_mapped_topic() { let mut builder = builder_with_topics(&[("cmd_vel", "/robot/cmd_vel")]); diff --git a/native/rust/dimos-module/src/transport.rs b/native/rust/dimos-module/src/transport.rs index 616a662d18..0698d563bd 100644 --- a/native/rust/dimos-module/src/transport.rs +++ b/native/rust/dimos-module/src/transport.rs @@ -22,6 +22,9 @@ use std::sync::Arc; /// happen inside it. pub type Dispatch = Arc; +/// Dispatch closure for a subscription that receives every DimOS channel. +pub type TopicDispatch = Arc; + /// Abstraction over the message transport used by a native module. /// /// New transport protocols should implement this trait. @@ -36,6 +39,17 @@ pub trait Transport: Send + Sync + 'static { on_msg: Dispatch, ) -> impl Future> + Send; + /// Deliver every DimOS message to `on_msg`, including messages on channels + /// that appear after the subscription starts. + fn subscribe_all(&self, _on_msg: TopicDispatch) -> impl Future> + Send { + async { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "transport does not support subscribe_all", + )) + } + } + /// Apply the per-channel publisher QoS the coordinator sends. The value is /// the `qos` object from the stdin config, or null when absent. Transports /// without per-topic QoS ignore it. @@ -54,6 +68,7 @@ pub(crate) trait DynTransport: Send + Sync + 'static { channel: &'a str, on_msg: Dispatch, ) -> BoxFuture<'a, io::Result<()>>; + fn subscribe_all_dyn(&self, on_msg: TopicDispatch) -> BoxFuture<'_, io::Result<()>>; fn set_publisher_qos_dyn(&self, qos: &serde_json::Value); } @@ -70,6 +85,10 @@ impl DynTransport for T { Box::pin(self.subscribe(channel, on_msg)) } + fn subscribe_all_dyn(&self, on_msg: TopicDispatch) -> BoxFuture<'_, io::Result<()>> { + Box::pin(self.subscribe_all(on_msg)) + } + fn set_publisher_qos_dyn(&self, qos: &serde_json::Value) { self.set_publisher_qos(qos) } @@ -96,6 +115,10 @@ impl Transport for SharedTransport { self.0.subscribe_dyn(channel, on_msg).await } + async fn subscribe_all(&self, on_msg: TopicDispatch) -> io::Result<()> { + self.0.subscribe_all_dyn(on_msg).await + } + fn set_publisher_qos(&self, qos: &serde_json::Value) { self.0.set_publisher_qos_dyn(qos) } diff --git a/native/rust/dimos-module/src/zenoh.rs b/native/rust/dimos-module/src/zenoh.rs index 9e455adb5f..ffb7e8dd74 100644 --- a/native/rust/dimos-module/src/zenoh.rs +++ b/native/rust/dimos-module/src/zenoh.rs @@ -24,7 +24,7 @@ use ::zenoh::Session; use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; -use crate::transport::{Dispatch, Transport}; +use crate::transport::{Dispatch, TopicDispatch, Transport}; pub(crate) const SESSION_KEY: &str = "session"; @@ -322,6 +322,17 @@ impl Transport for ZenohTransport { .map_err(to_io) } + async fn subscribe_all(&self, on_msg: TopicDispatch) -> io::Result<()> { + self.session + .declare_subscriber("dimos/**") + .callback(move |sample| { + on_msg(&sample.payload().to_bytes(), sample.key_expr().as_str()) + }) + .background() + .await + .map_err(to_io) + } + fn set_publisher_qos(&self, qos: &serde_json::Value) { let _ = self.qos.set(parse_channel_qos(qos)); } @@ -672,6 +683,41 @@ mod tests { assert_eq!(received, payload); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn subscribe_all_receives_new_channels() { + let transport = ZenohTransport::new().await.expect("open session"); + let channel = format!("dimos/subscribe_all/{}", std::process::id()); + let (tx, mut rx) = tokio::sync::mpsc::channel::<(String, Vec)>(8); + let sink: TopicDispatch = Arc::new(move |bytes: &[u8], topic: &str| { + let _ = tx.try_send((topic.to_string(), bytes.to_vec())); + }); + transport + .subscribe_all(sink) + .await + .expect("subscribe to all channels"); + + let payload = b"new channel"; + let received = tokio::time::timeout(Duration::from_secs(10), async { + 'publish: loop { + transport + .publish(&channel, payload.to_vec()) + .await + .expect("publish"); + while let Ok(Some(message)) = + tokio::time::timeout(Duration::from_millis(100), rx.recv()).await + { + if message.0 == channel { + break 'publish message; + } + } + } + }) + .await + .expect("new channel not delivered within timeout"); + + assert_eq!(received, (channel, payload.to_vec())); + } + #[test] fn parse_channel_qos_reads_set_fields() { let value = serde_json::json!({