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#include <sched.h>
61#include "libbb.h"
62
63typedef unsigned long ul;
64#define SZOF_UL (unsigned)(sizeof(ul))
65#define BITS_UL (unsigned)(sizeof(ul)*8)
66#define MASK_UL (unsigned)(sizeof(ul)*8 - 1)
67
68#if ENABLE_FEATURE_TASKSET_FANCY
69#define TASKSET_PRINTF_MASK "%s"
70
71static char *from_mask(const ul *mask, unsigned sz_in_bytes)
72{
73 char *str = xzalloc((sz_in_bytes+1) * 2);
74 char *p = str;
75 for (;;) {
76 ul v = *mask++;
77 if (SZOF_UL == 4)
78 p += sprintf(p, "%08lx", v);
79 if (SZOF_UL == 8)
80 p += sprintf(p, "%016lx", v);
81 if (SZOF_UL == 16)
82 p += sprintf(p, "%032lx", v);
83 sz_in_bytes -= SZOF_UL;
84 if ((int)sz_in_bytes <= 0)
85 break;
86 }
87 while (str[0] == '0' && str[1])
88 str++;
89 return str;
90}
91#else
92#define TASKSET_PRINTF_MASK "%lx"
93static unsigned long from_mask(ul *mask, unsigned sz_in_bytes UNUSED_PARAM)
94{
95 return *mask;
96}
97#endif
98
99static unsigned long *get_aff(int pid, unsigned *sz)
100{
101 int r;
102 unsigned long *mask = NULL;
103 unsigned sz_in_bytes = *sz;
104
105 for (;;) {
106 mask = xrealloc(mask, sz_in_bytes);
107 r = sched_getaffinity(pid, sz_in_bytes, (void*)mask);
108 if (r == 0)
109 break;
110 sz_in_bytes *= 2;
111 if (errno == EINVAL && (int)sz_in_bytes > 0)
112 continue;
113 bb_perror_msg_and_die("can't %cet pid %d's affinity", 'g', pid);
114 }
115
116 *sz = sz_in_bytes;
117 return mask;
118}
119
120#if ENABLE_FEATURE_TASKSET_CPULIST
121
122
123
124
125
126
127
128
129
130
131
132
133
134static void parse_cpulist(ul *mask, unsigned max, char *s)
135{
136 char *aff = s;
137 for (;;) {
138 unsigned bit, end;
139 unsigned stride = 1;
140
141 bit = end = bb_strtou(s, &s, 10);
142 if (*s == '-') {
143 s++;
144 end = bb_strtou(s, &s, 10);
145 if (*s == ':') {
146 s++;
147 stride = bb_strtou(s, &s, 10);
148 }
149 }
150 if ((*s != ',' && *s != '\0')
151 || bit > end
152 || end == UINT_MAX
153 || (stride - 1) > (UINT_MAX / 4)
154
155 ) {
156 bb_error_msg_and_die("bad affinity '%s'", aff);
157 }
158 while (bit <= end && bit < max) {
159 mask[bit / BITS_UL] |= (1UL << (bit & MASK_UL));
160 bit += stride;
161 }
162 if (*s == '\0')
163 break;
164 s++;
165 }
166}
167static void print_cpulist(const ul *mask, unsigned mask_size_in_bytes)
168{
169 const ul *mask_end;
170 const char *delim;
171 unsigned pos;
172 ul bit;
173
174 mask_end = mask + mask_size_in_bytes / sizeof(mask[0]);
175 delim = "";
176 pos = 0;
177 bit = 1;
178 for (;;) {
179 if (*mask & bit) {
180 unsigned onebit = pos + 1;
181 printf("%s%u", delim, pos);
182 do {
183 pos++;
184 bit <<= 1;
185 if (bit == 0) {
186 mask++;
187 if (mask >= mask_end)
188 break;
189 bit = 1;
190 }
191 } while (*mask & bit);
192 if (onebit != pos)
193 printf("-%u", pos - 1);
194 delim = ",";
195 }
196 pos++;
197 bit <<= 1;
198 if (bit == 0) {
199 mask++;
200 if (mask >= mask_end)
201 break;
202 bit = 1;
203 }
204 }
205 bb_putchar('\n');
206}
207#endif
208
209enum {
210 OPT_p = 1 << 0,
211 OPT_a = 1 << 1,
212 OPT_c = (1 << 2) * ENABLE_FEATURE_TASKSET_CPULIST,
213};
214
215static int process_pid_str(const char *pid_str, unsigned opts, char *aff)
216{
217 ul *mask;
218 unsigned mask_size_in_bytes;
219 const char *current_new;
220 pid_t pid = !pid_str ? 0 : xatou_range(pid_str, 1, INT_MAX);
221
222 mask_size_in_bytes = SZOF_UL;
223 current_new = "current";
224 print_aff:
225 mask = get_aff(pid, &mask_size_in_bytes);
226 if (opts & OPT_p) {
227#if ENABLE_FEATURE_TASKSET_CPULIST
228 if (opts & OPT_c) {
229 printf("pid %d's %s affinity list: ", pid, current_new);
230 print_cpulist(mask, mask_size_in_bytes);
231 } else
232#endif
233 printf("pid %d's %s affinity mask: "TASKSET_PRINTF_MASK"\n",
234 pid, current_new, from_mask(mask, mask_size_in_bytes));
235 if (!aff) {
236
237
238
239 return 0;
240 }
241 current_new = "new";
242 }
243 memset(mask, 0, mask_size_in_bytes);
244
245 if (!ENABLE_FEATURE_TASKSET_FANCY) {
246
247
248 if (aff[0] == '0' && (aff[1]|0x20) == 'x')
249 aff += 2;
250 mask[0] = xstrtoul(aff, 16);
251 }
252#if ENABLE_FEATURE_TASKSET_CPULIST
253 else if (opts & OPT_c) {
254 parse_cpulist(mask, mask_size_in_bytes * 8, aff);
255 }
256#endif
257 else {
258 unsigned i;
259 char *last_char;
260
261
262
263 if (aff[0] == '0' && (aff[1]|0x20) == 'x')
264 aff += 2;
265
266 i = 0;
267
268
269
270
271
272
273
274 last_char = strchrnul(aff, '\0');
275 while (last_char > aff) {
276 char c;
277 ul val;
278
279 last_char--;
280 c = *last_char;
281 if (isdigit(c))
282 val = c - '0';
283 else if ((c|0x20) >= 'a' && (c|0x20) <= 'f')
284 val = (c|0x20) - ('a' - 10);
285 else
286 bb_error_msg_and_die("bad affinity '%s'", aff);
287
288 if (i < mask_size_in_bytes * 8) {
289 mask[i / BITS_UL] |= val << (i & MASK_UL);
290
291 }
292
293
294
295
296
297
298
299
300
301 i += 4;
302 }
303 }
304
305
306 if (sched_setaffinity(pid, mask_size_in_bytes, (void*)mask))
307 bb_perror_msg_and_die("can't %cet pid %d's affinity", 's', pid);
308
309
310 if ((opts & OPT_p) && aff) {
311 aff = NULL;
312 goto print_aff;
313 }
314 return 0;
315}
316
317static int FAST_FUNC iter(const char *dn UNUSED_PARAM, struct dirent *ent, void *aff)
318{
319 if (isdigit(ent->d_name[0]))
320 return process_pid_str(ent->d_name, option_mask32, aff);
321 return 0;
322}
323
324int taskset_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
325int taskset_main(int argc UNUSED_PARAM, char **argv)
326{
327 const char *pid_str;
328 char *aff;
329 unsigned opts;
330
331
332
333
334
335
336 opts = getopt32(argv, "^+" "pa"IF_FEATURE_TASKSET_CPULIST("c")
337 "\0" "-1" );
338 argv += optind;
339
340 aff = *argv++;
341 if (!(opts & OPT_p)) {
342
343 if (!*argv)
344 bb_show_usage();
345 process_pid_str(NULL, opts, aff);
346 BB_EXECVP_or_die(argv);
347 }
348
349 pid_str = aff;
350 if (*argv)
351 pid_str = *argv;
352 else
353 aff = NULL;
354
355 if (opts & OPT_a) {
356 char *dn;
357 int r;
358
359 dn = xasprintf("/proc/%s/task", pid_str);
360 r = iterate_on_dir(dn, iter, aff);
361 IF_FEATURE_CLEAN_UP(free(dn);)
362 if (r == 0)
363 return r;
364
365 }
366 return process_pid_str(pid_str, opts, aff);
367}
368