Cache process handles for FormatPackageMessage (#360)

1 files changed, 150 insertions, 0 deletions

diff --git a/src/zenutil/openprocesscache.cpp b/src/zenutil/openprocesscache.cpp
new file mode 100644
index 000000000..6c8bf3de9
--- /dev/null
+++ b/src/zenutil/openprocesscache.cpp
@@ -0,0 +1,150 @@
+// Copyright Epic Games, Inc. All Rights Reserved.
+
+#include "zenutil/openprocesscache.h"
+#include <zencore/fmtutils.h>
+#include <zencore/logging.h>
+
+#if ZEN_PLATFORM_WINDOWS
+#	include <zencore/windows.h>
+#else
+#	include <sys/mman.h>
+#endif
+
+//////////////////////////////////////////////////////////////////////////
+
+namespace zen {
+
+OpenProcessCache::OpenProcessCache()
+#if ZEN_PLATFORM_WINDOWS
+: m_GcThread(&OpenProcessCache::GcWorker, this)
+#endif	// ZEN_PLATFORM_WINDOWS
+{
+}
+
+OpenProcessCache::~OpenProcessCache()
+{
+#if ZEN_PLATFORM_WINDOWS
+	try
+	{
+		m_GcExitEvent.Set();
+		if (m_GcThread.joinable())
+		{
+			m_GcThread.join();
+		}
+		RwLock::ExclusiveLockScope Lock(m_SessionsLock);
+		for (const auto& It : m_Sessions)
+		{
+			if (It.second.ProcessHandle == nullptr)
+			{
+				continue;
+			}
+			CloseHandle(It.second.ProcessHandle);
+		}
+		m_Sessions.clear();
+	}
+	catch (std::exception& Ex)
+	{
+		ZEN_ERROR("OpenProcessCache destructor failed with reason: `{}`", Ex.what());
+	}
+#endif	// ZEN_PLATFORM_WINDOWS
+}
+
+void*
+OpenProcessCache::GetProcessHandle(Oid SessionId, int ProcessPid)
+{
+	if (ProcessPid == 0)
+	{
+		return nullptr;
+	}
+#if ZEN_PLATFORM_WINDOWS
+	RwLock::ExclusiveLockScope Lock(m_SessionsLock);
+
+	void* DeadHandle = nullptr;
+	if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
+	{
+		if (ProcessPid == It->second.ProcessPid)
+		{
+			void* ProcessHandle = It->second.ProcessHandle;
+			ZEN_ASSERT(ProcessHandle != nullptr);
+			DWORD ExitCode = 0;
+			GetExitCodeProcess(It->second.ProcessHandle, &ExitCode);
+			if (ExitCode == STILL_ACTIVE)
+			{
+				return It->second.ProcessHandle;
+			}
+		}
+
+		// Remove the existing stale handle
+		ZEN_INFO("Closing stale target process pid {} for session: {}", It->second.ProcessPid, SessionId, It->second.ProcessHandle);
+		DeadHandle = It->second.ProcessHandle;
+		m_Sessions.erase(It);
+	}
+
+	// Cache new handle
+	void* ProcessHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_DUP_HANDLE, FALSE, ProcessPid);
+	ZEN_INFO("Opening target process pid {} for session {}: {}", ProcessPid, SessionId, ProcessHandle == nullptr ? "failed" : "success");
+	m_Sessions.insert_or_assign(SessionId, Process{.ProcessPid = ProcessPid, .ProcessHandle = ProcessHandle});
+	Lock.ReleaseNow();
+
+	if (DeadHandle != nullptr)
+	{
+		CloseHandle(DeadHandle);
+	}
+
+	return ProcessHandle;
+#else	// ZEN_PLATFORM_WINDOWS
+	ZEN_UNUSED(SessionId);
+	return nullptr;
+#endif	// ZEN_PLATFORM_WINDOWS
+}
+
+#if ZEN_PLATFORM_WINDOWS
+void
+OpenProcessCache::GCHandles()
+{
+	std::vector<void*>		   DeadHandles;
+	RwLock::ExclusiveLockScope Lock(m_SessionsLock);
+	for (auto& It : m_Sessions)
+	{
+		if (It.second.ProcessHandle == nullptr)
+		{
+			continue;
+		}
+		ZEN_ASSERT(It.second.ProcessPid != 0);
+		DWORD ExitCode = 0;
+		GetExitCodeProcess(It.second.ProcessHandle, &ExitCode);
+		if (ExitCode == STILL_ACTIVE)
+		{
+			continue;
+		}
+		ZEN_INFO("GC stale target process pid {} for session {}: {}", It.second.ProcessPid, It.first, It.second.ProcessHandle);
+		DeadHandles.push_back(It.second.ProcessHandle);
+		It.second.ProcessPid	= 0;
+		It.second.ProcessHandle = nullptr;
+	}
+	Lock.ReleaseNow();
+
+	for (auto Handle : DeadHandles)
+	{
+		CloseHandle(Handle);
+	}
+}
+
+void
+OpenProcessCache::GcWorker()
+{
+	while (!m_GcExitEvent.Wait(500))
+	{
+		try
+		{
+			GCHandles();
+		}
+		catch (std::exception& Ex)
+		{
+			ZEN_ERROR("gc of open process cache failed with reason: `{}`", Ex.what());
+		}
+	}
+}
+#endif	// ZEN_PLATFORM_WINDOWS
+
+}  // namespace zen