nproc: prepare for arbitrarily large CPU masks

function                                             old     new   delta
get_malloc_cpu_affinity                                -      76     +76
nproc_main                                           216     206     -10
process_pid_str                                      250     206     -44
------------------------------------------------------------------------------
(add/remove: 2/0 grow/shrink: 0/2 up/down: 76/-54)             Total: 22 bytes

Signed-off-by: Denys Vlasenko <vda.linux@googlemail.com>

5 files changed, 38 insertions, 27 deletions

diff --git a/coreutils/nproc.c b/coreutils/nproc.c
index f1d11fa5a..a0d818c59 100644
--- a/coreutils/nproc.c
+++ b/coreutils/nproc.c
@@ -23,13 +23,11 @@
 //usage:     "\n	--ignore=N	Exclude N CPUs"
 //usage:	)
 
-#include <sched.h>
 #include "libbb.h"
 
 int nproc_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
 int nproc_main(int argc UNUSED_PARAM, char **argv UNUSED_PARAM)
 {
-	unsigned long mask[1024];
 	int count = 0;
 #if ENABLE_LONG_OPTS
 	int ignore = 0;
@@ -52,9 +50,12 @@ int nproc_main(int argc UNUSED_PARAM, char **argv UNUSED_PARAM)
 		}
 	} else
 #endif
-	if (sched_getaffinity(0, sizeof(mask), (void*)mask) == 0) {
+	{
 		int i;
-		for (i = 0; i < ARRAY_SIZE(mask); i++) {
+		unsigned sz = 2 * 1024;
+		unsigned long *mask = get_malloc_cpu_affinity(0, &sz);
+		sz /= sizeof(long);
+		for (i = 0; i < sz; i++) {
 			unsigned long m = mask[i];
 			while (m) {
 				if (m & 1)
diff --git a/include/libbb.h b/include/libbb.h
index ef5d04713..67d29f843 100644
--- a/include/libbb.h
+++ b/include/libbb.h
@@ -2015,6 +2015,8 @@ int read_line_input(const char* prompt, char* command, int maxsize) FAST_FUNC;
 	read_line_input(prompt, command, maxsize)
 #endif
 
+unsigned long* FAST_FUNC get_malloc_cpu_affinity(int pid, unsigned *sz);
+
 #ifndef COMM_LEN
 # ifdef TASK_COMM_LEN
 enum { COMM_LEN = TASK_COMM_LEN };
diff --git a/libbb/Kbuild.src b/libbb/Kbuild.src
index c3b30003f..a0e2a6da7 100644
--- a/libbb/Kbuild.src
+++ b/libbb/Kbuild.src
@@ -10,6 +10,7 @@ lib-y:=
 
 INSERT
 
+lib-y += alloc_affinity.o
 lib-y += appletlib.o
 lib-y += ask_confirmation.o
 lib-y += bb_askpass.o
diff --git a/libbb/alloc_affinity.c b/libbb/alloc_affinity.c
new file mode 100644
index 000000000..b6d9f649a
--- /dev/null
+++ b/libbb/alloc_affinity.c
@@ -0,0 +1,29 @@
+/* vi: set sw=4 ts=4: */
+/*
+ * Utility routines.
+ *
+ * Copyright (C) 2024 Denys Vlasenko
+ *
+ * Licensed under GPLv2, see file LICENSE in this source tree.
+ */
+#include <sched.h>
+#include "libbb.h"
+
+unsigned long* FAST_FUNC get_malloc_cpu_affinity(int pid, unsigned *sz)
+{
+	unsigned long *mask = NULL;
+	unsigned sz_in_bytes = *sz;
+
+	for (;;) {
+		mask = xrealloc(mask, sz_in_bytes);
+		if (sched_getaffinity(pid, sz_in_bytes, (void*)mask) == 0)
+			break; /* got it */
+		sz_in_bytes *= 2;
+		if (errno == EINVAL && (int)sz_in_bytes > 0)
+			continue;
+		bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
+	}
+	//bb_error_msg("get mask[0]:%lx sz_in_bytes:%d", mask[0], sz_in_bytes);
+	*sz = sz_in_bytes;
+	return mask;
+}
diff --git a/util-linux/taskset.c b/util-linux/taskset.c
index 55c915e8d..a3aa06119 100644
--- a/util-linux/taskset.c
+++ b/util-linux/taskset.c
@@ -56,7 +56,6 @@
  * -a/--all-tasks (affect all threads)
  *	needs to get TIDs from /proc/PID/task/ and use _them_ as "pid" in sched_setaffinity(pid)
  */
-
 #include <sched.h>
 #include "libbb.h"
 
@@ -96,27 +95,6 @@ static unsigned long from_mask(ul *mask, unsigned sz_in_bytes UNUSED_PARAM)
 }
 #endif
 
-static unsigned long *get_aff(int pid, unsigned *sz)
-{
-	int r;
-	unsigned long *mask = NULL;
-	unsigned sz_in_bytes = *sz;
-
-	for (;;) {
-		mask = xrealloc(mask, sz_in_bytes);
-		r = sched_getaffinity(pid, sz_in_bytes, (void*)mask);
-		if (r == 0)
-			break;
-		sz_in_bytes *= 2;
-		if (errno == EINVAL && (int)sz_in_bytes > 0)
-			continue;
-		bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
-	}
-	//bb_error_msg("get mask[0]:%lx sz_in_bytes:%d", mask[0], sz_in_bytes);
-	*sz = sz_in_bytes;
-	return mask;
-}
-
 #if ENABLE_FEATURE_TASKSET_CPULIST
 /*
  * Parse the CPU list and set the mask accordingly.
@@ -222,7 +200,7 @@ static int process_pid_str(const char *pid_str, unsigned opts, char *aff)
 	mask_size_in_bytes = SZOF_UL;
 	current_new = "current";
  print_aff:
-	mask = get_aff(pid, &mask_size_in_bytes);
+	mask = get_malloc_cpu_affinity(pid, &mask_size_in_bytes);
 	if (opts & OPT_p) {
 #if ENABLE_FEATURE_TASKSET_CPULIST
 		if (opts & OPT_c) {