From ec9800eb9e8112e8c6eeeaf9feac7a32495a94fe Mon Sep 17 00:00:00 2001
From: Chen Lihui <lihui.chen@sony.com>
Date: Thu, 20 Jan 2022 10:10:14 +0800
Subject: [PATCH] fix one subscription can wait_for_message twice

Signed-off-by: Chen Lihui <lihui.chen@sony.com>
---
 rclcpp/include/rclcpp/wait_for_message.hpp   |  3 ++
 rclcpp/test/rclcpp/test_wait_for_message.cpp | 35 ++++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/rclcpp/include/rclcpp/wait_for_message.hpp b/rclcpp/include/rclcpp/wait_for_message.hpp
index 20c9df4db6..25c45ad782 100644
--- a/rclcpp/include/rclcpp/wait_for_message.hpp
+++ b/rclcpp/include/rclcpp/wait_for_message.hpp
@@ -18,6 +18,8 @@
 #include <memory>
 #include <string>
 
+#include "rcpputils/scope_exit.hpp"
+
 #include "rclcpp/node.hpp"
 #include "rclcpp/visibility_control.hpp"
 #include "rclcpp/wait_set.hpp"
@@ -54,6 +56,7 @@ bool wait_for_message(
 
   rclcpp::WaitSet wait_set;
   wait_set.add_subscription(subscription);
+  RCPPUTILS_SCOPE_EXIT(wait_set.remove_subscription(subscription); );
   wait_set.add_guard_condition(gc);
   auto ret = wait_set.wait(time_to_wait);
   if (ret.kind() != rclcpp::WaitResultKind::Ready) {
diff --git a/rclcpp/test/rclcpp/test_wait_for_message.cpp b/rclcpp/test/rclcpp/test_wait_for_message.cpp
index f0071698ca..dadc06f2da 100644
--- a/rclcpp/test/rclcpp/test_wait_for_message.cpp
+++ b/rclcpp/test/rclcpp/test_wait_for_message.cpp
@@ -72,3 +72,38 @@ TEST(TestUtilities, wait_for_message_indefinitely) {
 
   ASSERT_FALSE(received);
 }
+
+TEST(TestUtilities, wait_for_message_twice_one_sub) {
+  rclcpp::init(0, nullptr);
+
+  auto node = std::make_shared<rclcpp::Node>("wait_for_message_node3");
+
+  using MsgT = test_msgs::msg::Strings;
+  auto pub = node->create_publisher<MsgT>("wait_for_message_topic", 10);
+  auto sub = node->create_subscription<MsgT>(
+    "wait_for_message_topic", 1, [](const std::shared_ptr<const MsgT>) {});
+
+  MsgT out1;
+  MsgT out2;
+  auto received = false;
+  auto wait = std::async(
+    [&]() {
+      auto ret = rclcpp::wait_for_message(out1, sub, node->get_node_options().context(), 5s);
+      EXPECT_TRUE(ret);
+      ret = rclcpp::wait_for_message(out2, sub, node->get_node_options().context(), 5s);
+      EXPECT_TRUE(ret);
+      received = true;
+    });
+
+  for (auto i = 0u; i < 10 && received == false; ++i) {
+    pub->publish(*get_messages_strings()[0]);
+    std::this_thread::sleep_for(1s);
+  }
+
+  ASSERT_NO_THROW(wait.get());
+  ASSERT_TRUE(received);
+  EXPECT_EQ(out1, *get_messages_strings()[0]);
+  EXPECT_EQ(out2, *get_messages_strings()[0]);
+
+  rclcpp::shutdown();
+}