LXR toybox/toys/other/taskset.c

   1/* taskset.c - Retrieve or set the CPU affinity of a process.
   2 *
   3 * Copyright 2012 Elie De Brauwer <eliedebrauwer@gmail.com>
   4
   5USE_TASKSET(NEWTOY(taskset, "<1^pa", TOYFLAG_USR|TOYFLAG_BIN))
   6USE_NPROC(NEWTOY(nproc, "(all)", TOYFLAG_USR|TOYFLAG_BIN))
   7
   8config NPROC
   9  bool "nproc"
  10  default y
  11  help
  12    usage: nproc [--all]
  13
  14    Print number of processors.
  15
  16    --all       Show all processors, not just ones this task can run on
  17
  18config TASKSET
  19  bool "taskset"
  20  default y
  21  help
  22    usage: taskset [-ap] [mask] [PID | cmd [args...]]
  23
  24    Launch a new task which may only run on certain processors, or change
  25    the processor affinity of an existing PID.
  26
  27    Mask is a hex string where each bit represents a processor the process
  28    is allowed to run on. PID without a mask displays existing affinity.
  29
  30    -p  Set/get the affinity of given PID instead of a new command
  31    -a  Set/get the affinity of all threads of the PID
  32*/
  33
  34#define FOR_taskset
  35#include "toys.h"
  36
  37// mask is array of long which makes layout a bit weird on big endian systems
  38#define sched_setaffinity(pid, size, cpuset) \
  39  syscall(__NR_sched_setaffinity, (pid_t)pid, (size_t)size, (void *)cpuset)
  40#define sched_getaffinity(pid, size, cpuset) \
  41  syscall(__NR_sched_getaffinity, (pid_t)pid, (size_t)size, (void *)cpuset)
  42
  43static void do_taskset(pid_t pid, int quiet)
  44{
  45  unsigned long *mask = (unsigned long *)toybuf;
  46  char *s, *failed = "failed to %s pid %d's affinity";
  47  int i, j, k;
  48
  49  // loop through twice to display before/after affinity masks
  50  for (i=0; ; i++) {
  51    if (!quiet) {
  52      if (-1 == sched_getaffinity(pid, sizeof(toybuf), (void *)mask))
  53        perror_exit(failed, "get", pid);
  54
  55      printf("pid %d's %s affinity mask: ", pid, i ? "new" : "current");
  56
  57      for (j = sizeof(toybuf)/sizeof(long), k = 0; --j>=0;) {
  58        if (k) printf("%0*lx", (int)(2*sizeof(long)), mask[j]);
  59        else if (mask[j]) {
  60          k++;
  61          printf("%lx", mask[j]);
  62        }
  63      }
  64      putchar('\n');
  65    }
  66
  67    if (i || toys.optc < 2) return;
  68
  69    // Convert hex strong to mask[] bits
  70    memset(toybuf, 0, sizeof(toybuf));
  71    k = strlen(s = *toys.optargs);
  72    s += k;
  73    for (j = 0; j<k; j++) {
  74      unsigned long digit = *(--s) - '0';
  75
  76      if (digit > 9) digit = 10 + tolower(*s)-'a';
  77      if (digit > 15) error_exit("bad mask '%s'", *toys.optargs);
  78      mask[j/(2*sizeof(long))] |= digit << 4*(j&((2*sizeof(long))-1));
  79    }
  80
  81    if (-1 == sched_setaffinity(pid, sizeof(toybuf), (void *)mask))
  82      perror_exit(failed, "set", pid);
  83  }
  84}
  85
  86static int task_callback(struct dirtree *new)
  87{
  88  if (!new->parent) return DIRTREE_RECURSE|DIRTREE_SHUTUP|DIRTREE_PROC;
  89  do_taskset(atoi(new->name), 0);
  90
  91  return 0;
  92}
  93
  94void taskset_main(void)
  95{
  96  if (!FLAG(p)) {
  97    if (toys.optc < 2) error_exit("Needs 2 args");
  98    do_taskset(getpid(), 1);
  99    xexec(toys.optargs+1);
 100  } else {
 101    char *c, buf[33];
 102    pid_t pid = strtol(toys.optargs[toys.optc-1], &c, 10);
 103
 104    if (*c) error_exit("Not int %s", toys.optargs[toys.optc-1]);
 105
 106    if (FLAG(a)) {
 107      sprintf(buf, "/proc/%ld/task/", (long)pid);
 108      dirtree_read(buf, task_callback);
 109    } else do_taskset(pid, 0);
 110  }
 111}
 112
 113void nproc_main(void)
 114{
 115  unsigned i, j, nproc = 0;
 116  DIR *dd;
 117
 118  // This can only detect 32768 processors. Call getaffinity and count bits.
 119  if (!toys.optflags && -1!=sched_getaffinity(getpid(), 4096, toybuf))
 120    for (i = 0; i<4096; i++)
 121      if (toybuf[i]) for (j=0; j<8; j++) if (toybuf[i]&(1<<j)) nproc++;
 122
 123  // If getaffinity failed or --all, count cpu entries in sysfs
 124  // (/proc/cpuinfo filters out hot-unplugged CPUs, sysfs doesn't)
 125  if (!nproc && (dd = opendir("/sys/devices/system/cpu"))) {
 126    struct dirent *de;
 127    char *ss;
 128
 129    while ((de = readdir(dd))) {
 130      if (memcmp(de->d_name, "cpu", 3)) continue;
 131      for (ss = de->d_name+3; isdigit(*ss); ss++);
 132      if (!*ss) nproc++;
 133    }
 134    closedir(dd);
 135  }
 136
 137  xprintf("%u\n", nproc ? : 1);
 138}
 139