added benchmark utility command `bench` (#298)

currently this implements a way (`zen bench --purge`) to purge the standby lists on Windows. This basically flushes the file system cache and is useful to put your system in a consistent state before running benchmarks

5 files changed, 260 insertions, 1 deletions

diff --git a/src/zen/cmds/bench.cpp b/src/zen/cmds/bench.cpp
new file mode 100644
index 000000000..a2986ce16
--- /dev/null
+++ b/src/zen/cmds/bench.cpp
@@ -0,0 +1,217 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "bench.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;
+}
diff --git a/src/zen/cmds/bench.h b/src/zen/cmds/bench.h
new file mode 100644
index 000000000..8a8bd4a7c
--- /dev/null
+++ b/src/zen/cmds/bench.h
@@ -0,0 +1,20 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#pragma once
+
+#include "../zen.h"
+
+class BenchCommand : public ZenCmdBase
+{
+public:
+	BenchCommand();
+	~BenchCommand();
+
+	virtual int				  Run(const ZenCliOptions& GlobalOptions, int argc, char** argv) override;
+	virtual cxxopts::Options& Options() override { return m_Options; }
+
+private:
+	cxxopts::Options m_Options{"bench", "Benchmarking utility command"};
+	bool			 m_PurgeStandbyLists = false;
+	bool			 m_SingleProcess	 = false;
+};
diff --git a/src/zen/zen.cpp b/src/zen/zen.cpp
index 3cb5bec47..0df0ab834 100644
--- a/src/zen/zen.cpp
+++ b/src/zen/zen.cpp
@@ -5,7 +5,7 @@
 
 #include "zen.h"
 
-#include "chunk/chunk.h"
+#include "cmds/bench.h"
 #include "cmds/cache.h"
 #include "cmds/copy.h"
 #include "cmds/dedup.h"
@@ -179,6 +179,7 @@ main(int argc, char** argv)
 
 	auto _ = zen::MakeGuard([] { spdlog::shutdown(); });
 
+	BenchCommand			 BenchCmd;
 	CacheInfoCommand		 CacheInfoCmd;
 	CopyCommand				 CopyCmd;
 	CreateOplogCommand		 CreateOplogCmd;
@@ -219,6 +220,7 @@ main(int argc, char** argv)
 		const char* CmdSummary;
 	} Commands[] = {
 		// clang-format off
+		{"bench",				&BenchCmd,				"Utility command for benchmarking"},
 //		{"chunk",				&ChunkCmd,				"Perform chunking"},
 		{"cache-info",			&CacheInfoCmd,			"Info on cache, namespace or bucket"},
 		{"copy",				&CopyCmd,				"Copy file(s)"},
diff --git a/src/zencore/include/zencore/thread.h b/src/zencore/include/zencore/thread.h
index 09c25996f..cc1a692a0 100644
--- a/src/zencore/include/zencore/thread.h
+++ b/src/zencore/include/zencore/thread.h
@@ -216,6 +216,7 @@ public:
 	ZENCORE_API [[nodiscard]] bool IsRunning() const;
 	ZENCORE_API [[nodiscard]] bool IsValid() const;
 	ZENCORE_API bool			   Wait(int TimeoutMs = -1);
+	ZENCORE_API int				   WaitExitCode();
 	ZENCORE_API void			   Terminate(int ExitCode);
 	ZENCORE_API void			   Reset();
 	[[nodiscard]] inline int	   Pid() const { return m_Pid; }
diff --git a/src/zencore/thread.cpp b/src/zencore/thread.cpp
index fe81cc786..52b1e5d4e 100644
--- a/src/zencore/thread.cpp
+++ b/src/zencore/thread.cpp
@@ -679,6 +679,25 @@ ProcessHandle::Wait(int TimeoutMs)
 	ThrowLastError("Process::Wait failed"sv);
 }
 
+int
+ProcessHandle::WaitExitCode()
+{
+	Wait(-1);
+
+#if ZEN_PLATFORM_WINDOWS
+	DWORD ExitCode = 0;
+	GetExitCodeProcess(m_ProcessHandle, &ExitCode);
+
+	ZEN_ASSERT(ExitCode != STILL_ACTIVE);
+
+	return ExitCode;
+#else
+	ZEN_NOT_IMPLEMENTED();
+
+	return 0;
+#endif
+}
+
 //////////////////////////////////////////////////////////////////////////
 
 #if !ZEN_PLATFORM_WINDOWS || ZEN_WITH_TESTS