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         BUILD_BUG_ON(CPU_SETSIZE < 8*sizeof(int));
  79
  80         /* Take the least significant bits. Assume cpu_set_t is
  81          * implemented as an array of unsigned long or unsigned
  82          * int.
  83          */
  84         if (CPU_SETSIZE < 8*sizeof(long))
  85                 return *(unsigned*)mask;
  86         if (CPU_SETSIZE < 8*sizeof(long long))
  87                 return *(unsigned long*)mask;
  88 # if BB_BIG_ENDIAN
  89         if (sizeof(long long) > sizeof(long)) {
  90                 /* We can put two long in the long long, but they have to
  91                  * be swapped: the least significant word comes first in the
  92                  * array */
  93                 unsigned long *p = (void*)mask;
  94                 return p[0] + ((unsigned long long)p[1] << (8*sizeof(long)));
  95         }
  96 # endif
  97         return *(unsigned long long*)mask;
  98 }
  99 #endif
 100
 101
 102 int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
 103 int taskset_main(int argc UNUSED_PARAM, char **argv)
 104 {
 105         cpu_set_t mask;
 106         pid_t pid = 0;
 107         unsigned opt_p;
 108         const char *current_new;
 109         char *pid_str;
 110         char *aff = aff; /* for compiler */
 111
 112         /* NB: we mimic util-linux's taskset: -p does not take
 113          * an argument, i.e., "-pN" is NOT valid, only "-p N"!
 114          * Indeed, util-linux-2.13-pre7 uses:
 115          * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */
 116
 117         opt_complementary = "-1"; /* at least 1 arg */
 118         opt_p = getopt32(argv, "+p");
 119         argv += optind;
 120
 121         if (opt_p) {
 122                 pid_str = *argv++;
 123                 if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */
 124                         aff = pid_str;
 125                         pid_str = *argv; /* NB: *argv != NULL in this case */
 126                 }
 127                 /* else it was just "-p <pid>", and *argv == NULL */
 128                 pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1);
 129         } else {
 130                 aff = *argv++; /* <aff> <cmd...> */
 131                 if (!*argv)
 132                         bb_show_usage();
 133         }
 134
 135         current_new = "current\0new";
 136         if (opt_p) {
 137  print_aff:
 138                 if (sched_getaffinity(pid, sizeof(mask), &mask) < 0)
 139                         bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
 140                 printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
 141                                 pid, current_new, from_cpuset(&mask));
 142                 if (!*argv) {
 143                         /* Either it was just "-p <pid>",
 144                          * or it was "-p <aff> <pid>" and we came here
 145                          * for the second time (see goto below) */
 146                         return EXIT_SUCCESS;
 147                 }
 148                 *argv = NULL;
 149                 current_new += 8; /* "new" */
 150         }
 151
 152         /* Affinity was specified, translate it into cpu_set_t */
 153         CPU_ZERO(&mask);
 154         if (!ENABLE_FEATURE_TASKSET_FANCY) {
 155                 unsigned i;
 156                 unsigned long long m;
 157
 158                 /* Do not allow zero mask: */
 159                 m = xstrtoull_range(aff, 0, 1, ULLONG_MAX);
 160                 i = 0;
 161                 do {
 162                         if (m & 1)
 163                                 CPU_SET(i, &mask);
 164                         i++;
 165                         m >>= 1;
 166                 } while (m != 0);
 167         } else {
 168                 unsigned i;
 169                 char *last_byte;
 170                 char *bin;
 171                 uint8_t bit_in_byte;
 172
 173                 /* Cheap way to get "long enough" buffer */
 174                 bin = xstrdup(aff);
 175
 176                 if (aff[0] != '0' || (aff[1]|0x20) != 'x') {
 177 /* TODO: decimal/octal masks are still limited to 2^64 */
 178                         unsigned long long m = xstrtoull_range(aff, 0, 1, ULLONG_MAX);
 179                         bin += strlen(bin);
 180                         last_byte = bin - 1;
 181                         while (m) {
 182                                 *--bin = m & 0xff;
 183                                 m >>= 8;
 184                         }
 185                 } else {
 186                         /* aff is "0x.....", we accept very long masks in this form */
 187                         last_byte = hex2bin(bin, aff + 2, INT_MAX);
 188                         if (!last_byte) {
 189  bad_aff:
 190                                 bb_error_msg_and_die("bad affinity '%s'", aff);
 191                         }
 192                         last_byte--; /* now points to the last byte */
 193                 }
 194
 195                 i = 0;
 196                 bit_in_byte = 1;
 197                 while (last_byte >= bin) {
 198                         if (bit_in_byte & *last_byte) {
 199                                 if (i >= CPU_SETSIZE)
 200                                         goto bad_aff;
 201                                 CPU_SET(i, &mask);
 202                                 //bb_error_msg("bit %d set", i);
 203                         }
 204                         i++;
 205                         /* bit_in_byte is uint8_t! & 0xff is implied */
 206                         bit_in_byte = (bit_in_byte << 1);
 207                         if (!bit_in_byte) {
 208                                 bit_in_byte = 1;
 209                                 last_byte--;
 210                         }
 211                 }
 212         }
 213
 214         /* Set pid's or our own (pid==0) affinity */
 215         if (sched_setaffinity(pid, sizeof(mask), &mask))
 216                 bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);
 217
 218         if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */
 219                 goto print_aff; /* print new affinity and exit */
 220
 221         BB_EXECVP_or_die(argv);
 222 }