Remove systems if their parent entity is removed

include/gz/sim/EntityComponentManager.hh

@@ -828,6 +828,9 @@ namespace gz
       friend class GuiRunner;
       friend class SimulationRunner;
 
+      // Make SystemManager friend so it has access to removals
+      friend class SystemManager;
+
       // Make network managers friends so they have control over component
       // states. Like the runners, the managers are internal.
       friend class NetworkManagerPrimary;

src/SimulationRunner.cc

@@ -33,6 +33,7 @@
 #include <gz/msgs/world_stats.pb.h>
 
 #include <sdf/Root.hh>
+#include <vector>
 
 #include "gz/common/Profiler.hh"
 #include "gz/sim/components/Model.hh"
@@ -533,12 +534,15 @@ void SimulationRunner::ProcessSystemQueue()
 {
   auto pending = this->systemMgr->PendingCount();
 
-  if (0 == pending)
+  if (0 == pending && !this->threadsNeedCleanUp)
     return;
 
-  // If additional systems are to be added, stop the existing threads.
+  // If additional systems are to be added or have been removed,
+  // stop the existing threads.
   this->StopWorkerThreads();
 
+  this->threadsNeedCleanUp = false;
+
   this->systemMgr->ActivatePendingSystems();
 
   unsigned int threadCount =
@@ -567,10 +571,12 @@ void SimulationRunner::ProcessSystemQueue()
         this->postUpdateStartBarrier->Wait();
         if (this->postUpdateThreadsRunning)
         {
-          system->PostUpdate(this->currentInfo, this->entityCompMgr);
+          system->PostUpdate(this->currentInfo,
+            this->entityCompMgr);
         }
         this->postUpdateStopBarrier->Wait();
       }
+
       gzdbg << "Exiting postupdate worker thread ("
         << id << ")" << std::endl;
     }));
@@ -598,13 +604,13 @@ void SimulationRunner::UpdateSystems()
   {
     GZ_PROFILE("PreUpdate");
     for (auto& system : this->systemMgr->SystemsPreUpdate())
-      system->PreUpdate(this->currentInfo, this->entityCompMgr);
+      system.system->PreUpdate(this->currentInfo, this->entityCompMgr);
   }
 
   {
     GZ_PROFILE("Update");
     for (auto& system : this->systemMgr->SystemsUpdate())
-      system->Update(this->currentInfo, this->entityCompMgr);
+      system.system->Update(this->currentInfo, this->entityCompMgr);
   }
 
   {
@@ -889,6 +895,7 @@ void SimulationRunner::Step(const UpdateInfo &_info)
   // Update all the systems.
   this->UpdateSystems();
 
+
   if (!this->Paused() && this->requestedRunToSimTime &&
        this->requestedRunToSimTime.value() > this->simTimeEpoch &&
        this->currentInfo.simTime >= this->requestedRunToSimTime.value())
@@ -922,6 +929,8 @@ void SimulationRunner::Step(const UpdateInfo &_info)
   this->ProcessRecreateEntitiesCreate();
 
   // Process entity removals.
+  this->systemMgr->ProcessRemovedEntities(this->entityCompMgr,
+    this->threadsNeedCleanUp);
   this->entityCompMgr.ProcessRemoveEntityRequests();
 
   // Process components removals

src/SimulationRunner.hh

@@ -543,6 +543,9 @@ namespace gz
       /// at the appropriate time.
       private: std::unique_ptr<msgs::WorldControlState> newWorldControlState;
 
+      /// \brief Set if we need to remove systems due to entity removal
+      private: bool threadsNeedCleanUp;
+
       private: bool resetInitiated{false};
       friend class LevelManager;
     };

src/SimulationRunner_TEST.cc

@@ -16,7 +16,6 @@
 */
 
 #include <gtest/gtest.h>
-
 #include <tinyxml2.h>
 
 #include <gz/msgs/clock.pb.h>
@@ -111,7 +110,6 @@ void rootClockCb(const msgs::Clock &_msg)
   rootClockMsgs.push_back(_msg);
 }
 
-
 /////////////////////////////////////////////////
 TEST_P(SimulationRunnerTest, CreateEntities)
 {
@@ -1484,8 +1482,7 @@ TEST_P(SimulationRunnerTest,
   EXPECT_TRUE(runner.EntityCompMgr().EntityHasComponentType(sphereEntity,
       componentId)) << componentId;
 
-  // Remove entities that have plugin - this is not unloading or destroying
-  // the plugin though!
+  // Remove entities that have plugin
   auto entityCount = runner.EntityCompMgr().EntityCount();
   const_cast<EntityComponentManager &>(
       runner.EntityCompMgr()).RequestRemoveEntity(boxEntity);
@@ -1533,8 +1530,16 @@ TEST_P(SimulationRunnerTest,
   SimulationRunner runner(rootWithout.WorldByIndex(0), systemLoader,
       serverConfig);
 
-  // 1 model plugin from SDF and 2 world plugins from config
-  ASSERT_EQ(3u, runner.SystemCount());
+  // 1 model plugin from SDF and 1 world plugin from config
+  // and 1 model plugin from theconfig
+  EXPECT_EQ(3u, runner.SystemCount());
+  runner.SetPaused(false);
+  runner.Run(1);
+
+  // Remove the model. Only 1 world plugin should remain.
+  EXPECT_TRUE(runner.RequestRemoveEntity("box"));
+  runner.Run(2);
+  EXPECT_EQ(1u, runner.SystemCount());
 }
 
 /////////////////////////////////////////////////

src/SystemManager.cc

@@ -16,10 +16,13 @@
 */
 
 #include <list>
+#include <mutex>
 #include <set>
+#include <unordered_set>
 
 #include <gz/common/StringUtils.hh>
 
+#include "SystemInternal.hh"
 #include "gz/sim/components/SystemPluginInfo.hh"
 #include "gz/sim/Conversions.hh"
 #include "SystemManager.hh"
@@ -107,22 +110,32 @@ size_t SystemManager::ActivatePendingSystems()
     this->systems.push_back(system);
 
     if (system.configure)
-      this->systemsConfigure.push_back(system.configure);
+      this->systemsConfigure.emplace_back(
+        system.parentEntity,
+        system.configure);
 
     if (system.configureParameters)
-      this->systemsConfigureParameters.push_back(system.configureParameters);
+      this->systemsConfigureParameters.emplace_back(
+        system.parentEntity,
+        system.configureParameters);
 
     if (system.reset)
-      this->systemsReset.push_back(system.reset);
+      this->systemsReset.emplace_back(system.parentEntity,
+        system.reset);
 
     if (system.preupdate)
-      this->systemsPreupdate.push_back(system.preupdate);
+      this->systemsPreupdate.emplace_back(system.parentEntity,
+        system.preupdate);
 
     if (system.update)
-      this->systemsUpdate.push_back(system.update);
+      this->systemsUpdate.emplace_back(system.parentEntity,
+        system.update);
 
     if (system.postupdate)
+    {
       this->systemsPostupdate.push_back(system.postupdate);
+      this->postUpdateParents.push_back(system.parentEntity);
+    }
   }
 
   this->pendingSystems.clear();
@@ -156,6 +169,7 @@ void SystemManager::Reset(const UpdateInfo &_info, EntityComponentManager &_ecm)
   this->systemsPreupdate.clear();
   this->systemsUpdate.clear();
   this->systemsPostupdate.clear();
+  this->postUpdateParents.clear();
 
   std::vector<PluginInfo> pluginsToBeLoaded;
 
@@ -277,37 +291,42 @@ void SystemManager::AddSystemImpl(
 }
 
 //////////////////////////////////////////////////
-const std::vector<ISystemConfigure *>& SystemManager::SystemsConfigure()
+const std::vector<SystemIfaceWithParent<ISystemConfigure>>&
+  SystemManager::SystemsConfigure()
 {
   return this->systemsConfigure;
 }
 
-const std::vector<ISystemConfigureParameters *>&
+const std::vector<SystemIfaceWithParent<ISystemConfigureParameters>>&
 SystemManager::SystemsConfigureParameters()
 {
   return this->systemsConfigureParameters;
 }
 
 //////////////////////////////////////////////////
-const std::vector<ISystemReset *> &SystemManager::SystemsReset()
+const std::vector<SystemIfaceWithParent<ISystemReset>>&
+SystemManager::SystemsReset()
 {
   return this->systemsReset;
 }
 
 //////////////////////////////////////////////////
-const std::vector<ISystemPreUpdate *>& SystemManager::SystemsPreUpdate()
+const std::vector<SystemIfaceWithParent<ISystemPreUpdate>>&
+  SystemManager::SystemsPreUpdate()
 {
   return this->systemsPreupdate;
 }
 
 //////////////////////////////////////////////////
-const std::vector<ISystemUpdate *>& SystemManager::SystemsUpdate()
+const std::vector<SystemIfaceWithParent<ISystemUpdate>>&
+  SystemManager::SystemsUpdate()
 {
   return this->systemsUpdate;
 }
 
 //////////////////////////////////////////////////
-const std::vector<ISystemPostUpdate *>& SystemManager::SystemsPostUpdate()
+const std::vector<ISystemPostUpdate *>&
+  SystemManager::SystemsPostUpdate()
 {
   return this->systemsPostupdate;
 }
@@ -409,3 +428,90 @@ void SystemManager::ProcessPendingEntitySystems()
   }
   this->systemsToAdd.clear();
 }
+
+template <typename T>
+struct identity
+{
+    using type = T;
+};
+
+//////////////////////////////////////////////////
+/// TODO(arjo) - When we move to C++20 we can just use erase_if
+/// Removes all items that match the given predicate.
+/// This function runs in O(n) time and uses memory in-place
+template<typename Tp>
+void RemoveFromVectorIf(std::vector<Tp>& vec,
+  typename identity<std::function<bool(const Tp&)>>::type pred)
+{
+  vec.erase(std::remove_if(vec.begin(), vec.end(), pred), vec.end());
+}
+
+//////////////////////////////////////////////////
+void SystemManager::ProcessRemovedEntities(
+  const EntityComponentManager &_ecm,
+  bool &_needsCleanUp)
+{
+  // Note: This function has  O(n) time when an entity is removed
+  // where n is number of systems. Ideally we would only iterate
+  // over entities marked for removal but that would involve having
+  // a key value map. This can be marked as a future improvement.
+  if (!_ecm.HasEntitiesMarkedForRemoval())
+  {
+    return;
+  }
+
+  RemoveFromVectorIf(this->systemsReset,
+    [&](const SystemIfaceWithParent<ISystemReset>& system) {
+      return _ecm.IsMarkedForRemoval(system.parent);
+    });
+  RemoveFromVectorIf(this->systemsPreupdate,
+    [&](const SystemIfaceWithParent<ISystemPreUpdate>& system) {
+      return _ecm.IsMarkedForRemoval(system.parent);
+    });
+  RemoveFromVectorIf(this->systemsUpdate,
+    [&](const SystemIfaceWithParent<ISystemUpdate>& system) {
+      return _ecm.IsMarkedForRemoval(system.parent);
+    });
+
+  std::unordered_set<ISystemPostUpdate *> markedForRemoval;
+  for (std::size_t i = 0; i < this->systemsPostupdate.size(); i++)
+  {
+    if(_ecm.IsMarkedForRemoval(postUpdateParents[i]))
+    {
+      // If system with a PostUpdate is marked for removal
+      // mark all worker threads for removal.
+      _needsCleanUp = true;
+      markedForRemoval.insert(this->systemsPostupdate[i]);
+    }
+  }
+
+  RemoveFromVectorIf(this->systemsPostupdate,
+    [&](const auto system) {
+      if (markedForRemoval.count(system)) {
+        return true;
+      }
+      return false;
+    });
+
+  RemoveFromVectorIf(this->postUpdateParents,
+    [&](const Entity& entity){
+      return _ecm.IsMarkedForRemoval(entity);
+    }
+  );
+  RemoveFromVectorIf(this->systemsConfigure,
+    [&](const SystemIfaceWithParent<ISystemConfigure>& system) {
+      return _ecm.IsMarkedForRemoval(system.parent);
+    });
+  RemoveFromVectorIf(this->systemsConfigureParameters,
+    [&](const SystemIfaceWithParent<ISystemConfigureParameters>& system) {
+      return _ecm.IsMarkedForRemoval(system.parent);
+    });
+  RemoveFromVectorIf(this->systems, [&](const SystemInternal& system) {
+      return _ecm.IsMarkedForRemoval(system.parentEntity);
+    });
+  std::lock_guard lock(this->pendingSystemsMutex);
+  RemoveFromVectorIf(this->pendingSystems,
+    [&](const SystemInternal& system) {
+      return _ecm.IsMarkedForRemoval(system.parentEntity);
+    });
+}