From 3ac3ec2c2bffd4cbbaa81c28c7b075d90773ed18 Mon Sep 17 00:00:00 2001
From: Yadunund <yadunund@intrinsic.ai>
Date: Thu, 19 Sep 2024 09:18:47 +0800
Subject: [PATCH] Close the Zenoh session after destroying the last node

Signed-off-by: Yadunund <yadunund@intrinsic.ai>
---
 rmw_zenoh_cpp/src/detail/rmw_data_types.hpp |  7 +++++++
 rmw_zenoh_cpp/src/rmw_init.cpp              |  5 -----
 rmw_zenoh_cpp/src/rmw_zenoh.cpp             | 14 ++++++++++++++
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/rmw_zenoh_cpp/src/detail/rmw_data_types.hpp b/rmw_zenoh_cpp/src/detail/rmw_data_types.hpp
index b39728e2..a70ab759 100644
--- a/rmw_zenoh_cpp/src/detail/rmw_data_types.hpp
+++ b/rmw_zenoh_cpp/src/detail/rmw_data_types.hpp
@@ -24,6 +24,7 @@
 #include <optional>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <variant>
 #include <vector>
@@ -52,6 +53,12 @@ class rmw_context_impl_s final
   // An owned session.
   z_owned_session_t session;
 
+  // This is a temporary workaround for the unsoundness of rmw_shutdown
+  // reported in https://github.com/ros2/rmw_zenoh/issues/170.
+  // We keep track of all the nodes created in this session and only close
+  // the Zenoh session once the last node is destroyed.
+  std::unordered_set<const rmw_node_t *> session_nodes_;
+
   std::optional<z_owned_shm_provider_t> shm_provider;
 
   z_owned_subscriber_t graph_subscriber;
diff --git a/rmw_zenoh_cpp/src/rmw_init.cpp b/rmw_zenoh_cpp/src/rmw_init.cpp
index 7cc221de..43ccd871 100644
--- a/rmw_zenoh_cpp/src/rmw_init.cpp
+++ b/rmw_zenoh_cpp/src/rmw_init.cpp
@@ -412,11 +412,6 @@ rmw_shutdown(rmw_context_t * context)
   if (context->impl->shm_provider.has_value()) {
     z_drop(z_move(context->impl->shm_provider.value()));
   }
-  // Close the zenoh session
-  if (z_close(z_move(context->impl->session), NULL) != Z_OK) {
-    RMW_SET_ERROR_MSG("Error while closing zenoh session");
-    return RMW_RET_ERROR;
-  }
 
   context->impl->is_shutdown = true;
 
diff --git a/rmw_zenoh_cpp/src/rmw_zenoh.cpp b/rmw_zenoh_cpp/src/rmw_zenoh.cpp
index bd71b266..c9fa1f7b 100644
--- a/rmw_zenoh_cpp/src/rmw_zenoh.cpp
+++ b/rmw_zenoh_cpp/src/rmw_zenoh.cpp
@@ -284,6 +284,9 @@ rmw_create_node(
   node->context = context;
   node->data = node_data;
 
+  // Add this node to the set of session_nodes.
+  context->impl->session_nodes_.insert(node);
+
   free_token.cancel();
   free_node_data.cancel();
   destruct_node_data.cancel();
@@ -320,6 +323,17 @@ rmw_destroy_node(rmw_node_t * node)
     allocator->deallocate(node_data, allocator->state);
   }
 
+  // Erase the node from the set of session_nodes and close the Zenoh
+  // session if this is the last node.
+  node->context->impl->session_nodes_.erase(node);
+  if (node->context->impl->session_nodes_.empty()) {
+    // Close the zenoh session
+    if (z_close(z_move(node->context->impl->session), NULL) != Z_OK) {
+      RMW_SET_ERROR_MSG("Error while closing zenoh session");
+      return RMW_RET_ERROR;
+    }
+  }
+
   allocator->deallocate(const_cast<char *>(node->namespace_), allocator->state);
   allocator->deallocate(const_cast<char *>(node->name), allocator->state);
   allocator->deallocate(node, allocator->state);