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