From d4bf8ab3afe359a8383ab1474bb6b3ef6734bbbb Mon Sep 17 00:00:00 2001
From: Barry Xu <barry.xu@sony.com>
Date: Thu, 22 Apr 2021 18:05:06 +0800
Subject: [PATCH] Updated codes based on comments

Signed-off-by: Barry Xu <barry.xu@sony.com>
---
 rclcpp/include/rclcpp/context.hpp  | 13 ++++++-----
 rclcpp/include/rclcpp/executor.hpp |  2 +-
 rclcpp/src/rclcpp/context.cpp      | 35 +++++++++++++-----------------
 3 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/rclcpp/include/rclcpp/context.hpp b/rclcpp/include/rclcpp/context.hpp
index e511f93b4b..466758fb90 100644
--- a/rclcpp/include/rclcpp/context.hpp
+++ b/rclcpp/include/rclcpp/context.hpp
@@ -23,6 +23,7 @@
 #include <typeindex>
 #include <typeinfo>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -207,17 +208,17 @@ class Context : public std::enable_shared_from_this<Context>
    */
   RCLCPP_PUBLIC
   virtual
-  OnShutdownCallbackHandle::SharedPtr
+  OnShutdownCallbackHandle::WeakPtr
   on_shutdown(OnShutdownCallbackHandle::OnShutdownCallbackType callback);
 
   /// Remove a on_shutdown callback handle registered with on_shutdown.
   /**
-   * \param[in] callback_handle the on shutdown callback handle to be removed.
+   * \param[in] callback_handle the handle to be removed.
    */
   RCLCPP_PUBLIC
   virtual
   void
-  remove_on_shutdown_callback(OnShutdownCallbackHandle::SharedPtr callback_handle);
+  remove_on_shutdown_callback(const OnShutdownCallbackHandle::WeakPtr callback_handle);
 
   /// Return the shutdown callback handles as const.
   /**
@@ -225,7 +226,7 @@ class Context : public std::enable_shared_from_this<Context>
    * the list of "on shutdown" callback handles, i.e. on_shutdown().
    */
   RCLCPP_PUBLIC
-  const std::vector<OnShutdownCallbackHandle::WeakPtr> &
+  const std::unordered_set<OnShutdownCallbackHandle::SharedPtr> &
   get_on_shutdown_callbacks() const;
 
   /// Return the shutdown callback handles
@@ -234,7 +235,7 @@ class Context : public std::enable_shared_from_this<Context>
    * the list of "on shutdown" callback handles, i.e. on_shutdown().
    */
   RCLCPP_PUBLIC
-  std::vector<OnShutdownCallbackHandle::WeakPtr> &
+  std::unordered_set<OnShutdownCallbackHandle::SharedPtr> &
   get_on_shutdown_callbacks();
 
   /// Return the internal rcl context.
@@ -315,7 +316,7 @@ class Context : public std::enable_shared_from_this<Context>
   // attempt to acquire another sub context.
   std::recursive_mutex sub_contexts_mutex_;
 
-  std::vector<OnShutdownCallbackHandle::WeakPtr> on_shutdown_callbacks_;
+  std::unordered_set<OnShutdownCallbackHandle::SharedPtr> on_shutdown_callbacks_;
   std::mutex on_shutdown_callbacks_mutex_;
 
   /// Condition variable for timed sleep (see sleep_for).
diff --git a/rclcpp/include/rclcpp/executor.hpp b/rclcpp/include/rclcpp/executor.hpp
index 8fbf2417ed..2c5a167f1a 100644
--- a/rclcpp/include/rclcpp/executor.hpp
+++ b/rclcpp/include/rclcpp/executor.hpp
@@ -574,7 +574,7 @@ class Executor
   weak_nodes_ RCPPUTILS_TSA_GUARDED_BY(mutex_);
 
   /// shutdown callback handle registered to Context
-  rclcpp::OnShutdownCallbackHandle::SharedPtr shutdown_callback_handle_;
+  rclcpp::OnShutdownCallbackHandle::WeakPtr shutdown_callback_handle_;
 };
 
 }  // namespace rclcpp
diff --git a/rclcpp/src/rclcpp/context.cpp b/rclcpp/src/rclcpp/context.cpp
index ed3054dcb3..c2bbd791d2 100644
--- a/rclcpp/src/rclcpp/context.cpp
+++ b/rclcpp/src/rclcpp/context.cpp
@@ -19,6 +19,7 @@
 #include <sstream>
 #include <string>
 #include <vector>
+#include <unordered_set>
 #include <utility>
 
 #include "rcl/init.h"
@@ -317,11 +318,8 @@ Context::shutdown(const std::string & reason)
   // call each shutdown callback
   {
     std::lock_guard<std::mutex> lock(on_shutdown_callbacks_mutex_);
-    for (const auto & callback_handle_weak_ptr : on_shutdown_callbacks_) {
-      auto callback_handle = callback_handle_weak_ptr.lock();
-      if (callback_handle) {
-        callback_handle->callback();
-      }
+    for (const auto & callback_handle : on_shutdown_callbacks_) {
+      callback_handle->callback();
     }
   }
 
@@ -348,38 +346,35 @@ Context::shutdown(const std::string & reason)
   return true;
 }
 
-rclcpp::OnShutdownCallbackHandle::SharedPtr
+rclcpp::OnShutdownCallbackHandle::WeakPtr
 Context::on_shutdown(OnShutdownCallbackHandle::OnShutdownCallbackType callback)
 {
-  std::lock_guard<std::mutex> lock(on_shutdown_callbacks_mutex_);
   auto handle = std::make_shared<OnShutdownCallbackHandle>();
   handle->callback = callback;
-  on_shutdown_callbacks_.push_back(handle);
+  {
+    std::lock_guard<std::mutex> lock(on_shutdown_callbacks_mutex_);
+    on_shutdown_callbacks_.emplace(handle);
+  }
   return handle;
 }
 
 void
-Context::remove_on_shutdown_callback(OnShutdownCallbackHandle::SharedPtr callback_handle)
+Context::remove_on_shutdown_callback(const OnShutdownCallbackHandle::WeakPtr callback_handle)
 {
-  std::lock_guard<std::mutex> lock(on_shutdown_callbacks_mutex_);
-  auto it = std::find_if(
-    on_shutdown_callbacks_.begin(),
-    on_shutdown_callbacks_.end(),
-    [callback_handle](const auto & weak_handle) {
-      return callback_handle.get() == weak_handle.lock().get();
-    });
-  if (it != on_shutdown_callbacks_.end()) {
-    on_shutdown_callbacks_.erase(it);
+  auto callback_handle_shared_ptr = callback_handle.lock();
+  if (callback_handle_shared_ptr) {
+    std::lock_guard<std::mutex> lock(on_shutdown_callbacks_mutex_);
+    on_shutdown_callbacks_.erase(callback_handle_shared_ptr);
   }
 }
 
-const std::vector<rclcpp::OnShutdownCallbackHandle::WeakPtr> &
+const std::unordered_set<rclcpp::OnShutdownCallbackHandle::SharedPtr> &
 Context::get_on_shutdown_callbacks() const
 {
   return on_shutdown_callbacks_;
 }
 
-std::vector<rclcpp::OnShutdownCallbackHandle::WeakPtr> &
+std::unordered_set<rclcpp::OnShutdownCallbackHandle::SharedPtr> &
 Context::get_on_shutdown_callbacks()
 {
   return on_shutdown_callbacks_;