From 09de26185e8d8ed307ba82e47ec7126da85fa564 Mon Sep 17 00:00:00 2001
From: Kryggird <plt.yaea@proton.me>
Date: Fri, 22 Dec 2023 16:11:45 +0000
Subject: [PATCH] Account for cpu affinity in g_get_num_processors

---
 glib/gthread.c      | 20 ++++++++++++++++++++
 glib/tests/thread.c | 42 ++++++++++++++++++++++++++++++++++++++++++
 meson.build         |  3 +++
 3 files changed, 65 insertions(+)

diff --git a/glib/gthread.c b/glib/gthread.c
index a64b1ce14..e7012fd11 100644
--- a/glib/gthread.c
+++ b/glib/gthread.c
@@ -47,7 +47,11 @@
 
 #ifdef G_OS_UNIX
 #include <unistd.h>
+
+#if defined(THREADS_POSIX) && defined(HAVE_PTHREAD_GETAFFINITY_NP)
+#include <pthread.h>
 #endif
+#endif /* G_OS_UNIX */
 
 #ifndef G_OS_WIN32
 #include <sys/time.h>
@@ -1072,6 +1076,22 @@ g_get_num_processors (void)
 
   if (count > 0)
     return count;
+#elif defined(_SC_NPROCESSORS_ONLN) && defined(THREADS_POSIX) && defined(HAVE_PTHREAD_GETAFFINITY_NP)
+  {
+    int idx;
+    int ncores = MIN (sysconf (_SC_NPROCESSORS_ONLN), CPU_SETSIZE);
+    cpu_set_t cpu_mask;
+    CPU_ZERO (&cpu_mask);
+
+    int af_count = 0;
+    int err = pthread_getaffinity_np (pthread_self (), sizeof (cpu_mask), &cpu_mask);
+    if (!err)
+      for (idx = 0; idx < ncores && idx < CPU_SETSIZE; ++idx)
+        af_count += CPU_ISSET (idx, &cpu_mask);
+
+    int count = (af_count > 0) ? af_count : ncores;
+    return count;
+  }
 #elif defined(_SC_NPROCESSORS_ONLN)
   {
     int count;
diff --git a/glib/tests/thread.c b/glib/tests/thread.c
index 25c62cc6d..6e8779612 100644
--- a/glib/tests/thread.c
+++ b/glib/tests/thread.c
@@ -212,6 +212,47 @@ test_thread6 (void)
   g_thread_join (thread);
 }
 
+#if defined(_SC_NPROCESSORS_ONLN) && defined(THREADS_POSIX) && defined(HAVE_PTHREAD_GETAFFINITY_NP)
+static gpointer
+thread7_func (gpointer data)
+{
+  int idx = 0, err;
+  int ncores = sysconf (_SC_NPROCESSORS_ONLN);
+
+  cpu_set_t old_mask, new_mask;
+
+  err = pthread_getaffinity_np (pthread_self (), sizeof (old_mask), &old_mask);
+  CPU_ZERO (&new_mask);
+  g_assert_cmpint (err, ==, 0);
+
+  for (idx = 0; idx < ncores; ++idx)
+    if (CPU_ISSET (idx, &old_mask))
+      {
+        CPU_SET (idx, &new_mask);
+        break;
+      }
+
+  err = pthread_setaffinity_np (pthread_self (), sizeof (new_mask), &new_mask);
+  g_assert_cmpint (err, ==, 0);
+
+  int af_count = g_get_num_processors ();
+  return GINT_TO_POINTER (af_count);
+}
+#endif
+
+static void
+test_thread7 (void)
+{
+#if defined(_SC_NPROCESSORS_ONLN) && defined(THREADS_POSIX) && defined(HAVE_PTHREAD_GETAFFINITY_NP)
+  GThread *thread = g_thread_new ("mask", thread7_func, NULL);
+  gpointer result = g_thread_join (thread);
+
+  g_assert_cmpint (GPOINTER_TO_INT (result), ==, 1);
+#else
+  g_test_skip ("Skipping because pthread_getaffinity_np() is not available");
+#endif
+}
+
 int
 main (int argc, char *argv[])
 {
@@ -223,6 +264,7 @@ main (int argc, char *argv[])
   g_test_add_func ("/thread/thread4", test_thread4);
   g_test_add_func ("/thread/thread5", test_thread5);
   g_test_add_func ("/thread/thread6", test_thread6);
+  g_test_add_func ("/thread/thread7", test_thread7);
 
   return g_test_run ();
 }
diff --git a/meson.build b/meson.build
index 451590bec..f4dc2e5ee 100644
--- a/meson.build
+++ b/meson.build
@@ -2048,6 +2048,9 @@ else
   if cc.has_header_symbol('pthread.h', 'pthread_getname_np', prefix : pthread_prefix)
     glib_conf.set('HAVE_PTHREAD_GETNAME_NP', 1)
   endif
+  if cc.has_header_symbol('pthread.h', 'pthread_getaffinity_np', prefix : pthread_prefix)
+    glib_conf.set('HAVE_PTHREAD_GETAFFINITY_NP', 1)
+  endif
 
   # Assume that pthread_setname_np is available in some form; same as configure
   if cc.links(pthread_prefix + '''