Process Cost Updates in batches to prevent locking the VertexBuffer over and over again

Author: JustusBraun

rviz_mesh_tools_plugins/include/rviz_mesh_tools_plugins/MeshDisplay.hpp

@@ -85,6 +85,7 @@
 #include <message_filters/cache.h>
 
 #include <rviz_mesh_tools_plugins/RVizMessageFilter.hpp>
+#include <rviz_mesh_tools_plugins/ThreadSaveQueue.hpp>
 
 #include <rviz_rendering/mesh_loader.hpp>
 
@@ -156,6 +157,9 @@ class MeshDisplay : public rviz_common::Display
 
   /**
    * @brief Periodically called from rviz
+   *
+   * @param wall_dt Time since last update in NANOSECONDS (rviz doc says seconds)
+   * @param ros_dt Same as wall_dt but in ROS Time also NANOSECONDS
    */
   void update(float wall_dt, float ros_dt) override;
 
@@ -376,6 +380,11 @@ private Q_SLOTS:
    */
   void vertexCostsUpdateCallback(const mesh_msgs::msg::MeshVertexCostsSparseStamped::ConstSharedPtr& update);
 
+  /**
+   * @brief Apply all cached vertex cost updates
+   */
+  void applyCachedCostUpdates();
+
   /**
    * @brief Requests vertex colors from the specified service
    * @param uuid Mesh UUID
@@ -439,6 +448,11 @@ private Q_SLOTS:
   /// Cache for vertex costs updates
   std::shared_ptr<message_filters::Cache<mesh_msgs::msg::MeshVertexCostsSparseStamped>> m_costsUpdateMsgCache;
 
+  /// Cache for batching cost updates
+  ThreadSafeQueue<mesh_msgs::msg::MeshVertexCostsSparseStamped::ConstSharedPtr> m_costsUpdateCache;
+  float m_timeSinceLastUpdateApply;
+
+
   /// Counter for the received messages
   uint32_t m_messagesReceived;
   /// Uuid of the last received message

rviz_mesh_tools_plugins/include/rviz_mesh_tools_plugins/ThreadSaveQueue.hpp

@@ -0,0 +1,46 @@
+#pragma once
+
+#include <queue>
+#include <mutex>
+
+namespace rviz_mesh_tools_plugins
+{
+
+
+template <typename T>
+class ThreadSafeQueue: protected std::queue<T>
+{
+public:
+  T pop()
+  {
+    std::lock_guard lock(mutex_);
+    T elem = this->std::queue<T>::front();
+    this->std::queue<T>::pop();
+    return elem;
+  }
+
+  std::vector<T> popAll()
+  {
+    std::lock_guard lock(mutex_);
+    std::vector<T> data;
+    data.reserve(this->size());
+
+    while(!this->empty())
+    {
+      data.push_back(this->std::queue<T>::front());
+      this->std::queue<T>::pop();
+    }
+    return data;
+  }
+
+  void push(const T& elem)
+  {
+    std::lock_guard lock(mutex_);
+    this->std::queue<T>::push(elem);
+  }
+
+private:
+  std::mutex mutex_;
+};
+
+} // namespace rviz_mesh_tools_plugins

rviz_mesh_tools_plugins/src/MeshDisplay.cpp

@@ -82,6 +82,8 @@
 #include "rclcpp/executor.hpp"
 #include "rmw/validate_full_topic_name.h"
 
+#include <lvr2/attrmaps/AttrMaps.hpp>
+
 
 using namespace std::chrono_literals;
 using std::placeholders::_1;
@@ -102,6 +104,7 @@ namespace rviz_mesh_tools_plugins
 MeshDisplay::MeshDisplay() 
 : rviz_common::Display()
 , m_ignoreMsgs(false)
+, m_timeSinceLastUpdateApply(0.0)
 , m_meshQos(rclcpp::QoS(1).transient_local())
 , m_vertexColorsQos(rclcpp::QoS(5))
 , m_vertexCostsQos(rclcpp::QoS(5))
@@ -310,9 +313,18 @@ void MeshDisplay::onInitialize()
 
 void MeshDisplay::update(float wall_dt, float ros_dt) 
 {
-  (void) wall_dt;
   (void) ros_dt;
 
+  if (m_timeSinceLastUpdateApply >= 1.0e9)
+  {
+    this->applyCachedCostUpdates();
+    m_timeSinceLastUpdateApply = 0.0;
+  }
+  else
+  {
+    m_timeSinceLastUpdateApply += wall_dt;
+  }
+
   this->transformMesh();
 }
 
@@ -1107,74 +1119,127 @@ void MeshDisplay::vertexCostsUpdateCallback(
     return;
   }
 
-  // Vertex cost have to be applied to the cached layer
-  auto result = m_costCache.find(update->type);
-  // Unknown layer
-  if (result == m_costCache.end())
-  {
-    //TODO: We could create an empty layer
-    return;
-  }
+  this->m_costsUpdateCache.push(update);
+}
+
+void MeshDisplay::applyCachedCostUpdates()
+{
+
+  const auto updates = m_costsUpdateCache.popAll();
+
+  RCLCPP_DEBUG(rclcpp::get_logger("rviz_mesh_tools_plugins"), "Processing %lu updates", updates.size());
 
-  // Update data in the cache
-  const auto& data = update->mesh_vertex_costs;
-  for (size_t i = 0; i < data.vertices.size(); i++)
+  // Space to combine updates per layer
+  std::map<std::string, lvr2::SparseVertexMap<float>> combined;
+  
+  for (const auto& update: updates)
   {
-    result->second[data.vertices[i]] = data.costs[i];
+    // Vertex cost have to be applied to the cached layer
+    auto result = combined.find(update->type);
+    // Unknown layer
+    if (result == combined.end())
+    {
+      result = combined.insert({update->type, lvr2::SparseVertexMap<float>()}).first;
+    }
+
+    // Update data in the cache
+    const auto& data = update->mesh_vertex_costs;
+    for (size_t i = 0; i < data.vertices.size(); i++)
+    {
+      result->second.insert(lvr2::VertexHandle(data.vertices[i]), data.costs[i]);
+    }
   }
 
-  // Update the mesh if we updated the current visual
-  if (update->type == m_selectVertexCostMap->getStdString())
+  for (const auto& [type, update]: combined)
   {
-    if (m_costUseCustomLimits->getBool())
+    // Vertex cost have to be applied to the cached layer
+    auto result = m_costCache.find(type);
+    // Unknown layer
+    if (result == m_costCache.end())
     {
-      if (m_costCache.count(m_selectVertexCostMap->getStdString()) != 0)
-      {
-        std::shared_ptr<MeshVisual> visual = getLatestVisual();
-        if (visual)
-        {
-          visual->updateVertexCosts(
-            data.vertices,
-            data.costs,
-            m_costColorType->getOptionInt(),
-            m_costLowerLimit->getFloat(),
-            m_costUpperLimit->getFloat()
-          );
-        }
-      }
+      //TODO: We could create an empty layer
+      return;
     }
-    else
+
+    // Update data in the cache
+    for (lvr2::VertexHandle v: update)
+    {
+      result->second[v.idx()] = update[v];
+    }
+
+    // Update the mesh if we updated the current visual
+    if (type == m_selectVertexCostMap->getStdString())
     {
-      if (m_costCache.count(m_selectVertexCostMap->getStdString()) != 0)
+      if (m_costUseCustomLimits->getBool())
       {
-        // The correct way is the get the limits from the whole layer
-        // TODO: Cache this with the layer itself
-        float maxCost = std::numeric_limits<float>::min();
-        float minCost = std::numeric_limits<float>::max();
-        for (float cost : result->second)
+        if (m_costCache.count(m_selectVertexCostMap->getStdString()) != 0)
         {
-          if (std::isfinite(cost) && cost > maxCost)
-            maxCost = cost;
-          if (std::isfinite(cost) && cost < minCost)
-            minCost = cost;
+          std::shared_ptr<MeshVisual> visual = getLatestVisual();
+          if (visual)
+          {
+            std::vector<uint32_t> vertices;
+            std::vector<float> costs;
+            vertices.reserve(update.numValues());
+            costs.reserve(update.numValues());
+            for (lvr2::VertexHandle v: update)
+            {
+              vertices.push_back(v.idx());
+              costs.push_back(update[v]);
+            }
+
+            visual->updateVertexCosts(
+              vertices,
+              costs,
+              m_costColorType->getOptionInt(),
+              m_costLowerLimit->getFloat(),
+              m_costUpperLimit->getFloat()
+            );
+          }
         }
-
-
-        std::shared_ptr<MeshVisual> visual = getLatestVisual();
-        if (visual)
+      }
+      else
+    {
+        if (m_costCache.count(m_selectVertexCostMap->getStdString()) != 0)
         {
-          visual->updateVertexCosts(
-            data.vertices,
-            data.costs,
-            m_costColorType->getOptionInt(),
-            minCost,
-            maxCost
-          );
+          // The correct way is the get the limits from the whole layer
+          // TODO: Cache this with the layer itself
+          float maxCost = std::numeric_limits<float>::min();
+          float minCost = std::numeric_limits<float>::max();
+          for (float cost : result->second)
+          {
+            if (std::isfinite(cost) && cost > maxCost)
+              maxCost = cost;
+            if (std::isfinite(cost) && cost < minCost)
+              minCost = cost;
+          }
+
+
+          std::shared_ptr<MeshVisual> visual = getLatestVisual();
+          if (visual)
+          {
+            std::vector<uint32_t> vertices;
+            std::vector<float> costs;
+            vertices.reserve(update.numValues());
+            costs.reserve(update.numValues());
+            for (lvr2::VertexHandle v: update)
+            {
+              vertices.push_back(v.idx());
+              costs.push_back(update[v]);
+            }
+            visual->updateVertexCosts(
+              vertices,
+              costs,
+              m_costColorType->getOptionInt(),
+              minCost,
+              maxCost
+            );
+          }
         }
       }
+      updateMeshMaterial();
     }
-    updateMeshMaterial();
   }
+
 }
 
 void MeshDisplay::requestVertexColors(std::string uuid)