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 //usage:#define taskset_trivial_usage
  10 //usage:       "[-p] [MASK] [PID | PROG ARGS]"
  11 //usage:#define taskset_full_usage "\n\n"
  12 //usage:       "Set or get CPU affinity\n"
  13 //usage:     "\nOptions:"
  14 //usage:     "\n        -p      Operate on an existing PID"
  15 //usage:
  16 //usage:#define taskset_example_usage
  17 //usage:       "$ taskset 0x7 ./dgemm_test&\n"
  18 //usage:       "$ taskset -p 0x1 $!\n"
  19 //usage:       "pid 4790's current affinity mask: 7\n"
  20 //usage:       "pid 4790's new affinity mask: 1\n"
  21 //usage:       "$ taskset 0x7 /bin/sh -c './taskset -p 0x1 $$'\n"
  22 //usage:       "pid 6671's current affinity mask: 1\n"
  23 //usage:       "pid 6671's new affinity mask: 1\n"
  24 //usage:       "$ taskset -p 1\n"
  25 //usage:       "pid 1's current affinity mask: 3\n"
  26
  27 #include <sched.h>
  28 #include "libbb.h"
  29
  30 #if ENABLE_FEATURE_TASKSET_FANCY
  31 #define TASKSET_PRINTF_MASK "%s"
  32 /* craft a string from the mask */
  33 static char *from_cpuset(cpu_set_t *mask)
  34 {
  35         int i;
  36         char *ret = NULL;
  37         char *str = xzalloc((CPU_SETSIZE / 4) + 1); /* we will leak it */
  38
  39         for (i = CPU_SETSIZE - 4; i >= 0; i -= 4) {
  40                 int val = 0;
  41                 int off;
  42                 for (off = 0; off <= 3; ++off)
  43                         if (CPU_ISSET(i + off, mask))
  44                                 val |= 1 << off;
  45                 if (!ret && val)
  46                         ret = str;
  47                 *str++ = bb_hexdigits_upcase[val] | 0x20;
  48         }
  49         return ret;
  50 }
  51 #else
  52 #define TASKSET_PRINTF_MASK "%llx"
  53 static unsigned long long from_cpuset(cpu_set_t *mask)
  54 {
  55         struct BUG_CPU_SETSIZE_is_too_small {
  56                 char BUG_CPU_SETSIZE_is_too_small[
  57                         CPU_SETSIZE < sizeof(int) ? -1 : 1];
  58         };
  59         char *p = (void*)mask;
  60
  61         /* Take the least significant bits. Careful!
  62          * Consider both CPU_SETSIZE=4 and CPU_SETSIZE=1024 cases
  63          */
  64 #if BB_BIG_ENDIAN
  65         /* For big endian, it means LAST bits */
  66         if (CPU_SETSIZE < sizeof(long))
  67                 p += CPU_SETSIZE - sizeof(int);
  68         else if (CPU_SETSIZE < sizeof(long long))
  69                 p += CPU_SETSIZE - sizeof(long);
  70         else
  71                 p += CPU_SETSIZE - sizeof(long long);
  72 #endif
  73         if (CPU_SETSIZE < sizeof(long))
  74                 return *(unsigned*)p;
  75         if (CPU_SETSIZE < sizeof(long long))
  76                 return *(unsigned long*)p;
  77         return *(unsigned long long*)p;
  78 }
  79 #endif
  80
  81
  82 int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
  83 int taskset_main(int argc UNUSED_PARAM, char **argv)
  84 {
  85         cpu_set_t mask;
  86         pid_t pid = 0;
  87         unsigned opt_p;
  88         const char *current_new;
  89         char *pid_str;
  90         char *aff = aff; /* for compiler */
  91
  92         /* NB: we mimic util-linux's taskset: -p does not take
  93          * an argument, i.e., "-pN" is NOT valid, only "-p N"!
  94          * Indeed, util-linux-2.13-pre7 uses:
  95          * getopt_long(argc, argv, "+pchV", ...), not "...p:..." */
  96
  97         opt_complementary = "-1"; /* at least 1 arg */
  98         opt_p = getopt32(argv, "+p");
  99         argv += optind;
 100
 101         if (opt_p) {
 102                 pid_str = *argv++;
 103                 if (*argv) { /* "-p <aff> <pid> ...rest.is.ignored..." */
 104                         aff = pid_str;
 105                         pid_str = *argv; /* NB: *argv != NULL in this case */
 106                 }
 107                 /* else it was just "-p <pid>", and *argv == NULL */
 108                 pid = xatoul_range(pid_str, 1, ((unsigned)(pid_t)ULONG_MAX) >> 1);
 109         } else {
 110                 aff = *argv++; /* <aff> <cmd...> */
 111                 if (!*argv)
 112                         bb_show_usage();
 113         }
 114
 115         current_new = "current\0new";
 116         if (opt_p) {
 117  print_aff:
 118                 if (sched_getaffinity(pid, sizeof(mask), &mask) < 0)
 119                         bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
 120                 printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
 121                                 pid, current_new, from_cpuset(&mask));
 122                 if (!*argv) {
 123                         /* Either it was just "-p <pid>",
 124                          * or it was "-p <aff> <pid>" and we came here
 125                          * for the second time (see goto below) */
 126                         return EXIT_SUCCESS;
 127                 }
 128                 *argv = NULL;
 129                 current_new += 8; /* "new" */
 130         }
 131
 132         { /* Affinity was specified, translate it into cpu_set_t */
 133                 unsigned i;
 134                 /* Do not allow zero mask: */
 135                 unsigned long long m = xstrtoull_range(aff, 0, 1, ULLONG_MAX);
 136                 enum { CNT_BIT = CPU_SETSIZE < sizeof(m)*8 ? CPU_SETSIZE : sizeof(m)*8 };
 137
 138                 CPU_ZERO(&mask);
 139                 for (i = 0; i < CNT_BIT; i++) {
 140                         unsigned long long bit = (1ULL << i);
 141                         if (bit & m)
 142                                 CPU_SET(i, &mask);
 143                 }
 144         }
 145
 146         /* Set pid's or our own (pid==0) affinity */
 147         if (sched_setaffinity(pid, sizeof(mask), &mask))
 148                 bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);
 149
 150         if (!argv[0]) /* "-p <aff> <pid> [...ignored...]" */
 151                 goto print_aff; /* print new affinity and exit */
 152
 153         BB_EXECVP_or_die(argv);
 154 }