From 32828e4e52f7297fdce3ff1515b77b905c7e93e0 Mon Sep 17 00:00:00 2001
From: Kawe Mazidjatari <48657826+Mauler125@users.noreply.github.com>
Date: Mon, 1 Apr 2024 00:46:21 +0200
Subject: [PATCH] Tier0: add new thread tools

- Add ThreadInMainOrServerFrameThread()
- Add ThreadCouldDoServerWork()
- Add ThreadJoinServerJob()
- Properly implement ThreadInServerFrameThread()
- Export all these functions for external DLL usage (plugins, etc)
---
 src/public/tier0/threadtools.h | 34 ++++++++++++++++++-------------
 src/tier0/threadtools.cpp      | 37 ++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/src/public/tier0/threadtools.h b/src/public/tier0/threadtools.h
index 67fa97e8..e47e2379 100644
--- a/src/public/tier0/threadtools.h
+++ b/src/public/tier0/threadtools.h
@@ -1,6 +1,9 @@
 #ifndef THREADTOOLS_H
 #define THREADTOOLS_H
 #include "dbg.h"
+#ifndef BUILDING_MATHLIB
+#include "jobthread.h"
+#endif // BUILDING_MATHLIB
 
 inline void ThreadSleep(unsigned nMilliseconds)
 {
@@ -116,16 +119,14 @@ FORCEINLINE ThreadId_t ThreadGetCurrentId()
 
 extern ThreadId_t* g_ThreadMainThreadID;
 extern ThreadId_t* g_ThreadServerFrameThreadID;
+inline JobID_t* g_CurrentServerFrameJobID;
+inline JobID_t* g_AllocatedServerFrameJobID;
 
-FORCEINLINE bool ThreadInMainThread()
-{
-	return (ThreadGetCurrentId() == (*g_ThreadMainThreadID));
-}
-
-FORCEINLINE bool ThreadInServerFrameThread()
-{
-	return (ThreadGetCurrentId() == (*g_ThreadServerFrameThreadID));
-}
+PLATFORM_INTERFACE bool ThreadInMainThread();
+PLATFORM_INTERFACE bool ThreadInServerFrameThread();
+PLATFORM_INTERFACE bool ThreadInMainOrServerFrameThread();
+PLATFORM_INTERFACE bool ThreadCouldDoServerWork();
+PLATFORM_INTERFACE void ThreadJoinServerJob();
 
 #endif // !BUILDING_MATHLIB
 
@@ -263,7 +264,7 @@ typedef CInterlockedIntT<unsigned> CInterlockedUInt;
 
 #ifndef BUILDING_MATHLIB
 //=============================================================================
-inline ThreadId_t(*v_DeclareCurrentThreadIsMainThread)(void);
+
 
 #endif // !BUILDING_MATHLIB
 
@@ -562,15 +563,20 @@ class VThreadTools : public IDetour
 {
 	virtual void GetAdr(void) const
 	{
-		LogFunAdr("DeclareCurrentThreadIsMainThread", v_DeclareCurrentThreadIsMainThread);
 		LogVarAdr("g_ThreadMainThreadID", g_ThreadMainThreadID);
 		LogVarAdr("g_ThreadServerFrameThreadID", g_ThreadServerFrameThreadID);
+		LogVarAdr("g_CurrentServerFrameJobID", g_CurrentServerFrameJobID);
+		LogVarAdr("g_AllocatedServerFrameJobID", g_AllocatedServerFrameJobID);
 	}
-	virtual void GetFun(void) const
+	virtual void GetFun(void) const { }
+	virtual void GetVar(void) const
 	{
-		g_GameDll.FindPatternSIMD("48 83 EC 28 FF 15 ?? ?? ?? ?? 89 05 ?? ?? ?? ?? 48 83 C4 28").GetPtr(v_DeclareCurrentThreadIsMainThread);
+		g_GameDll.FindPatternSIMD("66 89 54 24 ?? 53 55 56 57 41 54 48 81 EC ?? ?? ?? ??")
+			.FindPatternSelf("39 05").ResolveRelativeAddressSelf(2, 6).GetPtr(g_CurrentServerFrameJobID);
+
+		g_GameDll.FindPatternSIMD("48 83 EC 28 FF 15 ?? ?? ?? ?? 8B 0D ?? ?? ?? ??")
+			.FindPatternSelf("8B 0D").ResolveRelativeAddressSelf(2, 6).GetPtr(g_AllocatedServerFrameJobID);
 	}
-	virtual void GetVar(void) const { }
 	virtual void GetCon(void) const { }
 	virtual void Detour(const bool /*bAttach*/) const { }
 };
diff --git a/src/tier0/threadtools.cpp b/src/tier0/threadtools.cpp
index 952cd4cc..ecc5333e 100644
--- a/src/tier0/threadtools.cpp
+++ b/src/tier0/threadtools.cpp
@@ -7,6 +7,7 @@
 //===========================================================================//
 
 #include "tier0/threadtools.h"
+#include "tier0/jobthread.h"
 
 #define INIT_SEM_COUNT 0
 #define MAX_SEM_COUNT 1
@@ -198,6 +199,42 @@ void CThreadSpinRWLock::SpinLockForRead()
 	}
 }
 
+bool ThreadInMainThread()
+{
+	return (ThreadGetCurrentId() == (*g_ThreadMainThreadID));
+}
+
+bool ThreadInServerFrameThread()
+{
+	return (ThreadGetCurrentId() == (*g_ThreadServerFrameThreadID) 
+		&& JT_GetCurrentJob() == (*g_CurrentServerFrameJobID));
+}
+
+bool ThreadInMainOrServerFrameThread()
+{
+	return (ThreadInMainThread() || ThreadInServerFrameThread());
+}
+
+bool ThreadCouldDoServerWork()
+{
+	if (*g_ThreadServerFrameThreadID == -1)
+		return ThreadInMainThread();
+
+	return ThreadInServerFrameThread();
+}
+
+void ThreadJoinServerJob()
+{
+	if (ThreadCouldDoServerWork())
+		return; // No job to join
+
+	if (*g_AllocatedServerFrameJobID)
+	{
+		JT_WaitForJobAndOnlyHelpWithJobTypes(*g_AllocatedServerFrameJobID, NULL, 0xFFFFFFFFFFFFFFFF);
+		*g_AllocatedServerFrameJobID = 0;
+	}
+}
+
 // NOTE: originally the game exported 'ThreadInMainThread()' and ThreadInServerFrameThread(),
 // but since the game is built static, and all instances of said functions are inline, we had
 // to export the variable symbols instead and get them here to reimplement said functions.