1 files changed, 69 insertions, 39 deletions

diff --git a/util-linux/taskset.c b/util-linux/taskset.c
index b542f8c83..afe2f04d2 100644
--- a/util-linux/taskset.c
+++ b/util-linux/taskset.c
@@ -34,10 +34,12 @@
 //kbuild:lib-$(CONFIG_TASKSET) += taskset.o
 
 //usage:#define taskset_trivial_usage
-//usage:       "[-p] [HEXMASK] PID | PROG ARGS"
+//usage:       "[-ap] [HEXMASK"IF_FEATURE_TASKSET_CPULIST(" | -c LIST")"] { PID | PROG ARGS }"
 //usage:#define taskset_full_usage "\n\n"
 //usage:       "Set or get CPU affinity\n"
-//usage:     "\n        -p      Operate on an existing PID"
+//usage:     "\n        -p      Operate on PID"
+//usage:     "\n        -a      Operate on all threads"
+//usage:     "\n        -c      Affinity is a list, not mask"
 //usage:
 //usage:#define taskset_example_usage
 //usage:       "$ taskset 0x7 ./dgemm_test&\n"
@@ -205,42 +207,18 @@ static void print_cpulist(const ul *mask, unsigned mask_size_in_bytes)
 }
 #endif
 
-int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
-int taskset_main(int argc UNUSED_PARAM, char **argv)
+enum {
+        OPT_p = 1 << 0,
+        OPT_a = 1 << 1,
+        OPT_c = (1 << 2) * ENABLE_FEATURE_TASKSET_CPULIST,
+};
+
+static int process_pid_str(const char *pid_str, unsigned opts, char *aff)
 {
         ul *mask;
         unsigned mask_size_in_bytes;
-        pid_t pid = 0;
         const char *current_new;
-        char *aff;
-        unsigned opts;
-        enum {
-                OPT_p = 1 << 0,
-                OPT_c = (1 << 1) * ENABLE_FEATURE_TASKSET_CPULIST,
-        };
-
-        /* NB: we mimic util-linux's taskset: -p does not take
-         * an argument, i.e., "-pN" is NOT valid, only "-p N"!
-         * Indeed, util-linux-2.13-pre7 uses:
-         * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */
-
-        opts = getopt32(argv, "^+" "p"IF_FEATURE_TASKSET_CPULIST("c")
-                        "\0" "-1" /* at least 1 arg */);
-        argv += optind;
-
-        aff = *argv++;
-        if (opts & OPT_p) {
-                char *pid_str = aff;
-                if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */
-                        pid_str = *argv; /* NB: *argv != NULL in this case */
-                }
-                /* else it was just "-p <pid>", and *argv == NULL */
-                pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1);
-        } else {
-                /* <aff> <cmd...> */
-                if (!*argv)
-                        bb_show_usage();
-        }
+        pid_t pid = xatoi_positive(pid_str);
 
         mask_size_in_bytes = SZOF_UL;
         current_new = "current";
@@ -255,13 +233,12 @@ int taskset_main(int argc UNUSED_PARAM, char **argv)
 #endif
                         printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
                                 pid, current_new, from_mask(mask, mask_size_in_bytes));
-                if (*argv == NULL) {
+                if (!aff) {
                         /* Either it was just "-p <pid>",
                          * or it was "-p <aff> <pid>" and we came here
                          * for the second time (see goto below) */
-                        return EXIT_SUCCESS;
+                        return 0;
                 }
-                *argv = NULL;
                 current_new = "new";
         }
         memset(mask, 0, mask_size_in_bytes);
@@ -331,8 +308,61 @@ int taskset_main(int argc UNUSED_PARAM, char **argv)
                 bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);
         //bb_error_msg("set mask[0]:%lx", mask[0]);
 
-        if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */
+        if ((opts & OPT_p) && aff) { /* "-p <aff> <pid> [...ignored...]" */
+                aff = NULL;
                 goto print_aff; /* print new affinity and exit */
+        }
+        return 0;
+}
+
+static int FAST_FUNC iter(const char *dn UNUSED_PARAM, struct dirent *ent, void *aff)
+{
+        if (isdigit(ent->d_name[0]))
+                return process_pid_str(ent->d_name, option_mask32, aff);
+        return 0;
+}
+
+int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
+int taskset_main(int argc UNUSED_PARAM, char **argv)
+{
+        const char *pid_str;
+        char *aff;
+        unsigned opts;
 
-        BB_EXECVP_or_die(argv);
+        /* NB: we mimic util-linux's taskset: -p does not take
+         * an argument, i.e., "-pN" is NOT valid, only "-p N"!
+         * Indeed, util-linux-2.13-pre7 uses:
+         * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */
+
+        opts = getopt32(argv, "^+" "pa"IF_FEATURE_TASKSET_CPULIST("c")
+                        "\0" "-1" /* at least 1 arg */);
+        argv += optind;
+
+        aff = *argv++;
+        if (!(opts & OPT_p)) {
+                /* <aff> <cmd...> */
+                if (!*argv)
+                        bb_show_usage();
+                process_pid_str("0", opts, aff);
+                BB_EXECVP_or_die(argv);
+        }
+
+        pid_str = aff;
+        if (*argv) /* "-p <aff> <pid> ...rest.is.ignored..." */
+                pid_str = *argv;
+        else
+                aff = NULL;
+
+        if (opts & OPT_a) {
+                char *dn;
+                int r;
+
+                dn = xasprintf("/proc/%s/task", pid_str);
+                r = iterate_on_dir(dn, iter, aff);
+                IF_FEATURE_CLEAN_UP(free(dn);)
+                if (r == 0)
+                        return r; /* EXIT_SUCCESS */
+                /* else: no /proc/PID/task, act as if no -a was given */
+        }
+        return process_pid_str(pid_str, opts, aff);
 }