diff options
| author | Dan Engelbrecht <[email protected]> | 2023-09-22 08:22:06 -0400 |
|---|---|---|
| committer | GitHub <[email protected]> | 2023-09-22 14:22:06 +0200 |
| commit | c7d4dc6a4d13881028d566f5ce501335e47e48bf (patch) | |
| tree | 493110da583a8e5d97fe05e14f23469ee6244d2b /src/zen/cmds/bench_cmd.cpp | |
| parent | add trace command to enable/disable tracing at runtime (#416) (diff) | |
| download | archived-zen-c7d4dc6a4d13881028d566f5ce501335e47e48bf.tar.xz archived-zen-c7d4dc6a4d13881028d566f5ce501335e47e48bf.zip | |
Collect all zen admin-related commands into admin.h/.cpp (#418)
* move commands in scrub.h/cpp to admin_cmd.h/cpp
* move job command into admin_cmd.h/.cpp
* admin -> admin_cmd
* bench -> bench_cmd
* cache -> cache_cmd
* copy -> copy_cmd
* dedup -> dedup_cmd
* hash -> hash_cmd
* print -> print_cmd
* projectstore -> projectstore_cmd
* rpcreplay -> rpcreplay_cmd
* serve -> serve_cmd
* status -> status_cmd
* top -> top_cmd
* trace -> trace_cmd
* up -> up_cmd
* version -> version_cmd
Diffstat (limited to 'src/zen/cmds/bench_cmd.cpp')
| -rw-r--r-- | src/zen/cmds/bench_cmd.cpp | 217 |
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; +} |