bfstd: Take sched_getaffinity() into account in nproc()

3 files changed, 58 insertions, 5 deletions

diff --git a/build/has/sched-getaffinity.c b/build/has/sched-getaffinity.c
new file mode 100644
index 0000000..6f8fd98
--- /dev/null
+++ b/build/has/sched-getaffinity.c
@@ -0,0 +1,9 @@
+// Copyright © Tavian Barnes <tavianator@tavianator.com>
+// SPDX-License-Identifier: 0BSD
+
+#include <sched.h>
+
+int main(void) {
+	cpu_set_t set;
+	return sched_getaffinity(0, sizeof(set), &set);
+}
diff --git a/build/header.mk b/build/header.mk
index f8aee4b..19cb175 100644
--- a/build/header.mk
+++ b/build/header.mk
@@ -43,6 +43,7 @@ HEADERS := \
     gen/has/posix-spawn-addfchdir.h \
     gen/has/pthread-set-name-np.h \
     gen/has/pthread-setname-np.h \
+    gen/has/sched-getaffinity.h \
     gen/has/st-acmtim.h \
     gen/has/st-acmtimespec.h \
     gen/has/st-birthtim.h \
diff --git a/src/bfstd.c b/src/bfstd.c
index aee6930..5d60199 100644
--- a/src/bfstd.c
+++ b/src/bfstd.c
@@ -775,14 +775,57 @@ long xsysconf(int name) {
 	return ret;
 }
 
+#if BFS_HAS_SCHED_GETAFFINITY
+/** Get the CPU count in an affinity mask of the given size. */
+static long bfs_sched_getaffinity(size_t size) {
+	cpu_set_t set, *pset = &set;
+
+	if (size > sizeof(set)) {
+		pset = malloc(size);
+		if (!pset) {
+			return -1;
+		}
+	}
+
+	long ret = -1;
+	if (sched_getaffinity(0, size, pset) == 0) {
+		ret = CPU_COUNT_S(size, pset);
+	}
+
+	if (pset != &set) {
+		free(pset);
+	}
+	return ret;
+}
+#endif
+
 long nproc(void) {
-	long nproc = xsysconf(_SC_NPROCESSORS_ONLN);
+	long ret = 0;
 
-	if (nproc < 1) {
-		return 1;
-	} else {
-		return nproc;
+#if BFS_HAS_SCHED_GETAFFINITY
+	size_t size = sizeof(cpu_set_t);
+	do {
+		// sched_getaffinity(2) says:
+		//
+		//     When working on systems with large kernel CPU affinity masks, one must
+		//     dynamically allocate the mask argument (see CPU_ALLOC(3)).  Currently,
+		//     the only way to do this is by probing for the size of the required mask
+		//     using sched_getaffinity() calls with increasing mask sizes (until the
+		//     call does not fail with the error EINVAL).
+		ret = bfs_sched_getaffinity(size);
+		size *= 2;
+	} while (ret < 0 && errno == EINVAL);
+#endif
+
+	if (ret < 1) {
+		ret = xsysconf(_SC_NPROCESSORS_ONLN);
 	}
+
+	if (ret < 1) {
+		ret = 1;
+	}
+
+	return ret;
 }
 
 size_t asciilen(const char *str) {