From c956e958e0a386f24e6865ad62ee4fe640f93b18 Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Wed, 17 Apr 2024 12:42:01 +0200
Subject: zen startup hardening (#49)

- Feature: `zen up` command improvements
	- --`port` allows you to specify a base port when starting an instance
	- --`base-dir` allows you to specify a base directory for the zenserver executable if it is not located next to the zen.exe executable
- Feature: `zen down`
	- --`port` allows you to specify a base port when shutting down an instance
	- --`base-dir` allows you to specify a base directory for the zenserver executable if it is not located next to the zen.exe executable
	- --`force` if regular shutdown fails it tries to find a running zenserver.exe process and terminate it
	- If it fails to attach to the running server it now waits for it to exit when setting the RequestExit shared memory flag
- Improvement: zenserver now checks the RequestExit flag in the shared memory and exist gracefully if it is set
- Improvement: When adding a sponsor process to a running zenserver instance, we wait for it to be picked up from the shared memory section to determine success/fail
---
 src/zenutil/zenserverprocess.cpp | 84 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 74 insertions(+), 10 deletions(-)

(limited to 'src/zenutil/zenserverprocess.cpp')

diff --git a/src/zenutil/zenserverprocess.cpp b/src/zenutil/zenserverprocess.cpp
index 7725d0af6..384371df3 100644
--- a/src/zenutil/zenserverprocess.cpp
+++ b/src/zenutil/zenserverprocess.cpp
@@ -9,6 +9,7 @@
 #include <zencore/session.h>
 #include <zencore/string.h>
 #include <zencore/thread.h>
+#include <zencore/timer.h>
 
 #include <atomic>
 
@@ -240,11 +241,15 @@ ZenServerState::Lookup(int DesiredListenPort)
 {
 	for (int i = 0; i < m_MaxEntryCount; ++i)
 	{
-		if (m_Data[i].DesiredListenPort == DesiredListenPort)
+		uint16_t EntryPort = m_Data[i].DesiredListenPort;
+		if (EntryPort != 0)
 		{
-			if (IsProcessRunning(m_Data[i].Pid))
+			if (DesiredListenPort == 0 || (EntryPort == DesiredListenPort))
 			{
-				return &m_Data[i];
+				if (IsProcessRunning(m_Data[i].Pid))
+				{
+					return &m_Data[i];
+				}
 			}
 		}
 	}
@@ -352,27 +357,56 @@ ZenServerState::ZenServerEntry::SignalShutdownRequest()
 	Flags |= uint16_t(FlagsEnum::kShutdownPlease);
 }
 
+bool
+ZenServerState::ZenServerEntry::IsShutdownRequested() const
+{
+	return (Flags.load() & static_cast<uint16_t>(FlagsEnum::kShutdownPlease)) != 0;
+}
+
 void
 ZenServerState::ZenServerEntry::SignalReady()
 {
 	Flags |= uint16_t(FlagsEnum::kIsReady);
 }
 
+bool
+ZenServerState::ZenServerEntry::IsReady() const
+{
+	return (Flags.load() & static_cast<uint16_t>(FlagsEnum::kIsReady)) != 0;
+}
+
 bool
 ZenServerState::ZenServerEntry::AddSponsorProcess(uint32_t PidToAdd)
 {
-	for (std::atomic<uint32_t>& PidEntry : SponsorPids)
+	uint32_t ServerPid	   = Pid.load();
+	auto	 WaitForPickup = [&](uint32_t AddedSlotIndex) {
+		Stopwatch Timer;
+		while (SponsorPids[AddedSlotIndex] == PidToAdd)
+		{
+			// Sponsor processes are checked every second, so 2 second wait time should be enough
+			if (Timer.GetElapsedTimeMs() > 2000)
+			{
+				return false;
+			}
+			if (!IsProcessRunning(ServerPid))
+			{
+				return false;
+			}
+			Sleep(100);
+		}
+		return true;
+	};
+	for (uint32_t SponsorIndex = 0; SponsorIndex < 8; SponsorIndex++)
 	{
-		if (PidEntry.load(std::memory_order_relaxed) == PidToAdd)
+		if (SponsorPids[SponsorIndex].load(std::memory_order_relaxed) == PidToAdd)
 		{
-			// Success, the because pid is already in the list
-			return true;
+			return WaitForPickup(SponsorIndex);
 		}
 		uint32_t Expected = 0;
-		if (PidEntry.compare_exchange_strong(Expected, PidToAdd))
+		if (SponsorPids[SponsorIndex].compare_exchange_strong(Expected, PidToAdd))
 		{
 			// Success!
-			return true;
+			return WaitForPickup(SponsorIndex);
 		}
 	}
 
@@ -661,7 +695,13 @@ ZenServerInstance::AttachToRunningServer(int BasePort)
 	}
 	else
 	{
-		State.Snapshot([&](const ZenServerState::ZenServerEntry& InEntry) { Entry = &InEntry; });
+		State.Snapshot([&](const ZenServerState::ZenServerEntry& InEntry) {
+			if (IsProcessRunning(InEntry.Pid.load()))
+			{
+				ZEN_INFO("Found entry pid {}, baseport {}", InEntry.Pid.load(), InEntry.DesiredListenPort.load());
+				Entry = &InEntry;
+			}
+		});
 	}
 
 	if (!Entry)
@@ -670,6 +710,8 @@ ZenServerInstance::AttachToRunningServer(int BasePort)
 		throw std::runtime_error("No server found");
 	}
 
+	ZEN_INFO("Found entry pid {}, baseport {}", Entry->Pid.load(), Entry->DesiredListenPort.load());
+
 	std::error_code Ec;
 	m_Process.Initialize(Entry->Pid, Ec);
 
@@ -786,4 +828,26 @@ ZenServerInstance::IsRunning()
 	return m_Process.IsRunning();
 }
 
+bool
+ZenServerInstance::Terminate()
+{
+	const std::filesystem::path BaseDir	   = m_Env.ProgramBaseDir();
+	const std::filesystem::path Executable = BaseDir / "zenserver" ZEN_EXE_SUFFIX_LITERAL;
+	ProcessHandle				RunningProcess;
+	std::error_code				Ec = FindProcess(Executable, RunningProcess);
+	if (Ec)
+	{
+		throw std::system_error(Ec, fmt::format("failed to look up running server executable '{}'", Executable));
+	}
+	if (RunningProcess.IsValid())
+	{
+		if (RunningProcess.Terminate(0))
+		{
+			return true;
+		}
+		return false;
+	}
+	return true;
+}
+
 }  // namespace zen
-- 
cgit v1.2.3