1 files changed, 66 insertions, 21 deletions

diff --git a/zenserver/zenserver.cpp b/zenserver/zenserver.cpp
index 3b56d8683..aa4a42fd7 100644
--- a/zenserver/zenserver.cpp
+++ b/zenserver/zenserver.cpp
@@ -84,9 +84,12 @@
 
 class ZenServer
 {
+	ZenServerState::ZenServerEntry* m_ServerEntry = nullptr;
+
 public:
-	void Initialize(ZenServiceConfig& ServiceConfig, int BasePort, int ParentPid)
+	void Initialize(ZenServiceConfig& ServiceConfig, int BasePort, int ParentPid, ZenServerState::ZenServerEntry* ServerEntry)
 	{
+		m_ServerEntry = ServerEntry;
 		using namespace fmt::literals;
 		ZEN_INFO(ZEN_APP_NAME " initializing");
 
@@ -94,23 +97,29 @@ public:
 
 		if (ParentPid)
 		{
-			m_Process.Initialize(ParentPid);
+			zen::ProcessHandle OwnerProcess;
+			OwnerProcess.Initialize(ParentPid);
 
-			if (!m_Process.IsValid())
+			if (!OwnerProcess.IsValid())
 			{
 				ZEN_WARN("Unable to initialize process handle for specified parent pid #{}", ParentPid);
+
+				// If the pid is not reachable should we just shut down immediately? the intended owner process
+				// could have been killed or somehow crashed already
 			}
 			else
 			{
 				ZEN_INFO("Using parent pid #{} to control process lifetime", ParentPid);
 			}
+
+			m_ProcessMonitor.AddPid(ParentPid);
 		}
 
 		// Initialize/check mutex based on base port
 
 		std::string MutexName = "zen_{}"_format(BasePort);
 
-		if (zen::NamedMutex::Exists(MutexName) || (m_ServerMutex.Create(MutexName) == false))
+		if (zen::NamedMutex::Exists(MutexName) || ((m_ServerMutex.Create(MutexName) == false)))
 		{
 			throw std::runtime_error("Failed to create mutex '{}' - is another instance already running?"_format(MutexName).c_str());
 		}
@@ -267,7 +276,7 @@ public:
 
 	void Run()
 	{
-		if (m_Process.IsValid())
+		if (m_ProcessMonitor.IsActive())
 		{
 			EnqueueTimer();
 		}
@@ -282,7 +291,7 @@ public:
 			ZEN_INFO("        \\/   \\/     \\/          \\/                         \\/ ");
 		}
 
-		ZEN_INFO(ZEN_APP_NAME " now running");
+		ZEN_INFO(ZEN_APP_NAME " now running (pid: {})", zen::GetCurrentProcessId());
 
 #if USE_SENTRY
 		sentry_clear_modulecache();
@@ -332,13 +341,33 @@ public:
 
 	void CheckOwnerPid()
 	{
-		if (m_Process.IsRunning())
+		// Pick up any new "owner" processes
+
+		std::set<uint32_t> AddedPids;
+
+		for (auto& PidEntry : m_ServerEntry->SponsorPids)
+		{
+			if (uint32_t ThisPid = PidEntry.load(std::memory_order::memory_order_relaxed))
+			{
+				if (PidEntry.compare_exchange_strong(ThisPid, 0))
+				{
+					if (AddedPids.insert(ThisPid).second)
+					{
+						m_ProcessMonitor.AddPid(ThisPid);
+
+						ZEN_INFO("added process with pid #{} as a sponsor process", ThisPid);
+					}
+				}
+			}
+		}
+
+		if (m_ProcessMonitor.IsRunning())
 		{
 			EnqueueTimer();
 		}
 		else
 		{
-			ZEN_INFO(ZEN_APP_NAME " exiting since parent process id {} is gone", m_Process.Pid());
+			ZEN_INFO(ZEN_APP_NAME " exiting since sponsor processes are all gone");
 
 			RequestExit(0);
 		}
@@ -366,7 +395,7 @@ private:
 	std::jthread		  m_IoRunner;
 	asio::io_context	  m_IoContext;
 	asio::steady_timer	  m_PidCheckTimer{m_IoContext};
-	zen::ProcessHandle	  m_Process;
+	zen::ProcessMonitor	  m_ProcessMonitor;
 	zen::NamedMutex		  m_ServerMutex;
 
 	zen::Ref<zen::HttpServer>						 m_Http;
@@ -411,27 +440,41 @@ main(int argc, char* argv[])
 	auto _ = zen::MakeGuard([&] { sentry_close(); });
 #endif
 
-	// Prototype config system, let's see how this pans out
+	try
+	{
+		// Prototype config system, we'll see how this pans out
+		//
+		// TODO: we need to report any parse errors here
 
-	ParseServiceConfig(GlobalOptions.DataDir, /* out */ ServiceConfig);
+		ParseServiceConfig(GlobalOptions.DataDir, /* out */ ServiceConfig);
 
-	ZEN_INFO("zen cache server starting on port {}", GlobalOptions.BasePort);
+		ZEN_INFO("zen cache server starting on port {}", GlobalOptions.BasePort);
 
-	try
-	{
 		ZenServerState ServerState;
 		ServerState.Initialize();
 		ServerState.Sweep();
 
-		if (ZenServerState::ZenServerEntry* Entry = ServerState.Lookup(GlobalOptions.BasePort))
+		ZenServerState::ZenServerEntry* Entry = ServerState.Lookup(GlobalOptions.BasePort);
+
+		if (Entry)
 		{
 			// Instance already running for this port? Should double check pid
 
 			ZEN_WARN("Looks like there is already a process listening to this port (pid: {})", Entry->Pid);
+
+			if (GlobalOptions.OwnerPid)
+			{
+				Entry->AddSponsorProcess(GlobalOptions.OwnerPid);
+
+				std::exit(0);
+			}
 		}
-		else
+
+		Entry = ServerState.Register(GlobalOptions.BasePort);
+
+		if (GlobalOptions.OwnerPid)
 		{
-			ServerState.Register(GlobalOptions.BasePort);
+			Entry->AddSponsorProcess(GlobalOptions.OwnerPid);
 		}
 
 		std::unique_ptr<std::thread>	 ShutdownThread;
@@ -445,15 +488,17 @@ main(int argc, char* argv[])
 		Server.SetDataRoot(GlobalOptions.DataDir);
 		Server.SetTestMode(GlobalOptions.IsTest);
 		Server.SetDedicatedMode(GlobalOptions.IsDedicated);
-		Server.Initialize(ServiceConfig, GlobalOptions.BasePort, GlobalOptions.OwnerPid);
+		Server.Initialize(ServiceConfig, GlobalOptions.BasePort, GlobalOptions.OwnerPid, Entry);
 
 		// Monitor shutdown signals
 
 		ShutdownThread.reset(new std::thread{[&] {
 			ZEN_INFO("shutdown monitor thread waiting for shutdown signal '{}'", ShutdownEventName);
-			ShutdownEvent->Wait();
-			ZEN_INFO("shutdown signal received");
-			Server.RequestExit(0);
+			if (ShutdownEvent->Wait())
+			{
+				ZEN_INFO("shutdown signal received");
+				Server.RequestExit(0);
+			}
 		}});
 
 		// If we have a parent process, establish the mechanisms we need