Remove systems if their parent entity is removed (#2232)

n particular if a user despawns an entity, the associated plugin gets removed. This should prevent issues like #2165. TBH I'm not sure if this is the right way forward as a system should technically be able to access any entity in a traditional ECS.

The PR has now been reworked and greatly simplified. All we do is stop all worker threads if an entity is removed and then recreate remaining threads.

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 =
@@ -922,6 +926,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"
@@ -122,7 +125,9 @@ size_t SystemManager::ActivatePendingSystems()
       this->systemsUpdate.push_back(system.update);
 
     if (system.postupdate)
+    {
       this->systemsPostupdate.push_back(system.postupdate);
+    }
   }
 
   this->pendingSystems.clear();
@@ -409,3 +414,131 @@ 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;
+  }
+
+  std::unordered_set<ISystemReset *> resetSystemsToBeRemoved;
+  std::unordered_set<ISystemPreUpdate *> preupdateSystemsToBeRemoved;
+  std::unordered_set<ISystemUpdate *> updateSystemsToBeRemoved;
+  std::unordered_set<ISystemPostUpdate *> postupdateSystemsToBeRemoved;
+  std::unordered_set<ISystemConfigure *> configureSystemsToBeRemoved;
+  std::unordered_set<ISystemConfigureParameters *>
+    configureParametersSystemsToBeRemoved;
+  for (const auto &system : this->systems)
+  {
+    if (_ecm.IsMarkedForRemoval(system.parentEntity))
+    {
+      if (system.reset)
+      {
+        resetSystemsToBeRemoved.insert(system.reset);
+      }
+      if (system.preupdate)
+      {
+        preupdateSystemsToBeRemoved.insert(system.preupdate);
+      }
+      if (system.update)
+      {
+        updateSystemsToBeRemoved.insert(system.update);
+      }
+      if (system.postupdate)
+      {
+        postupdateSystemsToBeRemoved.insert(system.postupdate);
+        // If system with a PostUpdate is marked for removal
+        // mark all worker threads for removal.
+        _needsCleanUp = true;
+      }
+      if (system.configure)
+      {
+        configureSystemsToBeRemoved.insert(system.configure);
+      }
+      if (system.configureParameters)
+      {
+        configureParametersSystemsToBeRemoved.insert(
+          system.configureParameters);
+      }
+    }
+  }
+
+  RemoveFromVectorIf(this->systemsReset,
+    [&](const auto system) {
+      if (resetSystemsToBeRemoved.count(system)) {
+        return true;
+      }
+      return false;
+    });
+  RemoveFromVectorIf(this->systemsPreupdate,
+    [&](const auto& system) {
+      if (preupdateSystemsToBeRemoved.count(system)) {
+        return true;
+      }
+      return false;
+    });
+  RemoveFromVectorIf(this->systemsUpdate,
+    [&](const auto& system) {
+      if (updateSystemsToBeRemoved.count(system)) {
+        return true;
+      }
+      return false;
+    });
+
+  RemoveFromVectorIf(this->systemsPostupdate,
+    [&](const auto& system) {
+      if (postupdateSystemsToBeRemoved.count(system)) {
+        return true;
+      }
+      return false;
+    });
+  RemoveFromVectorIf(this->systemsConfigure,
+    [&](const auto& system) {
+      if (configureSystemsToBeRemoved.count(system)) {
+        return true;
+      }
+      return false;
+    });
+  RemoveFromVectorIf(this->systemsConfigureParameters,
+    [&](const auto& system) {
+      if (configureParametersSystemsToBeRemoved.count(system)) {
+        return true;
+      }
+      return false;
+    });
+  RemoveFromVectorIf(this->systems,
+    [&](const SystemInternal& _arg) {
+      return _ecm.IsMarkedForRemoval(_arg.parentEntity);
+    });
+
+  std::lock_guard lock(this->pendingSystemsMutex);
+  RemoveFromVectorIf(this->pendingSystems,
+    [&](const SystemInternal& _system) {
+      return _ecm.IsMarkedForRemoval(_system.parentEntity);
+    });
+}

src/SystemManager.hh

@@ -145,6 +145,14 @@ namespace gz
       /// \brief Process system messages and add systems to entities
       public: void ProcessPendingEntitySystems();
 
+      /// \brief Remove systems that are attached to removed entities
+      /// \param[in] _entityCompMgr - ECM with entities marked for removal
+      /// \param[out] _needsCleanUp - Set to true if a system with a
+      /// PostUpdate was removed, and its thread needs to be terminated
+      public: void ProcessRemovedEntities(
+        const EntityComponentManager &_entityCompMgr,
+        bool &_needsCleanUp);
+
       /// \brief Implementation for AddSystem functions that takes an SDF
       /// element. This calls the AddSystemImpl that accepts an SDF Plugin.
       /// \param[in] _system Generic representation of a system.

src/SystemManager_TEST.cc

@@ -214,6 +214,71 @@ TEST(SystemManager, AddSystemEcm)
   EXPECT_EQ(1u, systemMgr.SystemsPostUpdate().size());
 }
 
+/////////////////////////////////////////////////
+TEST(SystemManager, AddAndRemoveSystemEcm)
+{
+  auto loader = std::make_shared<SystemLoader>();
+
+  auto ecm = EntityComponentManager();
+  auto eventManager = EventManager();
+
+  auto paramRegistry = std::make_unique<
+    gz::transport::parameters::ParametersRegistry>("SystemManager_TEST");
+  SystemManager systemMgr(
+    loader, &ecm, &eventManager, std::string(), paramRegistry.get());
+
+  EXPECT_EQ(0u, systemMgr.ActiveCount());
+  EXPECT_EQ(0u, systemMgr.PendingCount());
+  EXPECT_EQ(0u, systemMgr.TotalCount());
+  EXPECT_EQ(0u, systemMgr.SystemsConfigure().size());
+  EXPECT_EQ(0u, systemMgr.SystemsPreUpdate().size());
+  EXPECT_EQ(0u, systemMgr.SystemsUpdate().size());
+  EXPECT_EQ(0u, systemMgr.SystemsPostUpdate().size());
+
+  auto configSystem = std::make_shared<SystemWithConfigure>();
+  systemMgr.AddSystem(configSystem, kNullEntity, nullptr);
+
+  auto entity = ecm.CreateEntity();
+
+  auto updateSystemWithChild = std::make_shared<SystemWithUpdates>();
+  systemMgr.AddSystem(updateSystemWithChild, entity, nullptr);
+
+  // Configure called during AddSystem
+  EXPECT_EQ(1, configSystem->configured);
+  EXPECT_EQ(1, configSystem->configuredParameters);
+
+  EXPECT_EQ(0u, systemMgr.ActiveCount());
+  EXPECT_EQ(2u, systemMgr.PendingCount());
+  EXPECT_EQ(2u, systemMgr.TotalCount());
+  EXPECT_EQ(0u, systemMgr.SystemsConfigure().size());
+  EXPECT_EQ(0u, systemMgr.SystemsPreUpdate().size());
+  EXPECT_EQ(0u, systemMgr.SystemsUpdate().size());
+  EXPECT_EQ(0u, systemMgr.SystemsPostUpdate().size());
+
+  systemMgr.ActivatePendingSystems();
+  EXPECT_EQ(2u, systemMgr.ActiveCount());
+  EXPECT_EQ(0u, systemMgr.PendingCount());
+  EXPECT_EQ(2u, systemMgr.TotalCount());
+  EXPECT_EQ(1u, systemMgr.SystemsConfigure().size());
+  EXPECT_EQ(1u, systemMgr.SystemsPreUpdate().size());
+  EXPECT_EQ(1u, systemMgr.SystemsUpdate().size());
+  EXPECT_EQ(1u, systemMgr.SystemsPostUpdate().size());
+
+  // Remove the entity
+  ecm.RequestRemoveEntity(entity);
+  bool needsCleanUp;
+  systemMgr.ProcessRemovedEntities(ecm, needsCleanUp);
+
+  EXPECT_TRUE(needsCleanUp);
+  EXPECT_EQ(1u, systemMgr.ActiveCount());
+  EXPECT_EQ(0u, systemMgr.PendingCount());
+  EXPECT_EQ(1u, systemMgr.TotalCount());
+  EXPECT_EQ(1u, systemMgr.SystemsConfigure().size());
+  EXPECT_EQ(0u, systemMgr.SystemsPreUpdate().size());
+  EXPECT_EQ(0u, systemMgr.SystemsUpdate().size());
+  EXPECT_EQ(0u, systemMgr.SystemsPostUpdate().size());
+}
+
 /////////////////////////////////////////////////
 TEST(SystemManager, AddSystemWithInfo)
 {