From 5a68a246e750359819d63bcff5ef97dd9c7788fc Mon Sep 17 00:00:00 2001 From: Denys Vlasenko Date: Fri, 31 May 2024 11:56:40 +0200 Subject: 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 --- coreutils/nproc.c | 9 +++++---- include/libbb.h | 2 ++ libbb/Kbuild.src | 1 + libbb/alloc_affinity.c | 29 +++++++++++++++++++++++++++++ util-linux/taskset.c | 24 +----------------------- 5 files changed, 38 insertions(+), 27 deletions(-) create mode 100644 libbb/alloc_affinity.c 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 #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 +#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 #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) { -- cgit v1.2.3-55-g6feb