Add ClosableCallback type

Author: OlivierHecart

io/zenoh-transport/src/unicast/universal/rx.rs

@@ -11,7 +11,7 @@
 // Contributors:
 //   ZettaScale Zenoh Team, <zenoh@zettascale.tech>
 //
-use std::sync::{atomic::Ordering, MutexGuard};
+use std::sync::MutexGuard;
 
 use zenoh_buffers::ZSlice;
 use zenoh_codec::transport::frame::FrameReader;
@@ -108,11 +108,7 @@ impl TransportUnicastUniversal {
             // Drop invalid message and continue
             return Ok(());
         }
-        if let Some(callback) = self
-            .callback
-            .get()
-            .filter(|_| !self.closed.load(Ordering::Relaxed))
-        {
+        if let Some(callback) = self.callback.get() {
             for mut msg in frame {
                 self.trigger_callback(
                     callback.as_ref(),
@@ -193,11 +189,7 @@ impl TransportUnicastUniversal {
         if !more {
             // When shared-memory feature is disabled, msg does not need to be mutable
             if let Some(mut msg) = guard.defrag.defragment() {
-                if let Some(callback) = self
-                    .callback
-                    .get()
-                    .filter(|_| !self.closed.load(Ordering::Relaxed))
-                {
+                if let Some(callback) = self.callback.get() {
                     return self.trigger_callback(
                         callback.as_ref(),
                         msg.as_mut(),

io/zenoh-transport/src/unicast/universal/transport.rs

@@ -46,6 +46,35 @@ use crate::{
     TransportManager, TransportPeerEventHandler,
 };
 
+pub(crate) struct ClosableCallback {
+    callback: OnceLock<Arc<dyn TransportPeerEventHandler>>,
+    closed: AtomicBool,
+}
+
+impl ClosableCallback {
+    pub(crate) fn new() -> Self {
+        ClosableCallback {
+            callback: OnceLock::new(),
+            closed: AtomicBool::new(false),
+        }
+    }
+
+    pub(crate) fn set(&self, cb: Arc<dyn TransportPeerEventHandler>) {
+        let _ = self.callback.set(cb);
+    }
+
+    pub(crate) fn get(&self) -> Option<&Arc<dyn TransportPeerEventHandler>> {
+        self.callback
+            .get()
+            .filter(|_| !self.closed.load(Ordering::Relaxed))
+    }
+
+    pub(crate) fn close(&self) -> Option<&Arc<dyn TransportPeerEventHandler>> {
+        self.closed.store(true, Ordering::Relaxed);
+        self.callback.get()
+    }
+}
+
 /*************************************/
 /*        UNIVERSAL TRANSPORT        */
 /*************************************/
@@ -64,8 +93,7 @@ pub(crate) struct TransportUnicastUniversal {
     // The links associated to the channel
     pub(super) links: Arc<RwLock<TransportLinks>>,
     // The callback
-    pub(super) callback: Arc<OnceLock<Arc<dyn TransportPeerEventHandler>>>,
-    pub(super) closed: Arc<AtomicBool>,
+    pub(super) callback: Arc<ClosableCallback>,
     // Lock used to ensure no race in add_link method
     add_link_lock: Arc<AsyncMutex<()>>,
     // Mutex for notification
@@ -112,8 +140,7 @@ impl TransportUnicastUniversal {
             priority_rx: priority_rx.into_boxed_slice().into(),
             links: Arc::new(RwLock::new(TransportLinks::default())),
             add_link_lock: Arc::new(AsyncMutex::new(())),
-            callback: Arc::new(OnceLock::new()),
-            closed: Arc::new(AtomicBool::new(false)),
+            callback: Arc::new(ClosableCallback::new()),
             status: Arc::new(AsyncMutex::new(TransportStatus::Uninitialized)),
             #[cfg(feature = "stats")]
             stats,
@@ -138,8 +165,7 @@ impl TransportUnicastUniversal {
         // to avoid concurrent new_transport and closing/closed notifications
         let mut status_guard = self.get_status().await;
         *status_guard = TransportStatus::Closed;
-        self.closed.store(true, Ordering::Relaxed);
-        let callback = self.callback.get().cloned();
+        let callback = self.callback.close();
 
         // Close all the links
         let mut links = zwrite!(self.links).take();
@@ -323,7 +349,7 @@ impl TransportUnicastTrait for TransportUnicastUniversal {
     /*            ACCESSORS              */
     /*************************************/
     fn set_callback(&self, callback: Arc<dyn TransportPeerEventHandler>) {
-        let _ = self.callback.set(callback);
+        self.callback.set(callback)
     }
 
     async fn get_status(&self) -> AsyncMutexGuard<'_, TransportStatus> {