miscutils/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
   9 //config:config TASKSET
  10 //config:       bool "taskset"
  11 //config:       default n  # doesn't build on some non-x86 targets (m68k)
  12 //config:       help
  13 //config:         Retrieve or set a processes's CPU affinity.
  14 //config:         This requires sched_{g,s}etaffinity support in your libc.
  15 //config:
  16 //config:config FEATURE_TASKSET_FANCY
  17 //config:       bool "Fancy output"
  18 //config:       default y
  19 //config:       depends on TASKSET
  20 //config:       help
  21 //config:         Add code for fancy output. This merely silences a compiler-warning
  22 //config:         and adds about 135 Bytes. May be needed for machines with alot
  23 //config:         of CPUs.
  24
  25 //applet:IF_TASKSET(APPLET(taskset, BB_DIR_USR_BIN, BB_SUID_DROP))
  26 //kbuild:lib-$(CONFIG_TASKSET) += taskset.o
  27
  28 //usage:#define taskset_trivial_usage
  29 //usage:       "[-p] [MASK] [PID | PROG ARGS]"
  30 //usage:#define taskset_full_usage "\n\n"
  31 //usage:       "Set or get CPU affinity\n"
  32 //usage:     "\n        -p      Operate on an existing PID"
  33 //usage:
  34 //usage:#define taskset_example_usage
  35 //usage:       "$ taskset 0x7 ./dgemm_test&\n"
  36 //usage:       "$ taskset -p 0x1 $!\n"
  37 //usage:       "pid 4790's current affinity mask: 7\n"
  38 //usage:       "pid 4790's new affinity mask: 1\n"
  39 //usage:       "$ taskset 0x7 /bin/sh -c './taskset -p 0x1 $$'\n"
  40 //usage:       "pid 6671's current affinity mask: 1\n"
  41 //usage:       "pid 6671's new affinity mask: 1\n"
  42 //usage:       "$ taskset -p 1\n"
  43 //usage:       "pid 1's current affinity mask: 3\n"
  44 /*
  45  Not yet implemented:
  46  * -a/--all-tasks (affect all threads)
  47  * -c/--cpu-list  (specify CPUs via "1,3,5-7")
  48  */
  49
  50 #include <sched.h>
  51 #include "libbb.h"
  52
  53 #if ENABLE_FEATURE_TASKSET_FANCY
  54 #define TASKSET_PRINTF_MASK "%s"
  55 /* craft a string from the mask */
  56 static char *from_cpuset(cpu_set_t *mask)
  57 {
  58         int i;
  59         char *ret = NULL;
  60         char *str = xzalloc((CPU_SETSIZE / 4) + 1); /* we will leak it */
  61
  62         for (i = CPU_SETSIZE - 4; i >= 0; i -= 4) {
  63                 int val = 0;
  64                 int off;
  65                 for (off = 0; off <= 3; ++off)
  66                         if (CPU_ISSET(i + off, mask))
  67                                 val |= 1 << off;
  68                 if (!ret && val)
  69                         ret = str;
  70                 *str++ = bb_hexdigits_upcase[val] | 0x20;
  71         }
  72         return ret;
  73 }
  74 #else
  75 #define TASKSET_PRINTF_MASK "%llx"
  76 static unsigned long long from_cpuset(cpu_set_t *mask)
  77 {
  78         struct BUG_CPU_SETSIZE_is_too_small {
  79                 char BUG_CPU_SETSIZE_is_too_small[
  80                         CPU_SETSIZE < sizeof(int) ? -1 : 1];
  81         };
  82         char *p = (void*)mask;
  83
  84         /* Take the least significant bits. Careful!
  85          * Consider both CPU_SETSIZE=4 and CPU_SETSIZE=1024 cases
  86          */
  87 #if BB_BIG_ENDIAN
  88         /* For big endian, it means LAST bits */
  89         if (CPU_SETSIZE < sizeof(long))
  90                 p += CPU_SETSIZE - sizeof(int);
  91         else if (CPU_SETSIZE < sizeof(long long))
  92                 p += CPU_SETSIZE - sizeof(long);
  93         else
  94                 p += CPU_SETSIZE - sizeof(long long);
  95 #endif
  96         if (CPU_SETSIZE < sizeof(long))
  97                 return *(unsigned*)p;
  98         if (CPU_SETSIZE < sizeof(long long))
  99                 return *(unsigned long*)p;
 100         return *(unsigned long long*)p;
 101 }
 102 #endif
 103
 104
 105 int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
 106 int taskset_main(int argc UNUSED_PARAM, char **argv)
 107 {
 108         cpu_set_t mask;
 109         pid_t pid = 0;
 110         unsigned opt_p;
 111         const char *current_new;
 112         char *pid_str;
 113         char *aff = aff; /* for compiler */
 114
 115         /* NB: we mimic util-linux's taskset: -p does not take
 116          * an argument, i.e., "-pN" is NOT valid, only "-p N"!
 117          * Indeed, util-linux-2.13-pre7 uses:
 118          * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */
 119
 120         opt_complementary = "-1"; /* at least 1 arg */
 121         opt_p = getopt32(argv, "+p");
 122         argv += optind;
 123
 124         if (opt_p) {
 125                 pid_str = *argv++;
 126                 if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */
 127                         aff = pid_str;
 128                         pid_str = *argv; /* NB: *argv != NULL in this case */
 129                 }
 130                 /* else it was just "-p <pid>", and *argv == NULL */
 131                 pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1);
 132         } else {
 133                 aff = *argv++; /* <aff> <cmd...> */
 134                 if (!*argv)
 135                         bb_show_usage();
 136         }
 137
 138         current_new = "current\0new";
 139         if (opt_p) {
 140  print_aff:
 141                 if (sched_getaffinity(pid, sizeof(mask), &mask) < 0)
 142                         bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
 143                 printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
 144                                 pid, current_new, from_cpuset(&mask));
 145                 if (!*argv) {
 146                         /* Either it was just "-p <pid>",
 147                          * or it was "-p <aff> <pid>" and we came here
 148                          * for the second time (see goto below) */
 149                         return EXIT_SUCCESS;
 150                 }
 151                 *argv = NULL;
 152                 current_new += 8; /* "new" */
 153         }
 154
 155         /* Affinity was specified, translate it into cpu_set_t */
 156         CPU_ZERO(&mask);
 157         if (!ENABLE_FEATURE_TASKSET_FANCY) {
 158                 unsigned i;
 159                 unsigned long long m;
 160
 161                 /* Do not allow zero mask: */
 162                 m = xstrtoull_range(aff, 0, 1, ULLONG_MAX);
 163                 i = 0;
 164                 do {
 165                         if (m & 1)
 166                                 CPU_SET(i, &mask);
 167                         i++;
 168                         m >>= 1;
 169                 } while (m != 0);
 170         } else {
 171                 unsigned i;
 172                 char *last_byte;
 173                 char *bin;
 174                 uint8_t bit_in_byte;
 175
 176                 /* Cheap way to get "long enough" buffer */
 177                 bin = xstrdup(aff);
 178
 179                 if (aff[0] != '0' && (aff[1]|0x20) != 'x') {
 180 /* TODO: decimal/octal masks are still limited to 2^64 */
 181                         unsigned long long m = xstrtoull_range(aff, 0, 1, ULLONG_MAX);
 182                         bin += strlen(bin);
 183                         last_byte = bin - 1;
 184                         while (m) {
 185                                 *--bin = m & 0xff;
 186                                 m >>= 8;
 187                         }
 188                 } else {
 189                         /* aff is "0x.....", we accept very long masks in this form */
 190                         last_byte = hex2bin(bin, aff + 2, INT_MAX);
 191                         if (!last_byte) {
 192  bad_aff:
 193                                 bb_error_msg_and_die("bad affinity '%s'", aff);
 194                         }
 195                         last_byte--; /* now points to the last byte */
 196                 }
 197
 198                 i = 0;
 199                 bit_in_byte = 1;
 200                 while (last_byte >= bin) {
 201                         if (bit_in_byte & *last_byte) {
 202                                 if (i >= CPU_SETSIZE)
 203                                         goto bad_aff;
 204                                 CPU_SET(i, &mask);
 205                                 //bb_error_msg("bit %d set", i);
 206                         }
 207                         i++;
 208                         /* bit_in_byte is uint8_t! & 0xff is implied */
 209                         bit_in_byte = (bit_in_byte << 1);
 210                         if (!bit_in_byte) {
 211                                 bit_in_byte = 1;
 212                                 last_byte--;
 213                         }
 214                 }
 215         }
 216
 217         /* Set pid's or our own (pid==0) affinity */
 218         if (sched_setaffinity(pid, sizeof(mask), &mask))
 219                 bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);
 220
 221         if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */
 222                 goto print_aff; /* print new affinity and exit */
 223
 224         BB_EXECVP_or_die(argv);
 225 }