From 88d7a7fd153bf1355cdf798ffdac7443d0169afc Mon Sep 17 00:00:00 2001
From: Lasse Collin <lasse.collin@tukaani.org>
Date: Mon, 24 Oct 2016 18:51:36 +0300
Subject: [PATCH] tuklib_cpucores: Add support for sched_getaffinity().

It's available in glibc (GNU/Linux, GNU/kFreeBSD). It's better
than sysconf(_SC_NPROCESSORS_ONLN) because sched_getaffinity()
gives the number of cores available to the process instead of
the total number of cores online.

As a side effect, this commit fixes a bug on GNU/kFreeBSD where
configure would detect the FreeBSD-specific cpuset_getaffinity()
but it wouldn't actually work because on GNU/kFreeBSD it requires
using -lfreebsd-glue when linking. Now the glibc-specific function
will be used instead.

Thanks to Sebastian Andrzej Siewior for the original patch
and testing.
---
 m4/tuklib_cpucores.m4        | 30 +++++++++++++++++++++++++++++-
 src/common/tuklib_cpucores.c |  9 +++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/m4/tuklib_cpucores.m4 b/m4/tuklib_cpucores.m4
index 468c2db6..a2b09a72 100644
--- a/m4/tuklib_cpucores.m4
+++ b/m4/tuklib_cpucores.m4
@@ -10,6 +10,8 @@
 #
 #   Supported methods:
 #     - GetSystemInfo(): Windows (including Cygwin)
+#     - sched_getaffinity(): glibc (GNU/Linux, GNU/kFreeBSD)
+#     - cpuset_getaffinity(): FreeBSD
 #     - sysctl(): BSDs, OS/2
 #     - sysconf(): GNU/Linux, Solaris, Tru64, IRIX, AIX, QNX, Cygwin (but
 #       GetSystemInfo() is used on Cygwin)
@@ -45,8 +47,29 @@ compile error
 #endif
 ]])], [tuklib_cv_cpucores_method=special], [
 
+# glibc-based systems (GNU/Linux and GNU/kFreeBSD) have sched_getaffinity().
+# The CPU_COUNT() macro was added in glibc 2.9 so we try to link the
+# test program instead of merely compiling it. glibc 2.9 is old enough that
+# if someone uses the code on older glibc, the fallback to sysconf() should
+# be good enough.
+AC_LINK_IFELSE([AC_LANG_SOURCE([[
+#include <sched.h>
+int
+main(void)
+{
+	cpu_set_t cpu_mask;
+	sched_getaffinity(0, sizeof(cpu_mask), &cpu_mask);
+	return CPU_COUNT(&cpu_mask);
+}
+]])], [tuklib_cv_cpucores_method=sched_getaffinity], [
+
 # FreeBSD has both cpuset and sysctl. Look for cpuset first because
 # it's a better approach.
+#
+# This test would match on GNU/kFreeBSD too but it would require
+# -lfreebsd-glue when linking and thus in the current form this would
+# fail on GNU/kFreeBSD. The above test for sched_getaffinity() matches
+# on GNU/kFreeBSD so the test below should never run on that OS.
 AC_COMPILE_IFELSE([AC_LANG_SOURCE([[
 #include <sys/param.h>
 #include <sys/cpuset.h>
@@ -120,9 +143,14 @@ main(void)
 ]])], [tuklib_cv_cpucores_method=pstat_getdynamic], [
 
 	tuklib_cv_cpucores_method=unknown
-])])])])])])
+])])])])])])])
 
 case $tuklib_cv_cpucores_method in
+	sched_getaffinity)
+		AC_DEFINE([TUKLIB_CPUCORES_SCHED_GETAFFINITY], [1],
+			[Define to 1 if the number of available CPU cores
+			can be detected with sched_getaffinity()])
+		;;
 	cpuset)
 		AC_DEFINE([TUKLIB_CPUCORES_CPUSET], [1],
 			[Define to 1 if the number of available CPU cores
diff --git a/src/common/tuklib_cpucores.c b/src/common/tuklib_cpucores.c
index e235fd1c..c16e188d 100644
--- a/src/common/tuklib_cpucores.c
+++ b/src/common/tuklib_cpucores.c
@@ -18,6 +18,10 @@
 #	endif
 #	include <windows.h>
 
+// glibc >= 2.9
+#elif defined(TUKLIB_CPUCORES_SCHED_GETAFFINITY)
+#	include <sched.h>
+
 // FreeBSD
 #elif defined(TUKLIB_CPUCORES_CPUSET)
 #	include <sys/param.h>
@@ -49,6 +53,11 @@ tuklib_cpucores(void)
 	GetSystemInfo(&sysinfo);
 	ret = sysinfo.dwNumberOfProcessors;
 
+#elif defined(TUKLIB_CPUCORES_SCHED_GETAFFINITY)
+	cpu_set_t cpu_mask;
+	if (sched_getaffinity(0, sizeof(cpu_mask), &cpu_mask) == 0)
+		ret = CPU_COUNT(&cpu_mask);
+
 #elif defined(TUKLIB_CPUCORES_CPUSET)
 	cpuset_t set;
 	if (cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, -1,