From a36774be636c3429c6e73de33bf65a1bdbdcfb4b Mon Sep 17 00:00:00 2001 From: Tavian Barnes Date: Thu, 27 Feb 2025 13:39:04 -0500 Subject: bfstd: Take sched_getaffinity() into account in nproc() --- src/bfstd.c | 53 ++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 48 insertions(+), 5 deletions(-) (limited to 'src') 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) { -- cgit v1.2.3