Ensure that a waitable is only added to one wait set at a time

Author: nnmm

rclrs/src/error.rs

@@ -27,6 +27,8 @@ pub enum RclrsError {
         /// The error indicating the position of the nul byte.
         err: NulError,
     },
+    /// It was attempted to add a waitable to a wait set twice.
+    AlreadyAddedToWaitSet,
 }
 
 impl Display for RclrsError {
@@ -37,6 +39,12 @@ impl Display for RclrsError {
             RclrsError::StringContainsNul { s, .. } => {
                 write!(f, "Could not convert string '{}' to CString", s)
             }
+            RclrsError::AlreadyAddedToWaitSet => {
+                write!(
+                    f,
+                    "Could not add entity to wait set because it was already added to a wait set"
+                )
+            }
         }
     }
 }
@@ -68,6 +76,7 @@ impl Error for RclrsError {
             RclrsError::RclError { msg, .. } => msg.as_ref().map(|e| e as &dyn Error),
             RclrsError::UnknownRclError { msg, .. } => msg.as_ref().map(|e| e as &dyn Error),
             RclrsError::StringContainsNul { err, .. } => Some(err).map(|e| e as &dyn Error),
+            RclrsError::AlreadyAddedToWaitSet => None,
         }
     }
 }

rclrs/src/node/client.rs

@@ -2,10 +2,13 @@ use futures::channel::oneshot;
 use std::boxed::Box;
 use std::collections::HashMap;
 use std::ffi::CString;
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
 use crate::rcl_bindings::*;
-use crate::{MessageCow, Node, RclReturnCode, RclrsError, ToResult, WaitSet, Waitable};
+use crate::{
+    ExclusivityGuard, MessageCow, Node, RclReturnCode, RclrsError, ToResult, WaitSet, Waitable,
+};
 
 use parking_lot::Mutex;
 use rosidl_runtime_rs::Message;
@@ -27,6 +30,7 @@ where
     rcl_node_mtx: Arc<Mutex<rcl_node_t>>,
     requests: Mutex<HashMap<RequestId, RequestValue<T::Response>>>,
     futures: Arc<Mutex<HashMap<RequestId, oneshot::Sender<T::Response>>>>,
+    in_use_by_wait_set: Arc<AtomicBool>,
 }
 
 impl<T> Drop for Client<T>
@@ -48,6 +52,9 @@ where
     T: rosidl_runtime_rs::Service,
 {
     unsafe fn add_to_wait_set(self: Arc<Self>, wait_set: &mut WaitSet) -> Result<(), RclrsError> {
+        // First, ensure that self is not yet used by a wait set.
+        let client_exclusive =
+            ExclusivityGuard::new(Arc::clone(&self) as _, Arc::clone(&self.in_use_by_wait_set))?;
         // SAFETY: I'm not sure if it's required, but the client pointer will remain valid
         // for as long as the wait set exists, because it's stored in self.clients.
         // Passing in a null pointer for the third argument is explicitly allowed.
@@ -57,7 +64,7 @@ where
             std::ptr::null_mut(),
         )
         .ok()?;
-        wait_set.clients.push(self);
+        wait_set.clients.push(client_exclusive);
         Ok(())
     }
 
@@ -134,6 +141,7 @@ where
             futures: Arc::new(Mutex::new(
                 HashMap::<RequestId, oneshot::Sender<T::Response>>::new(),
             )),
+            in_use_by_wait_set: Arc::new(AtomicBool::new(false)),
         })
     }
 

rclrs/src/node/service.rs

@@ -1,9 +1,10 @@
 use std::boxed::Box;
 use std::ffi::CString;
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
 use crate::rcl_bindings::*;
-use crate::{Node, RclReturnCode, RclrsError, ToResult, WaitSet, Waitable};
+use crate::{ExclusivityGuard, Node, RclReturnCode, RclrsError, ToResult, WaitSet, Waitable};
 
 use rosidl_runtime_rs::Message;
 
@@ -27,6 +28,7 @@ where
     rcl_node_mtx: Arc<Mutex<rcl_node_t>>,
     /// The callback function that runs when a request was received.
     pub callback: Mutex<ServiceCallback<T::Request, T::Response>>,
+    in_use_by_wait_set: Arc<AtomicBool>,
 }
 
 impl<T> Drop for Service<T>
@@ -48,6 +50,9 @@ where
     T: rosidl_runtime_rs::Service,
 {
     unsafe fn add_to_wait_set(self: Arc<Self>, wait_set: &mut WaitSet) -> Result<(), RclrsError> {
+        // First, ensure that self is not yet used by a wait set.
+        let service_exclusive =
+            ExclusivityGuard::new(Arc::clone(&self) as _, Arc::clone(&self.in_use_by_wait_set))?;
         // SAFETY: I'm not sure if it's required, but the service pointer will remain valid
         // for as long as the wait set exists, because it's stored in self.clients.
         // Passing in a null pointer for the third argument is explicitly allowed.
@@ -57,7 +62,7 @@ where
             std::ptr::null_mut(),
         )
         .ok()?;
-        wait_set.services.push(self);
+        wait_set.services.push(service_exclusive);
         Ok(())
     }
 
@@ -136,6 +141,7 @@ where
             rcl_service_mtx: Mutex::new(service_handle),
             rcl_node_mtx: Arc::clone(&node.rcl_node_mtx),
             callback: Mutex::new(Box::new(callback)),
+            in_use_by_wait_set: Arc::new(AtomicBool::new(false)),
         })
     }
 

rclrs/src/node/subscription.rs

@@ -1,11 +1,12 @@
 use crate::error::{RclReturnCode, ToResult};
 use crate::{rcl_bindings::*, RclrsError};
-use crate::{Node, QoSProfile, WaitSet, Waitable};
+use crate::{ExclusivityGuard, Node, QoSProfile, WaitSet, Waitable};
 
 use std::boxed::Box;
 use std::ffi::CStr;
 use std::ffi::CString;
 use std::marker::PhantomData;
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 
 use rosidl_runtime_rs::{Message, RmwMessage};
@@ -36,6 +37,7 @@ where
     /// The callback function that runs when a message was received.
     pub callback: Mutex<Box<dyn FnMut(T) + 'static + Send>>,
     message: PhantomData<T>,
+    in_use_by_wait_set: Arc<AtomicBool>,
 }
 
 impl<T: Message> Drop for Subscription<T> {
@@ -54,6 +56,9 @@ where
     T: Message,
 {
     unsafe fn add_to_wait_set(self: Arc<Self>, wait_set: &mut WaitSet) -> Result<(), RclrsError> {
+        // First, ensure that self is not yet used by a wait set.
+        let subscription_exclusive =
+            ExclusivityGuard::new(Arc::clone(&self) as _, Arc::clone(&self.in_use_by_wait_set))?;
         // SAFETY: I'm not sure if it's required, but the subscription pointer will remain valid
         // for as long as the wait set exists, because it's stored in self.subscriptions.
         // Passing in a null pointer for the third argument is explicitly allowed.
@@ -63,7 +68,7 @@ where
             std::ptr::null_mut(),
         )
         .ok()?;
-        wait_set.subscriptions.push(self);
+        wait_set.subscriptions.push(subscription_exclusive);
         Ok(())
     }
 
@@ -139,6 +144,7 @@ where
             rcl_node_mtx: Arc::clone(&node.rcl_node_mtx),
             callback: Mutex::new(Box::new(callback)),
             message: PhantomData,
+            in_use_by_wait_set: Arc::new(AtomicBool::new(false)),
         })
     }
 

rclrs/src/wait.rs

@@ -19,12 +19,47 @@ use crate::error::{to_rclrs_result, RclReturnCode, RclrsError, ToResult};
 use crate::rcl_bindings::*;
 use crate::{ClientWaitable, Context, ServiceWaitable, SubscriptionWaitable};
 
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
 use std::time::Duration;
 use std::vec::Vec;
 
 use parking_lot::Mutex;
 
+/// A helper struct for tracking whether the waitable is currently in a wait set.
+///
+/// When this struct is constructed, which happens when adding an entity to the wait set,
+/// it checks that the atomic boolean is false and sets it to true.
+/// When it is dropped, which happens when it is removed from the wait set,
+/// or the wait set itself is dropped, it sets the atomic bool to false.
+pub(crate) struct ExclusivityGuard<T> {
+    in_use_by_wait_set: Arc<AtomicBool>,
+    waitable: T,
+}
+
+impl<T> Drop for ExclusivityGuard<T> {
+    fn drop(&mut self) {
+        self.in_use_by_wait_set.store(false, Ordering::Relaxed)
+    }
+}
+
+impl<T> ExclusivityGuard<T> {
+    pub fn new(waitable: T, in_use_by_wait_set: Arc<AtomicBool>) -> Result<Self, RclrsError> {
+        if in_use_by_wait_set
+            .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
+            .is_err()
+        {
+            return Err(RclrsError::AlreadyAddedToWaitSet);
+        }
+        Ok(Self {
+            in_use_by_wait_set,
+            waitable,
+        })
+    }
+}
+
 /// Trait to be implemented by entities that can be waited on, like a [`Subscription`][1].
 ///
 /// [1]: crate::Subscription
@@ -59,9 +94,9 @@ pub struct WaitSet {
     // The subscriptions that are currently registered in the wait set.
     // This correspondence is an invariant that must be maintained by all functions,
     // even in the error case.
-    pub(crate) subscriptions: Vec<Arc<dyn SubscriptionWaitable>>,
-    pub(crate) clients: Vec<Arc<dyn ClientWaitable>>,
-    pub(crate) services: Vec<Arc<dyn ServiceWaitable>>,
+    pub(crate) subscriptions: Vec<ExclusivityGuard<Arc<dyn SubscriptionWaitable>>>,
+    pub(crate) clients: Vec<ExclusivityGuard<Arc<dyn ClientWaitable>>>,
+    pub(crate) services: Vec<ExclusivityGuard<Arc<dyn ServiceWaitable>>>,
 }
 
 /// A list of entities that are ready, returned by [`WaitSet::wait`].
@@ -128,13 +163,14 @@ impl WaitSet {
 
     /// Adds a client to the wait set.
     ///
-    /// It is possible, but not useful, to add the same client twice.
-    ///
-    /// This will return an error if the number of clients in the wait set is larger than the
-    /// capacity set in [`WaitSet::new`].
+    /// # Errors
+    /// - If the client was already added to this wait set or another one,
+    ///   [`AlreadyAddedToWaitSet`][1] will be returned
+    /// - If the number of clients in the wait set is larger than the
+    ///   capacity set in [`WaitSet::new`], [`WaitSetFull`][2] will be returned
     ///
-    /// The same client must not be added to multiple wait sets, because that would make it
-    /// unsafe to simultaneously wait on those wait sets.
+    /// [1]: crate::RclrsError
+    /// [2]: crate::RclReturnCode
     pub fn add_client(&mut self, client: Arc<dyn ClientWaitable>) -> Result<(), RclrsError> {
         // SAFETY: The implementation of this trait for clients checks that the client
         // has not already been added to a different wait set.
@@ -143,13 +179,14 @@ impl WaitSet {
 
     /// Adds a service to the wait set.
     ///
-    /// It is possible, but not useful, to add the same service twice.
-    ///
-    /// This will return an error if the number of services in the wait set is larger than the
-    /// capacity set in [`WaitSet::new`].
+    /// # Errors
+    /// - If the service was already added to this wait set or another one,
+    ///   [`AlreadyAddedToWaitSet`][1] will be returned
+    /// - If the number of services in the wait set is larger than the
+    ///   capacity set in [`WaitSet::new`], [`WaitSetFull`][2] will be returned
     ///
-    /// The same service must not be added to multiple wait sets, because that would make it
-    /// unsafe to simultaneously wait on those wait sets.
+    /// [1]: crate::RclrsError
+    /// [2]: crate::RclReturnCode
     pub fn add_service(&mut self, service: Arc<dyn ServiceWaitable>) -> Result<(), RclrsError> {
         // SAFETY: The implementation of this trait for services checks that the service
         // has not already been added to a different wait set.
@@ -158,13 +195,14 @@ impl WaitSet {
 
     /// Adds a subscription to the wait set.
     ///
-    /// It is possible, but not useful, to add the same subscription twice.
-    ///
-    /// This will return an error if the number of subscriptions in the wait set is larger than the
-    /// capacity set in [`WaitSet::new`].
+    /// # Errors
+    /// - If the subscription was already added to this wait set or another one,
+    ///   [`AlreadyAddedToWaitSet`][1] will be returned
+    /// - If the number of subscriptions in the wait set is larger than the
+    ///   capacity set in [`WaitSet::new`], [`WaitSetFull`][2] will be returned
     ///
-    /// The same subscription must not be added to multiple wait sets, because that would make it
-    /// unsafe to simultaneously wait on those wait sets.
+    /// [1]: crate::RclrsError
+    /// [2]: crate::RclReturnCode
     pub fn add_subscription(
         &mut self,
         subscription: Arc<dyn SubscriptionWaitable>,
@@ -241,7 +279,9 @@ impl WaitSet {
             // https://github.com/ros2/rcl/blob/35a31b00a12f259d492bf53c0701003bd7f1745c/rcl/include/rcl/wait.h#L419
             let wait_set_entry = unsafe { *self.rcl_wait_set.subscriptions.add(i) };
             if !wait_set_entry.is_null() {
-                ready_entities.subscriptions.push(subscription.clone());
+                ready_entities
+                    .subscriptions
+                    .push(Arc::clone(&subscription.waitable));
             }
         }
         for (i, client) in self.clients.iter().enumerate() {
@@ -250,7 +290,7 @@ impl WaitSet {
             // https://github.com/ros2/rcl/blob/35a31b00a12f259d492bf53c0701003bd7f1745c/rcl/include/rcl/wait.h#L419
             let wait_set_entry = unsafe { *self.rcl_wait_set.clients.add(i) };
             if !wait_set_entry.is_null() {
-                ready_entities.clients.push(client.clone());
+                ready_entities.clients.push(Arc::clone(&client.waitable));
             }
         }
         for (i, service) in self.services.iter().enumerate() {
@@ -259,7 +299,7 @@ impl WaitSet {
             // https://github.com/ros2/rcl/blob/35a31b00a12f259d492bf53c0701003bd7f1745c/rcl/include/rcl/wait.h#L419
             let wait_set_entry = unsafe { *self.rcl_wait_set.services.add(i) };
             if !wait_set_entry.is_null() {
-                ready_entities.services.push(service.clone());
+                ready_entities.services.push(Arc::clone(&service.waitable));
             }
         }
         Ok(ready_entities)