Implement option for callback drop wait on undeclare for Publisher, Subscriber, Queryable, Querier, Matching listener and Session

Author: DenisBiryukov91

zenoh/src/api/admin.rs

@@ -172,6 +172,8 @@ pub(crate) fn init(session: WeakSession) {
             let session = session.clone();
             move |q| on_admin_query(&session, &prefix, &prefix, q)
         }),
+        #[cfg(feature = "unstable")]
+        None,
     );
 
     // Queryable simulating advanced publisher to allow advanced subscriber to receive historical data
@@ -185,6 +187,8 @@ pub(crate) fn init(session: WeakSession) {
             let session = session.clone();
             move |q| on_admin_query(&session, &adv_prefix, &prefix, q)
         }),
+        #[cfg(feature = "unstable")]
+        None,
     );
 
     // Subscribe to transport events and publish them to the adminspace

zenoh/src/api/builders/close.rs

@@ -34,6 +34,7 @@ use zenoh_runtime::ZRuntime;
 pub struct CloseBuilder<TCloseable: Closeable> {
     closee: TCloseable::TClosee,
     timeout: Duration,
+    close_args: <TCloseable::TClosee as Closee>::CloseArgs,
 }
 
 // NOTE: `Closeable` is only pub(crate) because it is a zenoh-internal trait, so we don't
@@ -44,6 +45,7 @@ impl<TCloseable: Closeable> CloseBuilder<TCloseable> {
         Self {
             closee: closeable.get_closee(),
             timeout: Duration::from_secs(10),
+            close_args: <TCloseable::TClosee as Closee>::CloseArgs::default(),
         }
     }
 
@@ -115,7 +117,7 @@ impl<TCloseable: Closeable> IntoFuture for CloseBuilder<TCloseable> {
     fn into_future(self) -> Self::IntoFuture {
         Box::pin(
             async move {
-                if tokio::time::timeout(self.timeout, self.closee.close_inner())
+                if tokio::time::timeout(self.timeout, self.closee.close_inner(self.close_args))
                     .await
                     .is_err()
                 {
@@ -220,10 +222,20 @@ impl<TOutput: Send + 'static> IntoFuture for NolocalJoinHandle<TOutput> {
 
 #[async_trait]
 pub(crate) trait Closee: Send + Sync + 'static {
-    async fn close_inner(&self);
+    type CloseArgs: Default + Send + Sync;
+    async fn close_inner(&self, close_arg: Self::CloseArgs);
 }
 
 pub(crate) trait Closeable {
     type TClosee: Closee;
     fn get_closee(&self) -> Self::TClosee;
 }
+
+impl CloseBuilder<crate::Session> {
+    /// Block in undeclare operation until all currently running zenoh entities' callbacks (if any) return.
+    #[zenoh_macros::unstable]
+    pub fn wait_until_callback_execution_ends(mut self) -> Self {
+        self.close_args.wait_until_callback_execution_ends = true;
+        self
+    }
+}

zenoh/src/api/builders/liveliness.rs

@@ -255,13 +255,17 @@ where
         let key_expr = self.key_expr?;
         let session = self.session;
         let (callback, handler) = self.handler.into_handler();
+        #[cfg(feature = "unstable")]
+        let callback_sync_group = crate::api::cancellation::SyncGroup::default();
         session
             .0
             .declare_liveliness_subscriber_inner(
                 &key_expr,
                 Locality::default(),
                 self.history,
                 callback,
+                #[cfg(feature = "unstable")]
+                callback_sync_group.notifier(),
             )
             .map(|sub_state| Subscriber {
                 inner: SubscriberInner {
@@ -272,6 +276,8 @@ where
                     undeclare_on_drop: true,
                 },
                 handler,
+                #[cfg(feature = "unstable")]
+                callback_sync_group,
             })
     }
 }
@@ -300,6 +306,8 @@ impl Wait for LivelinessSubscriberBuilder<'_, '_, Callback<Sample>, true> {
             Locality::default(),
             self.history,
             self.handler,
+            #[cfg(feature = "unstable")]
+            None,
         )?;
         Ok(())
     }
@@ -510,35 +518,15 @@ where
     Handler::Handler: Send,
 {
     fn wait(self) -> <Self as Resolvable>::To {
-        #[allow(unused_mut)] // mut is needed only for unstable cancellation_token
-        let (mut callback, receiver) = self.handler.into_handler();
-        #[cfg(feature = "unstable")]
-        let cancellation_token = if let Some(ct) = self.cancellation_token {
-            if let Some(notifier) = ct.notifier() {
-                callback.set_on_drop_notifier(notifier);
-                Some(ct)
-            } else {
-                return Ok(receiver);
-            }
-        } else {
-            None
-        };
-        #[allow(unused_variables)] // qid is only needed for unstable cancellation_token
-        self.session
-            .0
-            .liveliness_query(&self.key_expr?, self.timeout, callback)
-            .map(|qid| {
-                #[cfg(feature = "unstable")]
-                if let Some(cancellation_token) = cancellation_token {
-                    let weak_session = self.session.downgrade();
-                    let on_cancel = move || {
-                        let _ = weak_session.cancel_liveliness_query(qid); // fails only if no associated query exists - likely because it was already finalized
-                        Ok(())
-                    };
-                    cancellation_token.add_on_cancel_handler(Box::new(on_cancel));
-                }
-                receiver
-            })
+        let (callback, receiver) = self.handler.into_handler();
+        self.session.0.liveliness_query(
+            &self.key_expr?,
+            self.timeout,
+            callback,
+            #[cfg(feature = "unstable")]
+            self.cancellation_token.into_iter().collect(),
+        )?;
+        Ok(receiver)
     }
 }
 

zenoh/src/api/builders/matching_listener.rs

@@ -21,6 +21,8 @@ use std::{
 use zenoh_core::{Resolvable, Wait};
 use zenoh_result::ZResult;
 
+#[cfg(feature = "unstable")]
+use crate::api::cancellation::{SyncGroup, SyncGroupNotifier};
 use crate::{
     api::{
         handlers::{Callback, DefaultHandler, IntoHandler},
@@ -44,6 +46,8 @@ pub struct MatchingListenerBuilder<'a, Handler, const BACKGROUND: bool = false>
     pub(crate) matching_listeners: &'a Arc<Mutex<HashSet<Id>>>,
     pub(crate) matching_status_type: MatchingStatusType,
     pub handler: Handler,
+    #[cfg(feature = "unstable")]
+    pub(crate) parent_callback_sync_group_notifier: Option<SyncGroupNotifier>,
 }
 
 impl<'a> MatchingListenerBuilder<'a, DefaultHandler> {
@@ -140,6 +144,8 @@ impl<'a> MatchingListenerBuilder<'a, DefaultHandler> {
             matching_listeners: self.matching_listeners,
             matching_status_type: self.matching_status_type,
             handler,
+            #[cfg(feature = "unstable")]
+            parent_callback_sync_group_notifier: self.parent_callback_sync_group_notifier,
         }
     }
 }
@@ -179,6 +185,8 @@ impl<'a> MatchingListenerBuilder<'a, Callback<MatchingStatus>> {
             key_expr: self.key_expr,
             matching_status_type: self.matching_status_type,
             handler: self.handler,
+            #[cfg(feature = "unstable")]
+            parent_callback_sync_group_notifier: self.parent_callback_sync_group_notifier,
         }
     }
 }
@@ -197,12 +205,16 @@ where
     Handler::Handler: Send,
 {
     fn wait(self) -> <Self as Resolvable>::To {
+        #[cfg(feature = "unstable")]
+        let callback_sync_group = SyncGroup::default();
         let (callback, handler) = self.handler.into_handler();
         let state = self.session.declare_matches_listener_inner(
             self.key_expr,
             self.destination,
             self.matching_status_type,
             callback,
+            #[cfg(feature = "unstable")]
+            callback_sync_group.notifier(),
         )?;
         zlock!(self.matching_listeners).insert(state.id);
         Ok(MatchingListener {
@@ -213,6 +225,8 @@ where
                 undeclare_on_drop: true,
             },
             handler,
+            #[cfg(feature = "unstable")]
+            callback_sync_group,
         })
     }
 }
@@ -241,6 +255,8 @@ impl Wait for MatchingListenerBuilder<'_, Callback<MatchingStatus>, true> {
             self.destination,
             self.matching_status_type,
             self.handler,
+            #[cfg(feature = "unstable")]
+            self.parent_callback_sync_group_notifier,
         )?;
         zlock!(self.matching_listeners).insert(state.id);
         Ok(())

zenoh/src/api/builders/publisher.rs

@@ -21,6 +21,8 @@ use zenoh_protocol::core::CongestionControl;
 #[cfg(feature = "unstable")]
 use zenoh_protocol::core::Reliability;
 
+#[cfg(feature = "unstable")]
+use crate::api::cancellation::SyncGroup;
 #[cfg(feature = "unstable")]
 use crate::api::sample::SourceInfo;
 use crate::{
@@ -501,6 +503,8 @@ impl Wait for PublisherBuilder<'_, '_> {
             reliability: self.reliability,
             matching_listeners: Default::default(),
             undeclare_on_drop: true,
+            #[cfg(feature = "unstable")]
+            sync_group: SyncGroup::default(),
         })
     }
 }

zenoh/src/api/builders/querier.rs

@@ -25,7 +25,7 @@ use zenoh_result::ZResult;
 
 use super::sample::QoSBuilderTrait;
 #[cfg(feature = "unstable")]
-use crate::api::cancellation::CancellationTokenBuilderTrait;
+use crate::api::cancellation::{CancellationToken, CancellationTokenBuilderTrait};
 #[cfg(feature = "unstable")]
 use crate::api::query::ReplyKeyExpr;
 #[cfg(feature = "unstable")]
@@ -191,6 +191,8 @@ impl Wait for QuerierBuilder<'_, '_> {
             #[cfg(feature = "unstable")]
             accept_replies: self.accept_replies,
             matching_listeners: Default::default(),
+            #[cfg(feature = "unstable")]
+            cancellation_token: CancellationToken::default(),
         })
     }
 }
@@ -473,19 +475,7 @@ where
     Handler::Handler: Send,
 {
     fn wait(self) -> <Self as Resolvable>::To {
-        #[allow(unused_mut)] // mut is needed only for unstable cancellation_token
-        let (mut callback, receiver) = self.handler.into_handler();
-        #[cfg(feature = "unstable")]
-        let cancellation_token = if let Some(ct) = self.cancellation_token {
-            if let Some(notifier) = ct.notifier() {
-                callback.set_on_drop_notifier(notifier);
-                Some(ct)
-            } else {
-                return Ok(receiver);
-            }
-        } else {
-            None
-        };
+        let (callback, receiver) = self.handler.into_handler();
         #[allow(unused_mut)]
         // mut is only needed when building with "unstable" feature, which might add extra internal parameters on top of the user-provided ones
         let mut parameters = self.parameters.clone();
@@ -494,34 +484,26 @@ where
             parameters.set_reply_key_expr_any();
         }
         #[allow(unused_variables)] // qid is only needed for unstable cancellation_token
-        self.querier
-            .session
-            .query(
-                &self.querier.key_expr,
-                &parameters,
-                self.querier.target,
-                self.querier.consolidation,
-                self.querier.qos,
-                self.querier.destination,
-                self.querier.timeout,
-                self.value,
-                self.attachment,
-                #[cfg(feature = "unstable")]
-                self.source_info,
-                callback,
-            )
-            .map(|qid| {
-                #[cfg(feature = "unstable")]
-                if let Some(cancellation_token) = cancellation_token {
-                    let weak_session = self.querier.session.clone();
-                    let on_cancel = move || {
-                        let _ = weak_session.cancel_query(qid); // fails only if no associated query exists - likely because it was already finalized
-                        Ok(())
-                    };
-                    cancellation_token.add_on_cancel_handler(Box::new(on_cancel));
-                }
-                receiver
-            })
+        self.querier.session.query(
+            &self.querier.key_expr,
+            &parameters,
+            self.querier.target,
+            self.querier.consolidation,
+            self.querier.qos,
+            self.querier.destination,
+            self.querier.timeout,
+            self.value,
+            self.attachment,
+            #[cfg(feature = "unstable")]
+            self.source_info,
+            callback,
+            #[cfg(feature = "unstable")]
+            match self.cancellation_token {
+                Some(ct) => vec![self.querier.cancellation_token.clone(), ct],
+                None => vec![self.querier.cancellation_token.clone()],
+            },
+        )?;
+        Ok(receiver)
     }
 }
 

zenoh/src/api/builders/query.rs

@@ -374,52 +374,28 @@ where
     Handler::Handler: Send,
 {
     fn wait(self) -> <Self as Resolvable>::To {
-        #[allow(unused_mut)] // mut is needed only for unstable cancellation_token
-        let (mut callback, receiver) = self.handler.into_handler();
-        #[cfg(feature = "unstable")]
-        let cancellation_token = if let Some(ct) = self.cancellation_token {
-            if let Some(notifier) = ct.notifier() {
-                callback.set_on_drop_notifier(notifier);
-                Some(ct)
-            } else {
-                return Ok(receiver);
-            }
-        } else {
-            None
-        };
+        let (callback, receiver) = self.handler.into_handler();
         let Selector {
             key_expr,
             parameters,
         } = self.selector?;
-        #[allow(unused_variables)] // qid is only needed for unstable cancellation_token
-        self.session
-            .0
-            .query(
-                &key_expr,
-                &parameters,
-                self.target,
-                self.consolidation,
-                self.qos.into(),
-                self.destination,
-                self.timeout,
-                self.value,
-                self.attachment,
-                #[cfg(feature = "unstable")]
-                self.source_info,
-                callback,
-            )
-            .map(|qid| {
-                #[cfg(feature = "unstable")]
-                if let Some(cancellation_token) = cancellation_token {
-                    let weak_session = self.session.downgrade();
-                    let on_cancel = move || {
-                        let _ = weak_session.cancel_query(qid); // fails only if no associated query exists - likely because it was already finalized
-                        Ok(())
-                    };
-                    cancellation_token.add_on_cancel_handler(Box::new(on_cancel));
-                }
-                receiver
-            })
+        self.session.0.query(
+            &key_expr,
+            &parameters,
+            self.target,
+            self.consolidation,
+            self.qos.into(),
+            self.destination,
+            self.timeout,
+            self.value,
+            self.attachment,
+            #[cfg(feature = "unstable")]
+            self.source_info,
+            callback,
+            #[cfg(feature = "unstable")]
+            self.cancellation_token.into_iter().collect(),
+        )?;
+        Ok(receiver)
     }
 }