LXR busybox/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 */
  56static 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"
  76static 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
 105int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
 106int 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}
 226