Generalize callbacks for subscriptions (#260)

* Generalize callbacks to subscriptions

By implementing a trait (SubscriptionCallback) on valid function signatures,
and making subscriptions accept callbacks that implement this trait, we can now

* Receive both plain and boxed messages
* Optionally receive a MessageInfo along with the message, as the second argument
* Soon, receive a loaned message instead of an owned one

This corresponds to the functionality in any_subscription_callback.hpp in rclcpp.

Author: nnmm

rclrs/src/node.rs

@@ -6,7 +6,7 @@ pub use self::graph::*;
 use crate::rcl_bindings::*;
 use crate::{
     Client, ClientBase, Context, ParameterOverrideMap, Publisher, QoSProfile, RclrsError, Service,
-    ServiceBase, Subscription, SubscriptionBase, ToResult,
+    ServiceBase, Subscription, SubscriptionBase, SubscriptionCallback, ToResult,
 };
 
 use std::cmp::PartialEq;
@@ -227,15 +227,14 @@ impl Node {
     ///
     /// [1]: crate::Subscription
     // TODO: make subscription's lifetime depend on node's lifetime
-    pub fn create_subscription<T, F>(
+    pub fn create_subscription<T, Args>(
         &mut self,
         topic: &str,
         qos: QoSProfile,
-        callback: F,
+        callback: impl SubscriptionCallback<T, Args>,
     ) -> Result<Arc<Subscription<T>>, RclrsError>
     where
         T: Message,
-        F: FnMut(T) + 'static + Send,
     {
         let subscription = Arc::new(Subscription::<T>::new(self, topic, qos, callback)?);
         self.subscriptions

rclrs/src/subscription.rs

@@ -3,7 +3,6 @@ use crate::qos::QoSProfile;
 use crate::Node;
 use crate::{rcl_bindings::*, RclrsError};
 
-use std::boxed::Box;
 use std::ffi::CStr;
 use std::ffi::CString;
 use std::marker::PhantomData;
@@ -12,7 +11,11 @@ use std::sync::{Arc, Mutex, MutexGuard};
 
 use rosidl_runtime_rs::{Message, RmwMessage};
 
+mod callback;
+mod message_info;
 mod readonly_loaned_message;
+pub use callback::*;
+pub use message_info::*;
 pub use readonly_loaned_message::*;
 
 // SAFETY: The functions accessing this type, including drop(), shouldn't care about the thread
@@ -71,7 +74,7 @@ where
 {
     pub(crate) handle: Arc<SubscriptionHandle>,
     /// The callback function that runs when a message was received.
-    pub callback: Mutex<Box<dyn FnMut(T) + 'static + Send>>,
+    pub callback: Mutex<AnySubscriptionCallback<T>>,
     message: PhantomData<T>,
 }
 
@@ -80,17 +83,16 @@ where
     T: Message,
 {
     /// Creates a new subscription.
-    pub(crate) fn new<F>(
+    pub(crate) fn new<Args>(
         node: &Node,
         topic: &str,
         qos: QoSProfile,
-        callback: F,
+        callback: impl SubscriptionCallback<T, Args>,
     ) -> Result<Self, RclrsError>
     // This uses pub(crate) visibility to avoid instantiating this struct outside
     // [`Node::create_subscription`], see the struct's documentation for the rationale
     where
         T: Message,
-        F: FnMut(T) + 'static + Send,
     {
         // SAFETY: Getting a zero-initialized value is always safe.
         let mut rcl_subscription = unsafe { rcl_get_zero_initialized_subscription() };
@@ -129,7 +131,7 @@ where
 
         Ok(Self {
             handle,
-            callback: Mutex::new(Box::new(callback)),
+            callback: Mutex::new(callback.into_callback()),
             message: PhantomData,
         })
     }
@@ -171,22 +173,44 @@ where
     // |  rmw_take   |
     // +-------------+
     // ```
-    pub fn take(&self) -> Result<T, RclrsError> {
+    pub fn take(&self) -> Result<(T, MessageInfo), RclrsError> {
         let mut rmw_message = <T as Message>::RmwMsg::default();
+        let message_info = self.take_inner(&mut rmw_message)?;
+        Ok((T::from_rmw_message(rmw_message), message_info))
+    }
+
+    /// This is a version of take() that returns a boxed message.
+    ///
+    /// This can be more efficient for messages containing large arrays.
+    pub fn take_boxed(&self) -> Result<(Box<T>, MessageInfo), RclrsError> {
+        let mut rmw_message = Box::new(<T as Message>::RmwMsg::default());
+        let message_info = self.take_inner(&mut *rmw_message)?;
+        // TODO: This will still use the stack in general. Change signature of
+        // from_rmw_message to allow placing the result in a Box directly.
+        let message = Box::new(T::from_rmw_message(*rmw_message));
+        Ok((message, message_info))
+    }
+
+    // Inner function, to be used by both regular and boxed versions.
+    fn take_inner(
+        &self,
+        rmw_message: &mut <T as Message>::RmwMsg,
+    ) -> Result<MessageInfo, RclrsError> {
+        let mut message_info = unsafe { rmw_get_zero_initialized_message_info() };
         let rcl_subscription = &mut *self.handle.lock();
         unsafe {
             // SAFETY: The first two pointers are valid/initialized, and do not need to be valid
             // beyond the function call.
             // The latter two pointers are explicitly allowed to be NULL.
             rcl_take(
                 rcl_subscription,
-                &mut rmw_message as *mut <T as Message>::RmwMsg as *mut _,
-                std::ptr::null_mut(),
+                rmw_message as *mut <T as Message>::RmwMsg as *mut _,
+                &mut message_info,
                 std::ptr::null_mut(),
             )
             .ok()?
         };
-        Ok(T::from_rmw_message(rmw_message))
+        Ok(MessageInfo::from_rmw_message_info(&message_info))
     }
 }
 
@@ -233,19 +257,38 @@ where
     }
 
     fn execute(&self) -> Result<(), RclrsError> {
-        let msg = match self.take() {
-            Ok(msg) => msg,
+        // Immediately evaluated closure, to handle SubscriptionTakeFailed
+        // outside this match
+        match (|| {
+            match &mut *self.callback.lock().unwrap() {
+                AnySubscriptionCallback::Regular(cb) => {
+                    let (msg, _) = self.take()?;
+                    cb(msg)
+                }
+                AnySubscriptionCallback::Boxed(cb) => {
+                    let (msg, _) = self.take_boxed()?;
+                    cb(msg)
+                }
+                AnySubscriptionCallback::RegularWithMessageInfo(cb) => {
+                    let (msg, msg_info) = self.take()?;
+                    cb(msg, msg_info)
+                }
+                AnySubscriptionCallback::BoxedWithMessageInfo(cb) => {
+                    let (msg, msg_info) = self.take_boxed()?;
+                    cb(msg, msg_info)
+                }
+            }
+            Ok(())
+        })() {
             Err(RclrsError::RclError {
                 code: RclReturnCode::SubscriptionTakeFailed,
                 ..
             }) => {
                 // Spurious wakeup – this may happen even when a waitset indicated that this
                 // subscription was ready, so it shouldn't be an error.
-                return Ok(());
+                Ok(())
             }
-            Err(e) => return Err(e),
-        };
-        (*self.callback.lock().unwrap())(msg);
-        Ok(())
+            other => other,
+        }
     }
 }

rclrs/src/subscription/callback.rs

@@ -0,0 +1,113 @@
+use super::MessageInfo;
+
+/// A trait for allowed callbacks for subscriptions.
+///
+/// See [`AnySubscriptionCallback`] for a list of possible callback signatures.
+pub trait SubscriptionCallback<T, Args>: Send {
+    /// Converts the callback into an enum.
+    ///
+    /// User code never needs to call this function.
+    fn into_callback(self) -> AnySubscriptionCallback<T>;
+}
+
+/// An enum capturing the various possible function signatures for subscription callbacks.
+///
+/// The correct enum variant is deduced by the [`SubscriptionCallback`] trait.
+pub enum AnySubscriptionCallback<T> {
+    /// A callback with only the message as an argument.
+    Regular(Box<dyn FnMut(T) + Send>),
+    /// A callback with the message and the message info as arguments.
+    RegularWithMessageInfo(Box<dyn FnMut(T, MessageInfo) + Send>),
+    /// A callback with only the boxed message as an argument.
+    Boxed(Box<dyn FnMut(Box<T>) + Send>),
+    /// A callback with the boxed message and the message info as arguments.
+    BoxedWithMessageInfo(Box<dyn FnMut(Box<T>, MessageInfo) + Send>),
+}
+
+// We need one implementation per arity. This was inspired by Bevy's systems.
+impl<T, A0, Func> SubscriptionCallback<T, (A0,)> for Func
+where
+    Func: FnMut(A0) + Send + 'static,
+    (A0,): ArgTuple<T, Func = Box<dyn FnMut(A0) + Send>>,
+{
+    fn into_callback(self) -> AnySubscriptionCallback<T> {
+        <(A0,) as ArgTuple<T>>::into_callback_with_args(Box::new(self))
+    }
+}
+
+impl<T, A0, A1, Func> SubscriptionCallback<T, (A0, A1)> for Func
+where
+    Func: FnMut(A0, A1) + Send + 'static,
+    (A0, A1): ArgTuple<T, Func = Box<dyn FnMut(A0, A1) + Send>>,
+{
+    fn into_callback(self) -> AnySubscriptionCallback<T> {
+        <(A0, A1) as ArgTuple<T>>::into_callback_with_args(Box::new(self))
+    }
+}
+
+// Helper trait for SubscriptionCallback.
+//
+// For each tuple of args, it provides conversion from a function with
+// these args to the correct enum variant.
+trait ArgTuple<T> {
+    type Func;
+    fn into_callback_with_args(func: Self::Func) -> AnySubscriptionCallback<T>;
+}
+
+impl<T> ArgTuple<T> for (Box<T>,) {
+    type Func = Box<dyn FnMut(Box<T>) + Send>;
+    fn into_callback_with_args(func: Self::Func) -> AnySubscriptionCallback<T> {
+        AnySubscriptionCallback::Boxed(func)
+    }
+}
+
+impl<T> ArgTuple<T> for (Box<T>, MessageInfo) {
+    type Func = Box<dyn FnMut(Box<T>, MessageInfo) + Send>;
+    fn into_callback_with_args(func: Self::Func) -> AnySubscriptionCallback<T> {
+        AnySubscriptionCallback::BoxedWithMessageInfo(func)
+    }
+}
+
+impl<T> ArgTuple<T> for (T,) {
+    type Func = Box<dyn FnMut(T) + Send>;
+    fn into_callback_with_args(func: Self::Func) -> AnySubscriptionCallback<T> {
+        AnySubscriptionCallback::Regular(func)
+    }
+}
+
+impl<T> ArgTuple<T> for (T, MessageInfo) {
+    type Func = Box<dyn FnMut(T, MessageInfo) + Send>;
+    fn into_callback_with_args(func: Self::Func) -> AnySubscriptionCallback<T> {
+        AnySubscriptionCallback::RegularWithMessageInfo(func)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn callback_conversion() {
+        struct Message;
+        let cb = |_msg: Message| {};
+        assert!(matches!(
+            cb.into_callback(),
+            AnySubscriptionCallback::<Message>::Regular(_)
+        ));
+        let cb = |_msg: Message, _info: MessageInfo| {};
+        assert!(matches!(
+            cb.into_callback(),
+            AnySubscriptionCallback::<Message>::RegularWithMessageInfo(_)
+        ));
+        let cb = |_msg: Box<Message>| {};
+        assert!(matches!(
+            cb.into_callback(),
+            AnySubscriptionCallback::<Message>::Boxed(_)
+        ));
+        let cb = |_msg: Box<Message>, _info: MessageInfo| {};
+        assert!(matches!(
+            cb.into_callback(),
+            AnySubscriptionCallback::<Message>::BoxedWithMessageInfo(_)
+        ));
+    }
+}