added --powercycle option (#565)

* added --powercycle option. when this is passed in the zenserver process will shut down immediately after initialization is complete. This is primarily useful when benchmarking init/cleanup but could also be used to verify/clean up disk state
* moved EmptyStandbyList code to make it accessible to more commands

7 files changed, 164 insertions, 120 deletions

diff --git a/src/zen/bench.cpp b/src/zen/bench.cpp
new file mode 100644
index 000000000..614454ed5
--- /dev/null
+++ b/src/zen/bench.cpp
@@ -0,0 +1,134 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "bench.h"
+
+#include <zenbase/zenbase.h>
+#include <zencore/except.h>
+
+#if ZEN_PLATFORM_WINDOWS
+#	include <stdio.h>
+#	include <tchar.h>
+#	include <windows.h>
+#	include <exception>
+#	include <fmt/format.h>
+
+namespace zen::bench::util {
+
+// See https://www.geoffchappell.com/studies/windows/km/ntoskrnl/api/ex/sysinfo/set.htm
+
+typedef DWORD NTSTATUS;
+
+#	define NT_SUCCESS(Status)		  (((NTSTATUS)(Status)) >= 0)
+#	define STATUS_PRIVILEGE_NOT_HELD ((NTSTATUS)0xC0000061L)
+
+typedef enum _SYSTEM_INFORMATION_CLASS
+{
+	SystemMemoryListInformation =
+		80,	 // 80, q: SYSTEM_MEMORY_LIST_INFORMATION; s: SYSTEM_MEMORY_LIST_COMMAND (requires SeProfileSingleProcessPrivilege)
+} SYSTEM_INFORMATION_CLASS;
+
+// private
+typedef enum _SYSTEM_MEMORY_LIST_COMMAND
+{
+	MemoryCaptureAccessedBits,
+	MemoryCaptureAndResetAccessedBits,
+	MemoryEmptyWorkingSets,
+	MemoryFlushModifiedList,
+	MemoryPurgeStandbyList,
+	MemoryPurgeLowPriorityStandbyList,
+	MemoryCommandMax
+} SYSTEM_MEMORY_LIST_COMMAND;
+
+BOOL
+ObtainPrivilege(HANDLE TokenHandle, LPCSTR lpName, int flags)
+{
+	LUID			 Luid;
+	TOKEN_PRIVILEGES CurrentPriv;
+	TOKEN_PRIVILEGES NewPriv;
+
+	DWORD dwBufferLength = 16;
+	if (LookupPrivilegeValueA(0, lpName, &Luid))
+	{
+		NewPriv.PrivilegeCount			 = 1;
+		NewPriv.Privileges[0].Luid		 = Luid;
+		NewPriv.Privileges[0].Attributes = 0;
+
+		if (AdjustTokenPrivileges(TokenHandle,
+								  0,
+								  &NewPriv,
+								  DWORD((LPBYTE) & (NewPriv.Privileges[1]) - (LPBYTE)&NewPriv),
+								  &CurrentPriv,
+								  &dwBufferLength))
+		{
+			CurrentPriv.PrivilegeCount			 = 1;
+			CurrentPriv.Privileges[0].Luid		 = Luid;
+			CurrentPriv.Privileges[0].Attributes = flags != 0 ? 2 : 0;
+
+			return AdjustTokenPrivileges(TokenHandle, 0, &CurrentPriv, dwBufferLength, 0, 0);
+		}
+	}
+	return FALSE;
+}
+
+typedef NTSTATUS(WINAPI* NtSetSystemInformationFn)(INT, PVOID, ULONG);
+typedef NTSTATUS(WINAPI* NtQuerySystemInformationFn)(INT, PVOID, ULONG, PULONG);
+
+void
+EmptyStandByList()
+{
+	HMODULE NtDll = LoadLibrary(L"ntdll.dll");
+	if (!NtDll)
+	{
+		zen::ThrowLastError("Could not LoadLibrary ntdll");
+	}
+
+	HANDLE hToken;
+
+	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
+	{
+		zen::ThrowLastError("Could not open current process token");
+	}
+
+	if (!ObtainPrivilege(hToken, "SeProfileSingleProcessPrivilege", 1))
+	{
+		zen::ThrowLastError("Unable to obtain SeProfileSingleProcessPrivilege");
+	}
+
+	CloseHandle(hToken);
+
+	NtSetSystemInformationFn   NtSetSystemInformation	= (NtSetSystemInformationFn)GetProcAddress(NtDll, "NtSetSystemInformation");
+	NtQuerySystemInformationFn NtQuerySystemInformation = (NtQuerySystemInformationFn)GetProcAddress(NtDll, "NtQuerySystemInformation");
+
+	if (!NtSetSystemInformation || !NtQuerySystemInformation)
+	{
+		throw std::runtime_error("Failed to look up required ntdll functions");
+	}
+
+	SYSTEM_MEMORY_LIST_COMMAND MemoryListCommand = MemoryPurgeStandbyList;
+	NTSTATUS NtStatus = NtSetSystemInformation(SystemMemoryListInformation, &MemoryListCommand, sizeof(MemoryListCommand));
+
+	if (NtStatus == STATUS_PRIVILEGE_NOT_HELD)
+	{
+		throw elevation_required_exception("Insufficient privileges to execute the memory list command");
+	}
+	else if (!NT_SUCCESS(NtStatus))
+	{
+		throw std::runtime_error(fmt::format("Unable to execute the memory list command (status={})", NtStatus));
+	}
+}
+
+}  // namespace zen::bench::util
+
+#else
+
+namespace zen::bench::util {
+
+void
+EmptyStandByList()
+{
+	return;
+}
+
+}  // namespace zen::bench::util
+
+#endif
diff --git a/src/zen/bench.h b/src/zen/bench.h
new file mode 100644
index 000000000..6c03463ef
--- /dev/null
+++ b/src/zen/bench.h
@@ -0,0 +1,16 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#pragma once
+
+#include <stdexcept>
+
+namespace zen::bench::util {
+
+void EmptyStandByList();
+
+struct elevation_required_exception : public std::runtime_error
+{
+	explicit elevation_required_exception(const std::string& What) : std::runtime_error{What} {}
+};
+
+}  // namespace zen::bench::util
diff --git a/src/zen/cmds/bench_cmd.cpp b/src/zen/cmds/bench_cmd.cpp
index d0a513a1d..5c955e980 100644
--- a/src/zen/cmds/bench_cmd.cpp
+++ b/src/zen/cmds/bench_cmd.cpp
@@ -1,6 +1,7 @@
 // Copyright Epic Games, Inc. All Rights Reserved.
 
 #include "bench_cmd.h"
+#include "bench.h"
 
 #include <zencore/except.h>
 #include <zencore/filesystem.h>
@@ -11,126 +12,6 @@
 #include <zencore/thread.h>
 #include <zencore/timer.h>
 
-#if ZEN_PLATFORM_WINDOWS
-#	include <stdio.h>
-#	include <tchar.h>
-#	include <windows.h>
-#	include <exception>
-
-namespace zen::bench::util {
-
-// See https://www.geoffchappell.com/studies/windows/km/ntoskrnl/api/ex/sysinfo/set.htm
-
-typedef DWORD NTSTATUS;
-
-#	define NT_SUCCESS(Status)		  (((NTSTATUS)(Status)) >= 0)
-#	define STATUS_PRIVILEGE_NOT_HELD ((NTSTATUS)0xC0000061L)
-
-typedef enum _SYSTEM_INFORMATION_CLASS
-{
-	SystemMemoryListInformation =
-		80,	 // 80, q: SYSTEM_MEMORY_LIST_INFORMATION; s: SYSTEM_MEMORY_LIST_COMMAND (requires SeProfileSingleProcessPrivilege)
-} SYSTEM_INFORMATION_CLASS;
-
-// private
-typedef enum _SYSTEM_MEMORY_LIST_COMMAND
-{
-	MemoryCaptureAccessedBits,
-	MemoryCaptureAndResetAccessedBits,
-	MemoryEmptyWorkingSets,
-	MemoryFlushModifiedList,
-	MemoryPurgeStandbyList,
-	MemoryPurgeLowPriorityStandbyList,
-	MemoryCommandMax
-} SYSTEM_MEMORY_LIST_COMMAND;
-
-BOOL
-ObtainPrivilege(HANDLE TokenHandle, LPCSTR lpName, int flags)
-{
-	LUID			 Luid;
-	TOKEN_PRIVILEGES CurrentPriv;
-	TOKEN_PRIVILEGES NewPriv;
-
-	DWORD dwBufferLength = 16;
-	if (LookupPrivilegeValueA(0, lpName, &Luid))
-	{
-		NewPriv.PrivilegeCount			 = 1;
-		NewPriv.Privileges[0].Luid		 = Luid;
-		NewPriv.Privileges[0].Attributes = 0;
-
-		if (AdjustTokenPrivileges(TokenHandle,
-								  0,
-								  &NewPriv,
-								  DWORD((LPBYTE) & (NewPriv.Privileges[1]) - (LPBYTE)&NewPriv),
-								  &CurrentPriv,
-								  &dwBufferLength))
-		{
-			CurrentPriv.PrivilegeCount			 = 1;
-			CurrentPriv.Privileges[0].Luid		 = Luid;
-			CurrentPriv.Privileges[0].Attributes = flags != 0 ? 2 : 0;
-
-			return AdjustTokenPrivileges(TokenHandle, 0, &CurrentPriv, dwBufferLength, 0, 0);
-		}
-	}
-	return FALSE;
-}
-
-typedef NTSTATUS(WINAPI* NtSetSystemInformationFn)(INT, PVOID, ULONG);
-typedef NTSTATUS(WINAPI* NtQuerySystemInformationFn)(INT, PVOID, ULONG, PULONG);
-
-struct elevation_required_exception : public std::runtime_error
-{
-	explicit elevation_required_exception(const std::string& What) : std::runtime_error{What} {}
-};
-
-void
-EmptyStandByList()
-{
-	HMODULE NtDll = LoadLibrary(L"ntdll.dll");
-	if (!NtDll)
-	{
-		zen::ThrowLastError("Could not LoadLibrary ntdll");
-	}
-
-	HANDLE hToken;
-
-	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
-	{
-		zen::ThrowLastError("Could not open current process token");
-	}
-
-	if (!ObtainPrivilege(hToken, "SeProfileSingleProcessPrivilege", 1))
-	{
-		zen::ThrowLastError("Unable to obtain SeProfileSingleProcessPrivilege");
-	}
-
-	CloseHandle(hToken);
-
-	NtSetSystemInformationFn   NtSetSystemInformation	= (NtSetSystemInformationFn)GetProcAddress(NtDll, "NtSetSystemInformation");
-	NtQuerySystemInformationFn NtQuerySystemInformation = (NtQuerySystemInformationFn)GetProcAddress(NtDll, "NtQuerySystemInformation");
-
-	if (!NtSetSystemInformation || !NtQuerySystemInformation)
-	{
-		throw std::runtime_error("Failed to look up required ntdll functions");
-	}
-
-	SYSTEM_MEMORY_LIST_COMMAND MemoryListCommand = MemoryPurgeStandbyList;
-	NTSTATUS NtStatus = NtSetSystemInformation(SystemMemoryListInformation, &MemoryListCommand, sizeof(MemoryListCommand));
-
-	if (NtStatus == STATUS_PRIVILEGE_NOT_HELD)
-	{
-		throw elevation_required_exception("Insufficient privileges to execute the memory list command");
-	}
-	else if (!NT_SUCCESS(NtStatus))
-	{
-		throw std::runtime_error(fmt::format("Unable to execute the memory list command (status={})", NtStatus));
-	}
-}
-
-}  // namespace zen::bench::util
-
-#endif
-
 namespace zen {
 
 BenchCommand::BenchCommand()
diff --git a/src/zenserver/config.cpp b/src/zenserver/config.cpp
index 3640abfaf..92d6c7f7d 100644
--- a/src/zenserver/config.cpp
+++ b/src/zenserver/config.cpp
@@ -960,6 +960,9 @@ ParseCliOptions(int argc, char* argv[], ZenServerOptions& ServerOptions)
 						  "Specify a snapshot of server state to mirror into the persistence root at startup",
 						  cxxopts::value<std::string>(BaseSnapshotDir));
 	options.add_options()("content-dir", "Frontend content directory", cxxopts::value<std::string>(ContentDir));
+	options.add_options()("powercycle",
+						  "Exit immediately after initialization is complete",
+						  cxxopts::value<bool>(ServerOptions.IsPowerCycle));
 	options.add_options()("abslog", "Path to log file", cxxopts::value<std::string>(AbsLogFile));
 	options.add_options()("config", "Path to Lua config file", cxxopts::value<std::string>(ConfigFile));
 	options.add_options()("write-config", "Path to output Lua config file", cxxopts::value<std::string>(OutputConfigFile));
diff --git a/src/zenserver/config.h b/src/zenserver/config.h
index c7f6e1e2a..8135bf8f0 100644
--- a/src/zenserver/config.h
+++ b/src/zenserver/config.h
@@ -142,6 +142,7 @@ struct ZenServerOptions
 	bool					 UninstallService	= false;  // Flag used to initiate service uninstall (temporary)
 	bool					 IsDebug			= false;
 	bool					 IsCleanStart		= false;  // Indicates whether all state should be wiped on startup or not
+	bool					 IsPowerCycle		= false;  // When true, the process shuts down immediately after initialization
 	bool					 IsTest				= false;
 	bool					 IsDedicated		= false;  // Indicates a dedicated/shared instance, with larger resource requirements
 	bool					 ShouldCrash		= false;  // Option for testing crash handling
diff --git a/src/zenserver/zenserver.cpp b/src/zenserver/zenserver.cpp
index a50ff1b53..ba9ff4f88 100644
--- a/src/zenserver/zenserver.cpp
+++ b/src/zenserver/zenserver.cpp
@@ -117,6 +117,7 @@ ZenServer::Initialize(const ZenServerOptions& ServerOptions, ZenServerState::Zen
 	m_UseSentry				 = ServerOptions.NoSentry == false;
 	m_ServerEntry			 = ServerEntry;
 	m_DebugOptionForcedCrash = ServerOptions.ShouldCrash;
+	m_IsPowerCycle			 = ServerOptions.IsPowerCycle;
 	const int ParentPid		 = ServerOptions.OwnerPid;
 
 	if (ParentPid)
@@ -591,6 +592,13 @@ ZenServer::Run()
 
 	OnReady();
 
+	if (m_IsPowerCycle)
+	{
+		ZEN_INFO("Power cycle mode enabled -- shutting down");
+
+		RequestExit(0);
+	}
+
 	m_Http->Run(IsInteractiveMode);
 
 	SetNewState(kShuttingDown);
diff --git a/src/zenserver/zenserver.h b/src/zenserver/zenserver.h
index 7da536708..becc3437d 100644
--- a/src/zenserver/zenserver.h
+++ b/src/zenserver/zenserver.h
@@ -88,6 +88,7 @@ private:
 	ZenServerState::ZenServerEntry* m_ServerEntry	  = nullptr;
 	bool							m_IsDedicatedMode = false;
 	bool							m_TestMode		  = false;
+	bool							m_IsPowerCycle	  = false;
 	CbObject						m_RootManifest;
 	std::filesystem::path			m_DataRoot;
 	std::filesystem::path			m_ContentRoot;