Add support for simulation reset via a publicly callable API

include/gz/sim/Server.hh

@@ -338,6 +338,9 @@ namespace gz
       /// \brief Stop the server. This will stop all running simulations.
       public: void Stop();
 
+      /// \brief Reset all runners in this simulation
+      public: void ResetAll();
+
       /// \brief Private data
       private: std::unique_ptr<ServerPrivate> dataPtr;
     };

include/gz/sim/TestFixture.hh

@@ -96,6 +96,12 @@ class GZ_SIM_VISIBLE TestFixture
   public: TestFixture &OnPostUpdate(std::function<void(
       const UpdateInfo &, const EntityComponentManager &)> _cb);
 
+  /// \brief Wrapper around a system's update callback
+  /// \param[in] _cb Function to be called every update
+  /// \return Reference to self.
+  public: TestFixture &OnReset(std::function<void(
+      const UpdateInfo &, EntityComponentManager &)> _cb);
+
   /// \brief Finalize all the functions and add fixture to server.
   /// Finalize must be called before running the server, otherwise none of the
   /// `On*` functions will be called.

python/src/gz/sim/Server.cc

@@ -46,7 +46,9 @@ void defineSimServer(pybind11::object module)
   .def(
     "is_running",
     pybind11::overload_cast<>(&gz::sim::Server::Running, pybind11::const_),
-    "Get whether the server is running.");
+    "Get whether the server is running.")
+  .def("reset_all", &gz::sim::Server::ResetAll,
+    "Resets all simulation runners under this server.");
 }
 }  // namespace python
 }  // namespace sim

python/src/gz/sim/TestFixture.cc

@@ -83,6 +83,17 @@ defineSimTestFixture(pybind11::object module)
     ),
     pybind11::return_value_policy::reference,
     "Wrapper around a system's post-update callback"
+  )
+  .def(
+    "on_reset", WrapCallbacks(
+      [](TestFixture* self, std::function<void(
+          const UpdateInfo &, EntityComponentManager &)> _cb)
+      {
+        self->OnReset(_cb);
+      }
+    ),
+    pybind11::return_value_policy::reference,
+    "Wrapper around a system's post-update callback"
   );
   // TODO(ahcorde): This method is not compiling for the following reason:
   // The EventManager class has an unordered_map which holds a unique_ptr

src/Server.cc

@@ -485,6 +485,17 @@ bool Server::RequestRemoveEntity(const Entity _entity,
   return false;
 }
 
+//////////////////////////////////////////////////
+void Server::ResetAll()
+{
+  for (auto worldId = 0u;
+    worldId < this->dataPtr->simRunners.size();
+    worldId++)
+  {
+    this->dataPtr->simRunners[worldId]->Reset(true, false, false);
+  }
+}
+
 //////////////////////////////////////////////////
 void Server::Stop()
 {

src/SimulationRunner.cc

@@ -1661,3 +1661,17 @@ void SimulationRunner::CreateEntities(const sdf::World &_world)
   // Store the initial state of the ECM;
   this->initialEntityCompMgr.CopyFrom(this->entityCompMgr);
 }
+
+/////////////////////////////////////////////////
+void SimulationRunner::Reset(const bool _all,
+  const bool _time, const bool _model)
+{
+  WorldControl control;
+  std::lock_guard<std::mutex> lock(this->msgBufferMutex);
+  control.rewind = _all || _time;
+  if (_model)
+  {
+    gzwarn << "Model reset not supported" <<std::endl;
+  }
+  this->worldControls.push_back(control);
+}

src/SimulationRunner.hh

@@ -369,6 +369,9 @@ namespace gz
       /// \brief Set the next step to be blocking and paused.
       public: void SetNextStepAsBlockingPaused(const bool value);
 
+      /// \brief Reset the current simulation runner
+      public: void Reset(const bool all, const bool time, const bool model);
+
       /// \brief Updates the physics parameters of the simulation based on the
       /// Physics component of the world, if any.
       public: void UpdatePhysicsParams();

src/TestFixture.cc

@@ -29,7 +29,8 @@ class HelperSystem :
   public ISystemConfigure,
   public ISystemPreUpdate,
   public ISystemUpdate,
-  public ISystemPostUpdate
+  public ISystemPostUpdate,
+  public ISystemReset
 {
   // Documentation inherited
   public: void Configure(
@@ -50,6 +51,10 @@ class HelperSystem :
   public: void PostUpdate(const UpdateInfo &_info,
                 const EntityComponentManager &_ecm) override;
 
+  // Documentation inherited
+  public: void Reset(const UpdateInfo &_info,
+                            EntityComponentManager &_ecm) override;
+
   /// \brief Function to call every time  we configure a world
   public: std::function<void(const Entity &_entity,
                 const std::shared_ptr<const sdf::Element> &_sdf,
@@ -68,6 +73,10 @@ class HelperSystem :
   /// \brief Function to call every post-update
   public: std::function<void(const UpdateInfo &,
       const EntityComponentManager &)> postUpdateCallback;
+
+  /// \brief Reset callback
+  public: std::function<void(const UpdateInfo &,
+      EntityComponentManager &)> resetCallback;
 };
 
 /////////////////////////////////////////////////
@@ -105,6 +114,14 @@ void HelperSystem::PostUpdate(const UpdateInfo &_info,
     this->postUpdateCallback(_info, _ecm);
 }
 
+/////////////////////////////////////////////////
+void HelperSystem::Reset(const UpdateInfo &_info,
+    EntityComponentManager &_ecm)
+{
+  if (this->resetCallback)
+    this->resetCallback(_info, _ecm);
+}
+
 //////////////////////////////////////////////////
 class gz::sim::TestFixture::Implementation
 {
@@ -200,6 +217,16 @@ TestFixture &TestFixture::OnPostUpdate(std::function<void(
   return *this;
 }
 
+//////////////////////////////////////////////////
+TestFixture &TestFixture::OnReset(std::function<void(
+          const UpdateInfo &, EntityComponentManager &)> _cb)
+{
+  if (nullptr != this->dataPtr->helperSystem)
+    this->dataPtr->helperSystem->resetCallback = std::move(_cb);
+  return *this;
+}
+
+
 //////////////////////////////////////////////////
 std::shared_ptr<Server> TestFixture::Server() const
 {

src/TestFixture_TEST.cc

@@ -72,6 +72,8 @@ TEST_F(TestFixtureTest, Callbacks)
   unsigned int preUpdates{0u};
   unsigned int updates{0u};
   unsigned int postUpdates{0u};
+  unsigned int resets{0u};
+
   testFixture.
     OnConfigure([&](const Entity &_entity,
         const std::shared_ptr<const sdf::Element> &_sdf,
@@ -85,20 +87,33 @@ TEST_F(TestFixtureTest, Callbacks)
     {
       this->Updates(_info, _ecm);
       preUpdates++;
-      EXPECT_EQ(preUpdates, _info.iterations);
+      if (resets == 0)
+      {
+        EXPECT_EQ(preUpdates, _info.iterations);
+      }
     }).
     OnUpdate([&](const UpdateInfo &_info, EntityComponentManager &_ecm)
     {
       this->Updates(_info, _ecm);
       updates++;
-      EXPECT_EQ(updates, _info.iterations);
+      if (resets == 0)
+      {
+        EXPECT_EQ(updates, _info.iterations);
+      }
     }).
     OnPostUpdate([&](const UpdateInfo &_info,
         const EntityComponentManager &_ecm)
     {
       this->Updates(_info, _ecm);
       postUpdates++;
-      EXPECT_EQ(postUpdates, _info.iterations);
+      if (resets == 0)
+      {
+        EXPECT_EQ(postUpdates, _info.iterations);
+      }
+    }).
+    OnReset([&](const UpdateInfo &, EntityComponentManager &)
+    {
+      resets++;
     }).
     Finalize();
 
@@ -109,8 +124,19 @@ TEST_F(TestFixtureTest, Callbacks)
   EXPECT_EQ(expectedIterations, preUpdates);
   EXPECT_EQ(expectedIterations, updates);
   EXPECT_EQ(expectedIterations, postUpdates);
+  EXPECT_EQ(0u, resets);
+
+  testFixture.Server()->ResetAll();
+
+  testFixture.Server()->Run(true, expectedIterations, false);
+  EXPECT_EQ(1u, configures);
+  EXPECT_EQ(expectedIterations * 2 - 1, preUpdates);
+  EXPECT_EQ(expectedIterations * 2 - 1, updates);
+  EXPECT_EQ(expectedIterations * 2 - 1, postUpdates);
+  EXPECT_EQ(1u, resets);
 }
 
+
 /////////////////////////////////////////////////
 TEST_F(TestFixtureTest, LoadConfig)
 {