util-linux/taskset.c

   1 /* vi: set sw=4 ts=4: */
   2 /*
   3  * taskset - retrieve or set a processes' CPU affinity
   4  * Copyright (c) 2006 Bernhard Reutner-Fischer
   5  *
   6  * Licensed under GPLv2 or later, see file LICENSE in this source tree.
   7  */
   8 //config:config TASKSET
   9 //config:       bool "taskset (4.2 kb)"
  10 //config:       default y
  11 //config:       help
  12 //config:       Retrieve or set a processes's CPU affinity.
  13 //config:       This requires sched_{g,s}etaffinity support in your libc.
  14 //config:
  15 //config:config FEATURE_TASKSET_FANCY
  16 //config:       bool "Fancy output"
  17 //config:       default y
  18 //config:       depends on TASKSET
  19 //config:       help
  20 //config:       Needed for machines with more than 32-64 CPUs:
  21 //config:       affinity parameter 0xHHHHHHHHHHHHHHHHHHHH can be arbitrarily long
  22 //config:       in this case. Otherwise, it is limited to sizeof(long).
  23 //config:
  24 //config:config FEATURE_TASKSET_CPULIST
  25 //config:       bool "CPU list support (-c option)"
  26 //config:       default y
  27 //config:       depends on FEATURE_TASKSET_FANCY
  28 //config:       help
  29 //config:       Add support for taking/printing affinity as CPU list when '-c'
  30 //config:       option is used. For example, it prints '0-3,7' instead of mask '8f'.
  31
  32 //applet:IF_TASKSET(APPLET_NOEXEC(taskset, taskset, BB_DIR_USR_BIN, BB_SUID_DROP, taskset))
  33
  34 //kbuild:lib-$(CONFIG_TASKSET) += taskset.o
  35
  36 //usage:#define taskset_trivial_usage
  37 //usage:       "[-p] [HEXMASK] PID | PROG ARGS"
  38 //usage:#define taskset_full_usage "\n\n"
  39 //usage:       "Set or get CPU affinity\n"
  40 //usage:     "\n        -p      Operate on an existing PID"
  41 //usage:
  42 //usage:#define taskset_example_usage
  43 //usage:       "$ taskset 0x7 ./dgemm_test&\n"
  44 //usage:       "$ taskset -p 0x1 $!\n"
  45 //usage:       "pid 4790's current affinity mask: 7\n"
  46 //usage:       "pid 4790's new affinity mask: 1\n"
  47 //usage:       "$ taskset 0x7 /bin/sh -c './taskset -p 0x1 $$'\n"
  48 //usage:       "pid 6671's current affinity mask: 1\n"
  49 //usage:       "pid 6671's new affinity mask: 1\n"
  50 //usage:       "$ taskset -p 1\n"
  51 //usage:       "pid 1's current affinity mask: 3\n"
  52 /*
  53  * Not yet implemented:
  54  * -a/--all-tasks (affect all threads)
  55  *      needs to get TIDs from /proc/PID/task/ and use _them_ as "pid" in sched_setaffinity(pid)
  56  * -c/--cpu-list  (specify CPUs via "1,3,5-7")
  57  */
  58
  59 #include <sched.h>
  60 #include "libbb.h"
  61
  62 typedef unsigned long ul;
  63 #define SZOF_UL (unsigned)(sizeof(ul))
  64 #define BITS_UL (unsigned)(sizeof(ul)*8)
  65 #define MASK_UL (unsigned)(sizeof(ul)*8 - 1)
  66
  67 #if ENABLE_FEATURE_TASKSET_FANCY
  68 #define TASKSET_PRINTF_MASK "%s"
  69 /* craft a string from the mask */
  70 static char *from_mask(const ul *mask, unsigned sz_in_bytes)
  71 {
  72         char *str = xzalloc((sz_in_bytes+1) * 2); /* we will leak it */
  73         char *p = str;
  74         for (;;) {
  75                 ul v = *mask++;
  76                 if (SZOF_UL == 4)
  77                         p += sprintf(p, "%08lx", v);
  78                 if (SZOF_UL == 8)
  79                         p += sprintf(p, "%016lx", v);
  80                 if (SZOF_UL == 16)
  81                         p += sprintf(p, "%032lx", v); /* :) */
  82                 sz_in_bytes -= SZOF_UL;
  83                 if ((int)sz_in_bytes <= 0)
  84                         break;
  85         }
  86         while (str[0] == '0' && str[1])
  87                 str++;
  88         return str;
  89 }
  90 #else
  91 #define TASKSET_PRINTF_MASK "%lx"
  92 static unsigned long long from_mask(ul *mask, unsigned sz_in_bytes UNUSED_PARAM)
  93 {
  94         return *mask;
  95 }
  96 #endif
  97
  98 static unsigned long *get_aff(int pid, unsigned *sz)
  99 {
 100         int r;
 101         unsigned long *mask = NULL;
 102         unsigned sz_in_bytes = *sz;
 103
 104         for (;;) {
 105                 mask = xrealloc(mask, sz_in_bytes);
 106                 r = sched_getaffinity(pid, sz_in_bytes, (void*)mask);
 107                 if (r == 0)
 108                         break;
 109                 sz_in_bytes *= 2;
 110                 if (errno == EINVAL && (int)sz_in_bytes > 0)
 111                         continue;
 112                 bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
 113         }
 114         //bb_error_msg("get mask[0]:%lx sz_in_bytes:%d", mask[0], sz_in_bytes);
 115         *sz = sz_in_bytes;
 116         return mask;
 117 }
 118
 119 #if ENABLE_FEATURE_TASKSET_CPULIST
 120 /*
 121  * Parse the CPU list and set the mask accordingly.
 122  *
 123  * The list element can be either a CPU index or a range of CPU indices.
 124  * Example: "1,3,5-7". Stride can be specified: "0-7:2" is "0,2,4,6".
 125  * Note: leading and trailing whitespace is not allowed.
 126  *  util-linux 2.31 allows leading and sometimes trailing whitespace:
 127  *  ok:     taskset -c ' 1,  2'
 128  *  ok:     taskset -c ' 1 , 2'
 129  *  ok:     taskset -c ' 1-7: 2 ,8'
 130  *  not ok: taskset -c ' 1 '
 131  *  not ok: taskset -c ' 1-7: 2 '
 132  */
 133 static void parse_cpulist(ul *mask, unsigned max, char *s)
 134 {
 135         char *aff = s;
 136         for (;;) {
 137                 unsigned bit, end;
 138                 unsigned stride = 1;
 139
 140                 bit = end = bb_strtou(s, &s, 10);
 141                 if (*s == '-') {
 142                         s++;
 143                         end = bb_strtou(s, &s, 10);
 144                         if (*s == ':') {
 145                                 s++;
 146                                 stride = bb_strtou(s, &s, 10);
 147                         }
 148                 }
 149                 if ((*s != ',' && *s != '\0')
 150                  || bit > end
 151                  || end == UINT_MAX /* bb_strtou returns this on malformed / ERANGE numbers */
 152                  || (stride - 1) > (UINT_MAX / 4)
 153                 /* disallow 0, malformed input, and too large stride prone to overflows */
 154                 ) {
 155                         bb_error_msg_and_die("bad affinity '%s'", aff);
 156                 }
 157                 while (bit <= end && bit < max) {
 158                         mask[bit / BITS_UL] |= (1UL << (bit & MASK_UL));
 159                         bit += stride;
 160                 }
 161                 if (*s == '\0')
 162                         break;
 163                 s++;
 164         }
 165 }
 166 static void print_cpulist(const ul *mask, unsigned mask_size_in_bytes)
 167 {
 168         const ul *mask_end;
 169         const char *delim;
 170         unsigned pos;
 171         ul bit;
 172
 173         mask_end = mask + mask_size_in_bytes / sizeof(mask[0]);
 174         delim = "";
 175         pos = 0;
 176         bit = 1;
 177         for (;;) {
 178                 if (*mask & bit) {
 179                         unsigned onebit = pos + 1;
 180                         printf("%s%u", delim, pos);
 181                         do {
 182                                 pos++;
 183                                 bit <<= 1;
 184                                 if (bit == 0) {
 185                                         mask++;
 186                                         if (mask >= mask_end)
 187                                                 break;
 188                                         bit = 1;
 189                                 }
 190                         } while (*mask & bit);
 191                         if (onebit != pos)
 192                                 printf("-%u", pos - 1);
 193                         delim = ",";
 194                 }
 195                 pos++;
 196                 bit <<= 1;
 197                 if (bit == 0) {
 198                         mask++;
 199                         if (mask >= mask_end)
 200                                 break;
 201                         bit = 1;
 202                 }
 203         }
 204         bb_putchar('\n');
 205 }
 206 #endif
 207
 208 int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
 209 int taskset_main(int argc UNUSED_PARAM, char **argv)
 210 {
 211         ul *mask;
 212         unsigned mask_size_in_bytes;
 213         pid_t pid = 0;
 214         const char *current_new;
 215         char *aff;
 216         unsigned opts;
 217         enum {
 218                 OPT_p = 1 << 0,
 219                 OPT_c = (1 << 1) * ENABLE_FEATURE_TASKSET_CPULIST,
 220         };
 221
 222         /* NB: we mimic util-linux's taskset: -p does not take
 223          * an argument, i.e., "-pN" is NOT valid, only "-p N"!
 224          * Indeed, util-linux-2.13-pre7 uses:
 225          * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */
 226
 227         opts = getopt32(argv, "^+" "p"IF_FEATURE_TASKSET_CPULIST("c")
 228                         "\0" "-1" /* at least 1 arg */);
 229         argv += optind;
 230
 231         aff = *argv++;
 232         if (opts & OPT_p) {
 233                 char *pid_str = aff;
 234                 if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */
 235                         pid_str = *argv; /* NB: *argv != NULL in this case */
 236                 }
 237                 /* else it was just "-p <pid>", and *argv == NULL */
 238                 pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1);
 239         } else {
 240                 /* <aff> <cmd...> */
 241                 if (!*argv)
 242                         bb_show_usage();
 243         }
 244
 245         mask_size_in_bytes = SZOF_UL;
 246         current_new = "current";
 247  print_aff:
 248         mask = get_aff(pid, &mask_size_in_bytes);
 249         if (opts & OPT_p) {
 250 #if ENABLE_FEATURE_TASKSET_CPULIST
 251                 if (opts & OPT_c) {
 252                         printf("pid %d's %s affinity list: ", pid, current_new);
 253                         print_cpulist(mask, mask_size_in_bytes);
 254                 } else
 255 #endif
 256                         printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
 257                                 pid, current_new, from_mask(mask, mask_size_in_bytes));
 258                 if (*argv == NULL) {
 259                         /* Either it was just "-p <pid>",
 260                          * or it was "-p <aff> <pid>" and we came here
 261                          * for the second time (see goto below) */
 262                         return EXIT_SUCCESS;
 263                 }
 264                 *argv = NULL;
 265                 current_new = "new";
 266         }
 267         memset(mask, 0, mask_size_in_bytes);
 268
 269         if (!ENABLE_FEATURE_TASKSET_FANCY) {
 270                 /* Affinity was specified, translate it into mask */
 271                 /* it is always in hex, skip "0x" if it exists */
 272                 if (aff[0] == '0' && (aff[1]|0x20) == 'x')
 273                         aff += 2;
 274                 mask[0] = xstrtoul(aff, 16);
 275         }
 276 #if ENABLE_FEATURE_TASKSET_CPULIST
 277         else if (opts & OPT_c) {
 278                 parse_cpulist(mask, mask_size_in_bytes * 8, aff);
 279         }
 280 #endif
 281         else {
 282                 unsigned i;
 283                 char *last_char;
 284
 285                 /* Affinity was specified, translate it into mask */
 286                 /* it is always in hex, skip "0x" if it exists */
 287                 if (aff[0] == '0' && (aff[1]|0x20) == 'x')
 288                         aff += 2;
 289
 290                 i = 0; /* bit pos in mask[] */
 291
 292                 /* aff is ASCII hex string, accept very long masks in this form.
 293                  * Process hex string AABBCCDD... to ulong mask[]
 294                  * from the rightmost nibble, which is least-significant.
 295                  * Bits not fitting into mask[] are ignored: (example: 1234
 296                  * in 12340000000000000000000000000000000000000ff)
 297                  */
 298                 last_char = strchrnul(aff, '\0');
 299                 while (last_char > aff) {
 300                         char c;
 301                         ul val;
 302
 303                         last_char--;
 304                         c = *last_char;
 305                         if (isdigit(c))
 306                                 val = c - '0';
 307                         else if ((c|0x20) >= 'a' && (c|0x20) <= 'f')
 308                                 val = (c|0x20) - ('a' - 10);
 309                         else
 310                                 bb_error_msg_and_die("bad affinity '%s'", aff);
 311
 312                         if (i < mask_size_in_bytes * 8) {
 313                                 mask[i / BITS_UL] |= val << (i & MASK_UL);
 314                                 //bb_error_msg("bit %d set", i);
 315                         }
 316                         /* else:
 317                          * We can error out here, but we don't.
 318                          * For one, kernel itself ignores bits in mask[]
 319                          * which do not map to any CPUs:
 320                          * if mask[] has one 32-bit long element,
 321                          * but you have only 8 CPUs, all bits beyond first 8
 322                          * are ignored, silently.
 323                          * No point in making bits past 31th to be errors.
 324                          */
 325                         i += 4;
 326                 }
 327         }
 328
 329         /* Set pid's or our own (pid==0) affinity */
 330         if (sched_setaffinity(pid, mask_size_in_bytes, (void*)mask))
 331                 bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);
 332         //bb_error_msg("set mask[0]:%lx", mask[0]);
 333
 334         if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */
 335                 goto print_aff; /* print new affinity and exit */
 336
 337         BB_EXECVP_or_die(argv);
 338 }