From ed041b29cb6f25759d7a258a72400ad2d1c81d6b Mon Sep 17 00:00:00 2001
From: ChenYing Kuo <evshary@gmail.com>
Date: Mon, 16 Dec 2024 14:38:46 +0800
Subject: [PATCH] Don't close Zenoh session while the process is terminating.

Signed-off-by: ChenYing Kuo <evshary@gmail.com>
---
 rmw_zenoh_cpp/src/detail/zenoh_utils.cpp |  9 ++++++++-
 rmw_zenoh_cpp/src/detail/zenoh_utils.hpp | 19 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/rmw_zenoh_cpp/src/detail/zenoh_utils.cpp b/rmw_zenoh_cpp/src/detail/zenoh_utils.cpp
index be5721b7..c689d8bc 100644
--- a/rmw_zenoh_cpp/src/detail/zenoh_utils.cpp
+++ b/rmw_zenoh_cpp/src/detail/zenoh_utils.cpp
@@ -24,6 +24,10 @@
 
 namespace rmw_zenoh_cpp
 {
+// Initialize the static variable in ZenohSession
+std::once_flag ZenohSession::initFlag;
+std::atomic<bool> ZenohSession::is_exiting(false);
+
 /// Loan the zenoh session.
 ///=============================================================================
 const z_loaned_session_t * ZenohSession::loan()
@@ -35,7 +39,10 @@ const z_loaned_session_t * ZenohSession::loan()
 ///=============================================================================
 ZenohSession::~ZenohSession()
 {
-  z_close(z_loan_mut(inner_), NULL);
+  // Don't close Zenoh session while the process is terminating
+  if (!is_exiting.load()) {
+    z_close(z_loan_mut(inner_), NULL);
+  }
 }
 
 ///=============================================================================
diff --git a/rmw_zenoh_cpp/src/detail/zenoh_utils.hpp b/rmw_zenoh_cpp/src/detail/zenoh_utils.hpp
index ae2aed3a..af813e8f 100644
--- a/rmw_zenoh_cpp/src/detail/zenoh_utils.hpp
+++ b/rmw_zenoh_cpp/src/detail/zenoh_utils.hpp
@@ -17,8 +17,10 @@
 
 #include <zenoh.h>
 
+#include <atomic>
 #include <chrono>
 #include <functional>
+#include <mutex>
 #include <optional>
 
 #include "rmw/types.h"
@@ -32,12 +34,27 @@ class ZenohSession final
 {
 public:
   ZenohSession(z_owned_session_t sess)
-  : inner_(sess) {}
+  : inner_(sess) {
+    std::call_once(initFlag, [] {
+      // This atexit function should be triggered before ZenohSession destruction
+      // Refer to https://en.cppreference.com/w/cpp/utility/program/exit
+      atexit([] {
+        ZenohSession::is_exiting.store(true);
+      });
+    });
+  }
   const z_loaned_session_t * loan();
   ~ZenohSession();
 
 private:
   z_owned_session_t inner_;
+  // Used to ensure atexit function is only registered once.
+  static std::once_flag initFlag;
+  // The variable is used to identify whether the process is trying to exit or not.
+  // The atexit function we registered will set the flag and prevent us from closing
+  // Zenoh Session. Zenoh API can't be used in atexit function, because Tokio context
+  // is already destroyed. It will cause panic if we do so.
+  static std::atomic<bool> is_exiting;
 };
 
 ///=============================================================================