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