From 76294beae1019905c4b3208b4a139631dd4834f6 Mon Sep 17 00:00:00 2001
From: Nekotekina <nekotekina@gmail.com>
Date: Tue, 28 Apr 2020 21:49:13 +0300
Subject: [PATCH] Implement thread_ctrl::get_thread_affinity_mask()

---
 Utilities/Thread.cpp | 62 +++++++++++++++++++++++++++++++++++++++++++-
 Utilities/Thread.h   |  3 +++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/Utilities/Thread.cpp b/Utilities/Thread.cpp
index 171895b94b..e56b396cc8 100644
--- a/Utilities/Thread.cpp
+++ b/Utilities/Thread.cpp
@@ -2431,6 +2431,66 @@ void thread_ctrl::set_thread_affinity_mask(u64 mask)
 		}
 	}
 
-	pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cs);
+	if (int err = pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cs))
+	{
+		sig_log.error("Failed to set thread affinity 0x%x: error %d.", mask, err);
+	}
+#endif
+}
+
+u64 thread_ctrl::get_thread_affinity_mask()
+{
+#ifdef _WIN32
+	DWORD_PTR res, _sys;
+	if (!GetProcessAffinityMask(GetCurrentProcess(), &res, &_sys))
+	{
+		sig_log.error("Failed to get process affinity mask.");
+		return 0;
+	}
+
+	if (DWORD_PTR result = SetThreadAffinityMask(GetCurrentThread(), res))
+	{
+		if (res != result)
+		{
+			SetThreadAffinityMask(GetCurrentThread(), result);
+		}
+
+		return result;
+	}
+
+	sig_log.error("Failed to get thread affinity mask.");
+	return 0;
+#elif defined(__linux__) || defined(__DragonFly__) || defined(__FreeBSD__)
+	cpu_set_t cs;
+	CPU_ZERO(&cs);
+
+	if (int err = pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &cs))
+	{
+		sig_log.error("Failed to get thread affinity mask: error %d.", err);
+		return 0;
+	}
+
+	u64 result;
+
+	for (u32 core = 0; core < 64u; core++)
+	{
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wold-style-cast"
+		if (CPU_ISSET(core, &cs))
+#pragma GCC diagnostic pop
+		{
+			result |= 1ull << core;
+		}
+	}
+
+	if (result == 0)
+	{
+		sig_log.error("Thread affinity mask is out of u64 range.");
+		return 0;
+	}
+
+	return result;
+#else
+	return -1;
 #endif
 }
diff --git a/Utilities/Thread.h b/Utilities/Thread.h
index 40469f19d4..8e69ca0ad1 100644
--- a/Utilities/Thread.h
+++ b/Utilities/Thread.h
@@ -248,6 +248,9 @@ public:
 
 	// Sets the preferred affinity mask for this thread
 	static void set_thread_affinity_mask(u64 mask);
+
+	// Miscellaneous
+	static u64 get_thread_affinity_mask();
 };
 
 // Derived from the callable object Context, possibly a lambda