1 files changed, 70 insertions, 36 deletions

diff --git a/src/zenutil/openprocesscache.cpp b/src/zenutil/openprocesscache.cpp
index 6c8bf3de9..b8f01dd54 100644
--- a/src/zenutil/openprocesscache.cpp
+++ b/src/zenutil/openprocesscache.cpp
@@ -57,41 +57,65 @@ OpenProcessCache::GetProcessHandle(Oid SessionId, int ProcessPid)
 		return nullptr;
 	}
 #if ZEN_PLATFORM_WINDOWS
-	RwLock::ExclusiveLockScope Lock(m_SessionsLock);
+	ZEN_ASSERT(ProcessPid != 0);
+	{
+		RwLock::SharedLockScope Lock(m_SessionsLock);
 
-	void* DeadHandle = nullptr;
-	if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
+		if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
+		{
+			if (ProcessPid == It->second.ProcessPid)
+			{
+				void* ProcessHandle = It->second.ProcessHandle;
+				if (ProcessHandle == nullptr)
+				{
+					// Session is stale, just ignore it
+					return nullptr;
+				}
+				DWORD ExitCode = 0;
+				GetExitCodeProcess(ProcessHandle, &ExitCode);
+				if (ExitCode == STILL_ACTIVE)
+				{
+					return ProcessHandle;
+				}
+			}
+		}
+	}
+	void* ProcessHandle = nullptr;
+	void* DeadHandle	= nullptr;
 	{
-		if (ProcessPid == It->second.ProcessPid)
+		// Add/replace the session
+
+		RwLock::ExclusiveLockScope Lock(m_SessionsLock);
+		if (auto It = m_Sessions.find(SessionId); It != m_Sessions.end())
 		{
-			void* ProcessHandle = It->second.ProcessHandle;
-			ZEN_ASSERT(ProcessHandle != nullptr);
-			DWORD ExitCode = 0;
-			GetExitCodeProcess(It->second.ProcessHandle, &ExitCode);
-			if (ExitCode == STILL_ACTIVE)
+			// Check to see if someone beat us to it...
+			if ((It->second.ProcessHandle != nullptr) && (ProcessPid == It->second.ProcessPid))
 			{
-				return It->second.ProcessHandle;
+				DWORD ExitCode = 0;
+				GetExitCodeProcess(It->second.ProcessHandle, &ExitCode);
+				if (ExitCode == STILL_ACTIVE)
+				{
+					return It->second.ProcessHandle;
+				}
 			}
+
+			ZEN_INFO("Purging stale target process pid {} for session: {}", It->second.ProcessPid, SessionId, It->second.ProcessHandle);
+			DeadHandle = It->second.ProcessHandle;
+			m_Sessions.erase(It);
 		}
 
-		// 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
+		ProcessHandle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_DUP_HANDLE, FALSE, ProcessPid);
+		ZEN_INFO("Opened target process pid {} for session {}: {}", ProcessPid, SessionId, ProcessHandle == nullptr ? "failed" : "success");
+		m_Sessions.insert_or_assign(SessionId, Process{.ProcessPid = ProcessPid, .ProcessHandle = ProcessHandle});
 	}
 
-	// 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;
@@ -102,27 +126,37 @@ OpenProcessCache::GetProcessHandle(Oid SessionId, int ProcessPid)
 void
 OpenProcessCache::GCHandles()
 {
-	std::vector<void*>		   DeadHandles;
-	RwLock::ExclusiveLockScope Lock(m_SessionsLock);
-	for (auto& It : m_Sessions)
+	std::vector<void*> DeadHandles;
 	{
-		if (It.second.ProcessHandle == nullptr)
+		std::vector<Oid>		   DeadSessions;
+		RwLock::ExclusiveLockScope Lock(m_SessionsLock);
+		for (auto& It : m_Sessions)
 		{
-			continue;
+			if (It.second.ProcessHandle == nullptr)
+			{
+				// Stale session, remove it on second pass of GC
+				DeadSessions.push_back(It.first);
+				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);
+
+			// We just mark the session as "dead" for one GC pass to avoid re-opening a stale process
+			It.second.ProcessHandle = nullptr;
 		}
-		ZEN_ASSERT(It.second.ProcessPid != 0);
-		DWORD ExitCode = 0;
-		GetExitCodeProcess(It.second.ProcessHandle, &ExitCode);
-		if (ExitCode == STILL_ACTIVE)
+		for (auto SessionId : DeadSessions)
 		{
-			continue;
+			ZEN_INFO("Purging stale target process cache for session: {}", SessionId);
+			m_Sessions.erase(SessionId);
 		}
-		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)
 	{