From 7c403e88804b8d29ec23f581d033a661c8121829 Mon Sep 17 00:00:00 2001
From: Guillaume Autran <gautran@clearpath.ai>
Date: Thu, 19 Jul 2018 13:39:30 -0400
Subject: [PATCH] Improve runtime performance of `rmw_count_XXX` functions
 (#216)

Change the runtime performance of `rmw_count_subscribers` and `rmw_count_publishers` to be O(logN) instead of linear.
---
 .../src/namespace_prefix.cpp                  |  7 +++
 .../src/namespace_prefix.hpp                  |  4 ++
 rmw_fastrtps_shared_cpp/src/rmw_count.cpp     | 51 ++++++++++++++-----
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/rmw_fastrtps_shared_cpp/src/namespace_prefix.cpp b/rmw_fastrtps_shared_cpp/src/namespace_prefix.cpp
index 179c06c13..a44fecffe 100644
--- a/rmw_fastrtps_shared_cpp/src/namespace_prefix.cpp
+++ b/rmw_fastrtps_shared_cpp/src/namespace_prefix.cpp
@@ -51,3 +51,10 @@ _strip_ros_prefix_if_exists(const std::string & topic_name)
   }
   return topic_name;
 }
+
+/// Returns the list of ros prefixes
+const std::vector<std::string> &
+_get_all_ros_prefixes()
+{
+  return _ros_prefixes;
+}
diff --git a/rmw_fastrtps_shared_cpp/src/namespace_prefix.hpp b/rmw_fastrtps_shared_cpp/src/namespace_prefix.hpp
index a230f3cc9..6cea53ab8 100644
--- a/rmw_fastrtps_shared_cpp/src/namespace_prefix.hpp
+++ b/rmw_fastrtps_shared_cpp/src/namespace_prefix.hpp
@@ -34,4 +34,8 @@ _get_ros_prefix_if_exists(const std::string & topic_name);
 /// Returns the topic name stripped of and ROS specific prefix if exists.
 std::string
 _strip_ros_prefix_if_exists(const std::string & topic_name);
+
+/// Returns the list of ros prefixes
+const std::vector<std::string> &
+_get_all_ros_prefixes();
 #endif  // NAMESPACE_PREFIX_HPP_
diff --git a/rmw_fastrtps_shared_cpp/src/rmw_count.cpp b/rmw_fastrtps_shared_cpp/src/rmw_count.cpp
index 5e1ee2956..88d151a7e 100644
--- a/rmw_fastrtps_shared_cpp/src/rmw_count.cpp
+++ b/rmw_fastrtps_shared_cpp/src/rmw_count.cpp
@@ -23,6 +23,7 @@
 #include "rmw/types.h"
 
 #include "demangle.hpp"
+#include "namespace_prefix.hpp"
 #include "rmw_fastrtps_shared_cpp/custom_participant_info.hpp"
 #include "rmw_fastrtps_shared_cpp/rmw_common.hpp"
 #include "reader_info.hpp"
@@ -49,15 +50,28 @@ __rmw_count_publishers(
     return RMW_RET_ERROR;
   }
 
-  auto impl = static_cast<CustomParticipantInfo *>(node->data);
+  *count = 0;
+  auto ros_prefixes = _get_all_ros_prefixes();
+
+  // Build the list of all possible topic FQDN
+  std::vector<std::string> topic_fqdns;
+  topic_fqdns.push_back(topic_name);
+  if (topic_name[0] == '/') {
+    std::for_each(ros_prefixes.begin(), ros_prefixes.end(),
+      [&topic_fqdns, &topic_name](const std::string & prefix) {
+        topic_fqdns.push_back(prefix + topic_name);
+      });
+  }
 
+  auto impl = static_cast<CustomParticipantInfo *>(node->data);
   WriterInfo * slave_target = impl->secondaryPubListener;
+
   slave_target->mapmutex.lock();
-  *count = 0;
-  for (const auto & it : slave_target->topicNtypes) {
-    const auto topic_fqdn = _demangle_if_ros_topic(it.first);
-    if (topic_fqdn == topic_name) {
-      *count += it.second.size();
+  // Search and sum up the publisher counts
+  for (const auto & topic_fqdn : topic_fqdns) {
+    const auto & it = slave_target->topicNtypes.find(topic_fqdn);
+    if (it != slave_target->topicNtypes.end()) {
+      *count += it->second.size();
     }
   }
   slave_target->mapmutex.unlock();
@@ -89,15 +103,28 @@ __rmw_count_subscribers(
     return RMW_RET_ERROR;
   }
 
-  CustomParticipantInfo * impl = static_cast<CustomParticipantInfo *>(node->data);
+  *count = 0;
+  auto ros_prefixes = _get_all_ros_prefixes();
+
+  // Build the list of all possible topic FQDN
+  std::vector<std::string> topic_fqdns;
+  topic_fqdns.push_back(topic_name);
+  if (topic_name[0] == '/') {
+    std::for_each(ros_prefixes.begin(), ros_prefixes.end(),
+      [&topic_fqdns, &topic_name](const std::string & prefix) {
+        topic_fqdns.push_back(prefix + topic_name);
+      });
+  }
 
+  CustomParticipantInfo * impl = static_cast<CustomParticipantInfo *>(node->data);
   ReaderInfo * slave_target = impl->secondarySubListener;
-  *count = 0;
+
   slave_target->mapmutex.lock();
-  for (const auto & it : slave_target->topicNtypes) {
-    const auto topic_fqdn = _demangle_if_ros_topic(it.first);
-    if (topic_fqdn == topic_name) {
-      *count += it.second.size();
+  // Search and sum up the subscriber counts
+  for (const auto & topic_fqdn : topic_fqdns) {
+    const auto & it = slave_target->topicNtypes.find(topic_fqdn);
+    if (it != slave_target->topicNtypes.end()) {
+      *count += it->second.size();
     }
   }
   slave_target->mapmutex.unlock();