From 6b5ec2005384ba4c1615a788b8a99c7b67b8479a Mon Sep 17 00:00:00 2001
From: "Addisu Z. Taddese" <addisu@openrobotics.org>
Date: Wed, 31 Jul 2024 16:25:53 -0500
Subject: [PATCH] Improve signal handling (#2501)

Signed-off-by: Addisu Z. Taddese <addisu@openrobotics.org>
---
 src/Server.cc        | 5 +++++
 src/ServerPrivate.cc | 8 ++++++++
 src/ServerPrivate.hh | 4 ++++
 3 files changed, 17 insertions(+)

diff --git a/src/Server.cc b/src/Server.cc
index e2286945e7..68bf0637f8 100644
--- a/src/Server.cc
+++ b/src/Server.cc
@@ -214,6 +214,11 @@ Server::Server(const ServerConfig &_config)
     this->dataPtr->AddRecordPlugin(_config);
   }
 
+  // If we've received a signal before we create entities, the Stop event
+  // won't be propagated to them. Instead, we just quit early here.
+  if (this->dataPtr->signalReceived)
+    return;
+
   this->dataPtr->CreateEntities();
 
   // Set the desired update period, this will override the desired RTF given in
diff --git a/src/ServerPrivate.cc b/src/ServerPrivate.cc
index 80df07819e..89b5b30f03 100644
--- a/src/ServerPrivate.cc
+++ b/src/ServerPrivate.cc
@@ -112,6 +112,7 @@ ServerPrivate::~ServerPrivate()
 //////////////////////////////////////////////////
 void ServerPrivate::OnSignal(int _sig)
 {
+  this->signalReceived = true;
   gzdbg << "Server received signal[" << _sig  << "]\n";
   this->Stop();
 }
@@ -130,6 +131,13 @@ void ServerPrivate::Stop()
 bool ServerPrivate::Run(const uint64_t _iterations,
     std::optional<std::condition_variable *> _cond)
 {
+  // Return early if we've received a signal right before.
+  // The ServerPrivate signal handler would set `running=false`,
+  // but we immediately would set it to true here, which will essentially ignore
+  // the signal. Since we can't reliably use the `running` variable, we return
+  // if `signalReceived` is true
+  if (this->signalReceived)
+    return false;
   this->runMutex.lock();
   this->running = true;
   if (_cond)
diff --git a/src/ServerPrivate.hh b/src/ServerPrivate.hh
index 03128c56e8..24abe078a5 100644
--- a/src/ServerPrivate.hh
+++ b/src/ServerPrivate.hh
@@ -169,6 +169,10 @@ namespace gz
       /// \brief Our signal handler.
       public: gz::common::SignalHandler sigHandler;
 
+      /// \brief Set to true from signal handler. This will be used to
+      /// terminate the server where checking `running` is not sufficient.
+      public: std::atomic<bool> signalReceived{false};
+
       /// \brief Our system loader.
       public: SystemLoaderPtr systemLoader;