Allow system plugins to control update affinity in distributed simulation

examples/worlds/levels.sdf

@@ -4234,6 +4234,16 @@
       <ref>tile_2</ref>
     </level>
 
+    <level name="level2">
+      <pose>20 0 2.5 0 0 0</pose>
+      <geometry>
+        <box>
+          <size>40 40 1000</size>
+        </box>
+      </geometry>
+      <ref>tile_3</ref>
+    </level>
+
     <level name="level3">
       <pose>20 20 2.5 0 0 0</pose>
       <geometry>

include/gz/sim/Types.hh

@@ -20,6 +20,7 @@
 #include <chrono>
 #include <cstdint>
 #include <functional>
+#include <string>
 #include <utility>
 
 #include "gz/sim/Entity.hh"
@@ -59,6 +60,14 @@ namespace gz
       /// update state when paused is true.
       // cppcheck-suppress unusedStructMember
       bool paused{true};
+
+      /// \brief True if the simulation is distributed, which means that
+      /// there will be one or more NetworkManagerSecondary instances running.
+      bool isDistributed{false};
+
+      /// \brief The namespace of the NetworkManagerSecondary if running a
+      /// distributed simulation. Empty otherwise.
+      std::string secondaryNamespace;
     };
 
     /// \brief Possible states for a component.

src/SdfEntityCreator.cc

@@ -328,13 +328,16 @@ void SdfEntityCreator::CreateEntities(const sdf::World *_world,
       ++modelIndex)
   {
     const sdf::Model *model = _world->ModelByIndex(modelIndex);
+    gzdbg << "Try Create Default Level Model: " << model->Name() << std::endl;
     if (levelEntityNames.empty() ||
         levelEntityNames.find(model->Name()) != levelEntityNames.end())
 
     {
       Entity modelEntity = this->CreateEntities(model, false);
 
       this->SetParent(modelEntity, _worldEntity);
+      gzdbg << "Created Default Level Model: " << model->Name()
+            << "[" << modelEntity << "]" << std::endl;
     }
   }
 
@@ -375,6 +378,7 @@ void SdfEntityCreator::CreateEntities(const sdf::World *_world,
   {
     std::optional<Entity> parentEntity =
       this->dataPtr->ecm->EntityByName(_ref->Data());
+    gzdbg << "Try Find Performer Parent: " << _ref->Data() << std::endl;
     if (!parentEntity)
     {
       // Performers have not been created yet. Try to create the model
@@ -385,6 +389,8 @@ void SdfEntityCreator::CreateEntities(const sdf::World *_world,
         Entity modelEntity = this->CreateEntities(model, false);
         this->SetParent(modelEntity, _worldEntity);
         this->SetParent(_entity, modelEntity);
+        gzdbg << "Created Performer Model: " << model->Name()
+              << "[" << modelEntity << "]" << std::endl;
       }
       else if (_world->ActorNameExists(_ref->Data()))
       {

src/SimulationRunner.cc

@@ -329,6 +329,12 @@ void SimulationRunner::UpdateCurrentInfo()
 {
   GZ_PROFILE("SimulationRunner::UpdateCurrentInfo");
 
+  // Distributed
+  if (this->networkMgr != nullptr)
+  {
+    this->currentInfo.isDistributed = true;
+  }
+
   // Rewind
   if (this->requestedRewind)
   {
@@ -720,6 +726,14 @@ bool SimulationRunner::Run(const uint64_t _iterations)
     if (this->networkMgr->IsSecondary())
     {
       gzdbg << "Secondary running." << std::endl;
+
+      // Create entities if set. This needs to be called before updating
+      // the systems (so cannot be called in Step).
+      if (this->createEntities)
+      {
+        this->CreateEntities();
+      }
+
       while (!this->stopReceived)
       {
         std::this_thread::sleep_for(std::chrono::milliseconds(100));
@@ -731,7 +745,9 @@ bool SimulationRunner::Run(const uint64_t _iterations)
   // Keep track of wall clock time. Only start the realTimeWatch if this
   // runner is not paused.
   if (!this->currentInfo.paused)
+  {
     this->realTimeWatch.Start();
+  }
 
   // Variables for time keeping.
   std::chrono::steady_clock::time_point startTime;
@@ -786,7 +802,9 @@ bool SimulationRunner::Run(const uint64_t _iterations)
 
   // Create the clock publisher.
   if (!this->clockPub.Valid())
+  {
     this->clockPub = this->node->Advertise<msgs::Clock>("clock");
+  }
 
   // Create the global clock publisher.
   if (!this->rootClockPub.Valid())
@@ -932,17 +950,6 @@ void SimulationRunner::Step(const UpdateInfo &_info)
   GZ_PROFILE("SimulationRunner::Step");
   this->currentInfo = _info;
 
-  // The Run method does not check for entity creation each iteration
-  // on network secondaries, so check here to ensure that performers
-  // have parents assigned.
-  if (this->networkMgr && this->networkMgr->IsSecondary())
-  {
-    if (this->createEntities)
-    {
-      this->CreateEntities();
-    }
-  }
-
   // Process new ECM state information, typically sent from the GUI after
   // a change was made to the GUI's ECM.
   this->ProcessNewWorldControlState();

src/network/NetworkManagerPrimary.cc

@@ -231,7 +231,7 @@ void NetworkManagerPrimary::PopulateAffinities(
   // Previous performer-to-secondary mapping - may need updating
   std::map<Entity, std::string> pToSPrevious;
 
-  // Updated performer-to-level mapping - used to update affinities
+  // Updated level-to-performer mapping - used to update affinities
   std::map<Entity, std::set<Entity>> lToPNew;
 
   // All performers
@@ -265,20 +265,28 @@ void NetworkManagerPrimary::PopulateAffinities(
   // First assignment: distribute levels evenly across secondaries
   if (pToSPrevious.empty())
   {
+    // Ensure assignments are unique
+    std::map<Entity, std::string> pToS;
+
+    gzdbg << "Performer to secondary first assignment" << std::endl;
     auto secondaryIt = this->secondaries.begin();
 
     for (const auto &it : lToPNew)
     {
       for (const auto &performer : it.second)
       {
-        this->SetAffinity(performer, secondaryIt->second->prefix,
-            _msg.add_affinity());
+        if (pToS.find(performer) == pToS.end())
+        {
+          pToS[performer] = secondaryIt->second->prefix;
+          this->SetAffinity(performer, secondaryIt->second->prefix,
+              _msg.add_affinity());
+        }
 
         // Remove performers as they are assigned
         allPerformers.erase(performer);
       }
 
-      // Round-robin levels
+      // Round-robin secondaries
       secondaryIt++;
       if (secondaryIt == this->secondaries.end())
       {
@@ -289,10 +297,14 @@ void NetworkManagerPrimary::PopulateAffinities(
     // Also assign level-less performers
     for (auto performer : allPerformers)
     {
-      this->SetAffinity(performer, secondaryIt->second->prefix,
-          _msg.add_affinity());
+      if (pToS.find(performer) == pToS.end())
+      {
+        pToS[performer] = secondaryIt->second->prefix;
+        this->SetAffinity(performer, secondaryIt->second->prefix,
+            _msg.add_affinity());
+      }
 
-      // Round-robin performers
+      // Round-robin secondaries
       secondaryIt++;
       if (secondaryIt == this->secondaries.end())
       {

src/network/NetworkManagerSecondary.cc

@@ -20,13 +20,16 @@
 #include "gz/sim/private_msgs/simulation_step.pb.h"
 
 #include <algorithm>
+#include <map>
 #include <string>
+#include <vector>
 
 #include <gz/common/Console.hh>
 #include <gz/common/Util.hh>
 #include <gz/common/Profiler.hh>
 
 #include "gz/sim/components/ParentEntity.hh"
+#include "gz/sim/components/PerformerAffinity.hh"
 #include "gz/sim/Conversions.hh"
 #include "gz/sim/Entity.hh"
 #include "gz/sim/EntityComponentManager.hh"
@@ -95,6 +98,10 @@ bool NetworkManagerSecondary::OnControl(const private_msgs::PeerControl &_req,
 {
   this->enableSim = _req.enable_sim();
   _resp.set_enable_sim(this->enableSim);
+
+  gzdbg << "Received msg PeerControl: enable_sim: "
+        << this->enableSim << std::endl;
+
   return true;
 }
 
@@ -111,67 +118,114 @@ void NetworkManagerSecondary::OnStep(
            << std::endl;
   }
 
-  // Update affinities
+  // Performer-to-secondary mapping. Not using a set as we do not want sorting.
+  std::map<Entity, std::vector<std::string> > pToS;
   for (int i = 0; i < _msg.affinity_size(); ++i)
   {
     const auto &affinityMsg = _msg.affinity(i);
-    const auto &entityId = affinityMsg.entity().id();
-
-    if (affinityMsg.secondary_prefix() == this->Namespace())
+    const auto &perfEntity = affinityMsg.entity().id();
+    if (std::find(pToS[perfEntity].begin(), pToS[perfEntity].end(),
+        affinityMsg.secondary_prefix()) == pToS[perfEntity].end())
     {
-      this->performers.insert(entityId);
-
-      gzmsg << "Secondary [" << this->Namespace()
-             << "] assigned affinity to performer [" << entityId << "]."
-             << std::endl;
+      pToS[perfEntity].push_back(affinityMsg.secondary_prefix());
     }
-    // If performer has been assigned to another secondary, remove it
-    else
+  }
+
+  // Update affinities
+  for (auto [perfEntity, secondaries] : pToS)
+  {
+    for (auto secondary :  secondaries)
     {
-      auto parent =
-          this->dataPtr->ecm->Component<components::ParentEntity>(entityId);
-      if (parent != nullptr)
+      if (secondary == this->Namespace())
       {
-        this->dataPtr->ecm->RequestRemoveEntity(parent->Data());
-      }
+        this->performers.insert(perfEntity);
+
+        // Set performer affinity to this secondary
+        this->dataPtr->ecm->RemoveComponent<components::PerformerAffinity>(
+            perfEntity);
+        auto newAffinity = components::PerformerAffinity(this->Namespace());
+        this->dataPtr->ecm->CreateComponent(perfEntity, newAffinity);
 
-      if (this->performers.find(entityId) != this->performers.end())
-      {
         gzmsg << "Secondary [" << this->Namespace()
-               << "] unassigned affinity to performer [" << entityId << "]."
-               << std::endl;
-        this->performers.erase(entityId);
+              << "] assigned affinity to performer ["
+              << perfEntity << "]."
+              << std::endl;
+      }
+      else
+      {
+        // If performer assigned to this secondary is also assigned to another
+        // secondary, remove from performers and the parent model from the ECM.
+        if (this->performers.find(perfEntity) != this->performers.end())
+        {
+          gzmsg << "Secondary [" << this->Namespace()
+                << "] unassigned affinity to performer ["
+                << perfEntity << "]."
+                << std::endl;
+          this->performers.erase(perfEntity);
+
+          // Remove performer affinity from this this secondary
+          this->dataPtr->ecm->RemoveComponent<components::PerformerAffinity>(
+              perfEntity);
+
+          auto parent = this->dataPtr->ecm->Component<components::ParentEntity>(
+              perfEntity);
+          if (parent != nullptr)
+          {
+            this->dataPtr->ecm->RequestRemoveEntity(parent->Data());
+            gzmsg << "Secondary [" << this->Namespace()
+                  << "] request remove model entity ["
+                  << parent->Data() << "]."
+                  << std::endl;
+          }
+          else
+          {
+            gzerr << "Secondary [" << this->Namespace()
+                  << "] should remove parent for performer ["
+                  << perfEntity << "]."
+                  << std::endl;
+          }
+        }
       }
     }
   }
 
   // Update info
   auto info = convert<UpdateInfo>(_msg.stats());
 
+  // Set the secondaryNamespace
+  info.isDistributed = true;
+  info.secondaryNamespace = this->Namespace();
+
   // Step runner
   this->dataPtr->stepFunction(info);
 
   // Update state with all the performer's entities
   std::unordered_set<Entity> entities;
-  for (const auto &perf : this->performers)
+  for (const auto &perfEntity : this->performers)
   {
     // Performer model
-    auto parent = this->dataPtr->ecm->Component<components::ParentEntity>(perf);
+    auto parent =
+        this->dataPtr->ecm->Component<components::ParentEntity>(perfEntity);
     if (parent == nullptr)
     {
-      gzerr << "Failed to get parent for performer [" << perf << "]"
-             << std::endl;
+      gzerr << "Secondary: failed to get parent for performer ["
+            << perfEntity << "]"
+            << std::endl;
       continue;
     }
-    auto modelEntity = parent->Data();
 
+    // Also include the performer's model entity
+    auto modelEntity = parent->Data();
+    entities.insert(modelEntity);
     auto children = this->dataPtr->ecm->Descendants(modelEntity);
     entities.insert(children.begin(), children.end());
   }
 
   msgs::SerializedStateMap stateMsg;
   if (!entities.empty())
+  {
     this->dataPtr->ecm->State(stateMsg, entities);
+  }
   stateMsg.set_has_one_time_component_changes(
     this->dataPtr->ecm->HasOneTimeComponentChanges());
 

src/systems/imu/Imu.cc

@@ -230,7 +230,7 @@ void ImuPrivate::AddSensor(
   // frame like NED, that must be supplied to the IMU sensor,
   // otherwise orientations are reported w.r.t to the initial
   // orientation.
-  if (data.Element()->HasElement("imu")) {
+  if (data.Element() != nullptr && data.Element()->HasElement("imu")) {
     auto imuElementPtr = data.Element()->GetElement("imu");
     if (imuElementPtr->HasElement("orientation_reference_frame")) {
       double heading = 0.0;