1 /* taskset.c - Retrieve or set the CPU affinity of a process.
2 *
3 * Copyright 2012 Elie De Brauwer <[email protected]>
4
5 USE_TASKSET(NEWTOY(taskset, "<1^pa", TOYFLAG_USR|TOYFLAG_BIN))
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 // 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
do_taskset(pid_t pid,int quiet)43 static 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
task_callback(struct dirtree * new)86 static 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
taskset_main(void)94 void 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
nproc_main(void)113 void 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 (smemcmp(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