1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
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
44
45
46static 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
91static 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
99void 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
119void nproc_main(void)
120{
121 unsigned i, j, nproc = 0;
122
123
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
130 if (!nproc) nproc = sysconf(_SC_NPROCESSORS_CONF);
131
132 xprintf("%u\n", nproc);
133}
134