Merge branch 'main' of https://github.com/EpicGames/zen

1 files changed, 217 insertions, 0 deletions

diff --git a/src/zen/cmds/bench_cmd.cpp b/src/zen/cmds/bench_cmd.cpp
new file mode 100644
index 000000000..06b8967a3
--- /dev/null
+++ b/src/zen/cmds/bench_cmd.cpp
@@ -0,0 +1,217 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "bench_cmd.h"
+
+#include <zencore/except.h>
+#include <zencore/filesystem.h>
+#include <zencore/fmtutils.h>
+#include <zencore/logging.h>
+#include <zencore/string.h>
+#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
+
+BenchCommand::BenchCommand()
+{
+	m_Options.add_options()("h,help", "Print help");
+	m_Options.add_options()("purge",
+							"Purge standby memory (system cache)",
+							cxxopts::value<bool>(m_PurgeStandbyLists)->default_value("false"));
+	m_Options.add_options()("single", "Do not spawn child processes", cxxopts::value<bool>(m_SingleProcess)->default_value("false"));
+}
+
+BenchCommand::~BenchCommand() = default;
+
+int
+BenchCommand::Run(const ZenCliOptions& GlobalOptions, int argc, char** argv)
+{
+	ZEN_UNUSED(GlobalOptions);
+
+	if (!ParseOptions(argc, argv))
+	{
+		return 0;
+	}
+
+#if ZEN_PLATFORM_WINDOWS
+	if (m_PurgeStandbyLists)
+	{
+		bool Ok = false;
+
+		zen::Stopwatch Timer;
+
+		try
+		{
+			zen::bench::util::EmptyStandByList();
+
+			Ok = true;
+		}
+		catch (zen::bench::util::elevation_required_exception&)
+		{
+			ZEN_CONSOLE("purging standby lists requires elevation. Will try launch as elevated process");
+		}
+		catch (std::exception& Ex)
+		{
+			ZEN_CONSOLE("ERROR: {}", Ex.what());
+		}
+
+		if (!Ok && !m_SingleProcess)
+		{
+			try
+			{
+				zen::CreateProcOptions Cpo;
+				Cpo.Flags = zen::CreateProcOptions::Flag_Elevated | zen::CreateProcOptions::Flag_NewConsole;
+
+				std::filesystem::path CurExe{zen::GetRunningExecutablePath()};
+
+				if (zen::CreateProcResult Cpr = zen::CreateProc(CurExe, fmt::format("bench --purge --single"), Cpo))
+				{
+					zen::ProcessHandle ProcHandle;
+					ProcHandle.Initialize(Cpr);
+
+					int ExitCode = ProcHandle.WaitExitCode();
+
+					if (ExitCode == 0)
+					{
+						Ok = true;
+					}
+					else
+					{
+						ZEN_CONSOLE("ERROR: Elevated child process failed with return code {}", ExitCode);
+					}
+				}
+			}
+			catch (std::exception& Ex)
+			{
+				ZEN_CONSOLE("ERROR: {}", Ex.what());
+			}
+		}
+
+		if (Ok)
+		{
+			// TODO: could also add reporting on just how much memory was purged
+			ZEN_CONSOLE("purged standby lists! (took {})", zen::NiceTimeSpanMs(Timer.GetElapsedTimeMs()));
+		}
+	}
+#endif
+
+	return 0;
+}