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
7 config TASKSET
8 bool "taskset"
9 default y
10 help
11 usage: taskset [-ap] [mask] [PID | cmd [args...]]
12
13 Launch a new task which may only run on certain processors, or change
14 the processor affinity of an exisitng PID.
15
16 Mask is a hex string where each bit represents a processor the process
17 is allowed to run on. PID without a mask displays existing affinity.
18
19 -p Set/get the affinity of given PID instead of a new command.
20 -a Set/get the affinity of all threads of the PID.
21 */
22
23 #define FOR_taskset
24 #include "toys.h"
25
26 #include <sys/syscall.h>
27 #define sched_setaffinity(pid, size, cpuset) \
28 syscall(__NR_sched_setaffinity, (pid_t)pid, (size_t)size, (void *)cpuset)
29 #define sched_getaffinity(pid, size, cpuset) \
30 syscall(__NR_sched_getaffinity, (pid_t)pid, (size_t)size, (void *)cpuset)
31
32 // mask is an array of long, which makes the layout a bit weird on big
33 // endian systems but as long as it's consistent...
34
do_taskset(pid_t pid,int quiet)35 static void do_taskset(pid_t pid, int quiet)
36 {
37 unsigned long *mask = (unsigned long *)toybuf;
38 char *s = *toys.optargs, *failed = "failed to %s %d's affinity";
39 int i, j, k;
40
41 for (i=0; ; i++) {
42 if (!quiet) {
43 int j = sizeof(toybuf), flag = 0;
44
45 if (-1 == sched_getaffinity(pid, sizeof(toybuf), (void *)mask))
46 perror_exit(failed, "get", pid);
47
48 printf("pid %d's %s affinity mask: ", pid, i ? "new" : "current");
49
50 while (j--) {
51 int x = 255 & (mask[j/sizeof(long)] >> (8*(j&(sizeof(long)-1))));
52
53 if (flag) printf("%02x", x);
54 else if (x) {
55 flag++;
56 printf("%x", x);
57 }
58 }
59 putchar('\n');
60 }
61
62 if (i || toys.optc < 2) return;
63
64 memset(toybuf, 0, sizeof(toybuf));
65 k = strlen(s = *toys.optargs);
66 s += k;
67 for (j = 0; j<k; j++) {
68 unsigned long digit = *(--s) - '0';
69
70 if (digit > 9) digit = 10 + tolower(*s)-'a';
71 if (digit > 15) error_exit("bad mask '%s'", *toys.optargs);
72 mask[j/(2*sizeof(long))] |= digit << 4*(j&((2*sizeof(long))-1));
73 }
74
75 if (-1 == sched_setaffinity(pid, sizeof(toybuf), (void *)mask))
76 perror_exit(failed, "set", pid);
77 }
78 }
79
task_callback(struct dirtree * new)80 static int task_callback(struct dirtree *new)
81 {
82 if (!new->parent) return DIRTREE_RECURSE;
83 if (isdigit(*new->name)) do_taskset(atoi(new->name), 0);
84
85 return 0;
86 }
87
taskset_main(void)88 void taskset_main(void)
89 {
90 if (!(toys.optflags & FLAG_p)) {
91 if (toys.optc < 2) error_exit("Needs 2 args");
92 do_taskset(getpid(), 1);
93 xexec(toys.optargs+1);
94 } else {
95 char *c;
96 pid_t pid = strtol(toys.optargs[toys.optc-1], &c, 10);
97
98 if (*c) error_exit("Not int %s", toys.optargs[1]);
99
100 if (toys.optflags & FLAG_a) {
101 char buf[33];
102 sprintf(buf, "/proc/%ld/task/", (long)pid);
103 dirtree_read(buf, task_callback);
104 } else do_taskset(pid, 0);
105 }
106 }
107