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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
/* taskset.c - Retrieve or set the CPU affinity of a process.
*
* Copyright 2012 Elie De Brauwer <eliedebrauwer@gmail.com>
USE_TASKSET(NEWTOY(taskset, "<1^pa", TOYFLAG_BIN|TOYFLAG_STAYROOT))
USE_NPROC(NEWTOY(nproc, "(all)", TOYFLAG_USR|TOYFLAG_BIN))
config NPROC
bool "nproc"
default y
help
usage: nproc [--all]
Print number of processors.
--all Show all processors, not just ones this task can run on.
config TASKSET
bool "taskset"
default y
help
usage: taskset [-ap] [mask] [PID | cmd [args...]]
Launch a new task which may only run on certain processors, or change
the processor affinity of an exisitng PID.
Mask is a hex string where each bit represents a processor the process
is allowed to run on. PID without a mask displays existing affinity.
-p Set/get the affinity of given PID instead of a new command.
-a Set/get the affinity of all threads of the PID.
*/
#define FOR_taskset
#include "toys.h"
#include <sys/syscall.h>
#define sched_setaffinity(pid, size, cpuset) \
syscall(__NR_sched_setaffinity, (pid_t)pid, (size_t)size, (void *)cpuset)
#define sched_getaffinity(pid, size, cpuset) \
syscall(__NR_sched_getaffinity, (pid_t)pid, (size_t)size, (void *)cpuset)
GLOBALS(
int nproc;
)
// mask is an array of long, which makes the layout a bit weird on big
// endian systems but as long as it's consistent...
static void do_taskset(pid_t pid, int quiet)
{
unsigned long *mask = (unsigned long *)toybuf;
char *s = *toys.optargs, *failed = "failed to %s %d's affinity";
int i, j, k;
for (i=0; ; i++) {
if (!quiet) {
int j = sizeof(toybuf), flag = 0;
if (-1 == sched_getaffinity(pid, sizeof(toybuf), (void *)mask))
perror_exit(failed, "get", pid);
printf("pid %d's %s affinity mask: ", pid, i ? "new" : "current");
while (j--) {
int x = 255 & (mask[j/sizeof(long)] >> (8*(j&(sizeof(long)-1))));
if (flag) printf("%02x", x);
else if (x) {
flag++;
printf("%x", x);
}
}
putchar('\n');
}
if (i || toys.optc < 2) return;
memset(toybuf, 0, sizeof(toybuf));
k = strlen(s = *toys.optargs);
s += k;
for (j = 0; j<k; j++) {
unsigned long digit = *(--s) - '0';
if (digit > 9) digit = 10 + tolower(*s)-'a';
if (digit > 15) error_exit("bad mask '%s'", *toys.optargs);
mask[j/(2*sizeof(long))] |= digit << 4*(j&((2*sizeof(long))-1));
}
if (-1 == sched_setaffinity(pid, sizeof(toybuf), (void *)mask))
perror_exit(failed, "set", pid);
}
}
static int task_callback(struct dirtree *new)
{
if (!new->parent) return DIRTREE_RECURSE;
if (isdigit(*new->name)) do_taskset(atoi(new->name), 0);
return 0;
}
void taskset_main(void)
{
if (!(toys.optflags & FLAG_p)) {
if (toys.optc < 2) error_exit("Needs 2 args");
do_taskset(getpid(), 1);
xexec(toys.optargs+1);
} else {
char *c;
pid_t pid = strtol(toys.optargs[toys.optc-1], &c, 10);
if (*c) error_exit("Not int %s", toys.optargs[1]);
if (toys.optflags & FLAG_a) {
char buf[33];
sprintf(buf, "/proc/%ld/task/", (long)pid);
dirtree_read(buf, task_callback);
} else do_taskset(pid, 0);
}
}
static int do_nproc(struct dirtree *new)
{
if (!new->parent) return DIRTREE_RECURSE;
if (!strncmp(new->name, "cpu", 3) && isdigit(new->name[3])) TT.nproc++;
return 0;
}
void nproc_main(void)
{
int i, j;
// This can only detect 32768 processors. Call getaffinity and count bits.
if (!toys.optflags && -1!=sched_getaffinity(getpid(), 4096, toybuf)) {
for (i = 0; i<4096; i++)
if (toybuf[i])
for (j=0; j<8; j++)
if (toybuf[i]&(1<<j))
TT.nproc++;
}
// If getaffinity failed or --all, count cpu entries in proc
if (!TT.nproc) dirtree_read("/sys/devices/system/cpu", do_nproc);
xprintf("%d\n", TT.nproc);
}
|
