taskset: add support for taking/printing CPU list (-c option)

function                                             old     new   delta
taskset_main                                         511     855    +344

Based on patch by Fryderyk Wrobel <frd1996@gmail.com>

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

1 files changed, 118 insertions, 11 deletions

diff --git a/util-linux/taskset.c b/util-linux/taskset.c
index ed8878ad4..cebc20fea 100644
--- a/util-linux/taskset.c
+++ b/util-linux/taskset.c
@@ -20,6 +20,14 @@
 //config:	Needed for machines with more than 32-64 CPUs:
 //config:	affinity parameter 0xHHHHHHHHHHHHHHHHHHHH can be arbitrarily long
 //config:	in this case. Otherwise, it is limited to sizeof(long).
+//config:
+//config:config FEATURE_TASKSET_CPULIST
+//config:	bool "CPU list support (-c option)"
+//config:	default y
+//config:	depends on FEATURE_TASKSET_FANCY
+//config:	help
+//config:	Add support for taking/printing affinity as CPU list when '-c'
+//config:	option is used. For example, it prints '0-3,7' instead of mask '8f'.
 
 //applet:IF_TASKSET(APPLET_NOEXEC(taskset, taskset, BB_DIR_USR_BIN, BB_SUID_DROP, taskset))
 
@@ -108,26 +116,109 @@ static unsigned long *get_aff(int pid, unsigned *sz)
 	return mask;
 }
 
+#if ENABLE_FEATURE_TASKSET_CPULIST
+/*
+ * Parse the CPU list and set the mask accordingly.
+ *
+ * The list element can be either a CPU index or a range of CPU indices.
+ * Example: "1,3,5-7".
+ *
+ * note1: pattern specifiers after a range (e.g. 0-255:2/64) are not supported
+ * note2: leading/trailing white-spaces are not allowed
+ */
+static void parse_cpulist(ul *mask, unsigned max, char *s)
+{
+	char *aff = s;
+	for (;;) {
+		unsigned bit, end;
+
+		bit = end = bb_strtou(s, &s, 10);
+		if (*s == '-') {
+			s++;
+			end = bb_strtou(s, &s, 10);
+		}
+		if ((*s != ',' && *s != '\0')
+		 || bit > end
+		 || end == UINT_MAX /* bb_strtou returns this on malformed / ERANGE numbers */
+		) {
+			bb_error_msg_and_die("bad affinity '%s'", aff);
+		}
+		while (bit <= end && bit < max) {
+			mask[bit / BITS_UL] |= (1UL << (bit & MASK_UL));
+			bit++;
+		}
+		if (*s == '\0')
+			break;
+		s++;
+	}
+}
+static void print_cpulist(const ul *mask, unsigned mask_size_in_bytes)
+{
+	const ul *mask_end;
+	const char *delim;
+	unsigned pos;
+	ul bit;
+
+	mask_end = mask + mask_size_in_bytes / sizeof(mask[0]);
+	delim = "";
+	pos = 0;
+	bit = 1;
+	for (;;) {
+		if (*mask & bit) {
+			unsigned onebit = pos + 1;
+			printf("%s%u", delim, pos);
+			do {
+				pos++;
+				bit <<= 1;
+				if (bit == 0) {
+					mask++;
+					if (mask >= mask_end)
+						break;
+					bit = 1;
+				}
+			} while (*mask & bit);
+			if (onebit != pos)
+				printf("-%u", pos - 1);
+			delim = ",";
+		}
+		pos++;
+		bit <<= 1;
+		if (bit == 0) {
+			mask++;
+			if (mask >= mask_end)
+				break;
+			bit = 1;
+		}
+	}
+	bb_putchar('\n');
+}
+#endif
+
 int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
 int taskset_main(int argc UNUSED_PARAM, char **argv)
 {
 	ul *mask;
 	unsigned mask_size_in_bytes;
 	pid_t pid = 0;
-	unsigned opt_p;
 	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:..." */
 
-	opt_p = getopt32(argv, "^+" "p" "\0" "-1" /* at least 1 arg */);
+	opts = getopt32(argv, "^+" "p"IF_FEATURE_TASKSET_CPULIST("c")
+			"\0" "-1" /* at least 1 arg */);
 	argv += optind;
 
 	aff = *argv++;
-	if (opt_p) {
+	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 */
@@ -144,8 +235,14 @@ int taskset_main(int argc UNUSED_PARAM, char **argv)
 	current_new = "current";
  print_aff:
 	mask = get_aff(pid, &mask_size_in_bytes);
-	if (opt_p) {
-		printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
+	if (opts & OPT_p) {
+#if ENABLE_FEATURE_TASKSET_CPULIST
+		if (opts & OPT_c) {
+			printf("pid %d's %s affinity list: ", pid, current_new);
+			print_cpulist(mask, mask_size_in_bytes);
+		} else
+#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) {
 			/* Either it was just "-p <pid>",
@@ -158,17 +255,27 @@ int taskset_main(int argc UNUSED_PARAM, char **argv)
 	}
 	memset(mask, 0, mask_size_in_bytes);
 
-	/* Affinity was specified, translate it into mask */
-	/* it is always in hex, skip "0x" if it exists */
-	if (aff[0] == '0' && (aff[1]|0x20) == 'x')
-		aff += 2;
-
 	if (!ENABLE_FEATURE_TASKSET_FANCY) {
+		/* Affinity was specified, translate it into mask */
+		/* it is always in hex, skip "0x" if it exists */
+		if (aff[0] == '0' && (aff[1]|0x20) == 'x')
+			aff += 2;
 		mask[0] = xstrtoul(aff, 16);
-	} else {
+	}
+#if ENABLE_FEATURE_TASKSET_CPULIST
+	else if (opts & OPT_c) {
+		parse_cpulist(mask, mask_size_in_bytes * 8, aff);
+	}
+#endif
+	else {
 		unsigned i;
 		char *last_char;
 
+		/* Affinity was specified, translate it into mask */
+		/* it is always in hex, skip "0x" if it exists */
+		if (aff[0] == '0' && (aff[1]|0x20) == 'x')
+			aff += 2;
+
 		i = 0; /* bit pos in mask[] */
 
 		/* aff is ASCII hex string, accept very long masks in this form.