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
   5 USE_TASKSET(NEWTOY(taskset, "<1^pa", TOYFLAG_BIN|TOYFLAG_STAYROOT))
   6 USE_NPROC(NEWTOY(nproc, "(all)", TOYFLAG_USR|TOYFLAG_BIN))
   7
   8 config 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
  18 config 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 exisitng 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 #include <sys/syscall.h>
  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
  43 // mask is an array of long, which makes the layout a bit weird on big
  44 // endian systems but as long as it's consistent...
  45
  46 static void do_taskset(pid_t pid, int quiet)
  47 {
  48   unsigned long *mask = (unsigned long *)toybuf;
  49   char *s = *toys.optargs, *failed = "failed to %s %d's affinity";
  50   int i, j, k;
  51
  52   for (i=0; ; i++) {
  53     if (!quiet) {
  54       int j = sizeof(toybuf), flag = 0;
  55
  56       if (-1 == sched_getaffinity(pid, sizeof(toybuf), (void *)mask))
  57         perror_exit(failed, "get", pid);
  58
  59       printf("pid %d's %s affinity mask: ", pid, i ? "new" : "current");
  60
  61       while (j--) {
  62         int x = 255 & (mask[j/sizeof(long)] >> (8*(j&(sizeof(long)-1))));
  63
  64         if (flag) printf("%02x", x);
  65         else if (x) {
  66           flag++;
  67           printf("%x", x);
  68         }
  69       }
  70       putchar('\n');
  71     }
  72
  73     if (i || toys.optc < 2) return;
  74
  75     memset(toybuf, 0, sizeof(toybuf));
  76     k = strlen(s = *toys.optargs);
  77     s += k;
  78     for (j = 0; j<k; j++) {
  79       unsigned long digit = *(--s) - '0';
  80
  81       if (digit > 9) digit = 10 + tolower(*s)-'a';
  82       if (digit > 15) error_exit("bad mask '%s'", *toys.optargs);
  83       mask[j/(2*sizeof(long))] |= digit << 4*(j&((2*sizeof(long))-1));
  84     }
  85
  86     if (-1 == sched_setaffinity(pid, sizeof(toybuf), (void *)mask))
  87       perror_exit(failed, "set", pid);
  88   }
  89 }
  90
  91 static int task_callback(struct dirtree *new)
  92 {
  93   if (!new->parent) return DIRTREE_RECURSE;
  94   if (isdigit(*new->name)) do_taskset(atoi(new->name), 0);
  95
  96   return 0;
  97 }
  98
  99 void taskset_main(void)
 100 {
 101   if (!(toys.optflags & FLAG_p)) {
 102     if (toys.optc < 2) error_exit("Needs 2 args");
 103     do_taskset(getpid(), 1);
 104     xexec(toys.optargs+1);
 105   } else {
 106     char *c;
 107     pid_t pid = strtol(toys.optargs[toys.optc-1], &c, 10);
 108
 109     if (*c) error_exit("Not int %s", toys.optargs[1]);
 110
 111     if (toys.optflags & FLAG_a) {
 112       char buf[33];
 113       sprintf(buf, "/proc/%ld/task/", (long)pid);
 114       dirtree_read(buf, task_callback);
 115     } else do_taskset(pid, 0);
 116   }
 117 }
 118
 119 void nproc_main(void)
 120 {
 121   unsigned i, j, nproc = 0;
 122
 123   // This can only detect 32768 processors. Call getaffinity and count bits.
 124   if (!toys.optflags && -1!=sched_getaffinity(getpid(), 4096, toybuf)) {
 125     for (i = 0; i<4096; i++)
 126       if (toybuf[i]) for (j=0; j<8; j++) if (toybuf[i]&(1<<j)) nproc++;
 127   }
 128
 129   // If getaffinity failed or --all, count cpu entries in proc
 130   if (!nproc) nproc = sysconf(_SC_NPROCESSORS_CONF);
 131
 132   xprintf("%u\n", nproc);
 133 }