1 /*
2 * Terminology:
3 *
4 * cpuset - (libc) cpu_set_t data structure represents set of CPUs
5 * cpumask - string with hex mask (e.g. "0x00000001")
6 * cpulist - string with CPU ranges (e.g. "0-3,5,7,8")
7 *
8 * Based on code from taskset.c and Linux kernel.
9 *
10 * This file may be redistributed under the terms of the
11 * GNU Lesser General Public License.
12 *
13 * Copyright (C) 2010 Karel Zak <kzak@redhat.com>
14 */
15
16 #include <stdio.h>
17 #include <stdlib.h>
18 #include <unistd.h>
19 #include <sched.h>
20 #include <errno.h>
21 #include <string.h>
22 #include <ctype.h>
23 #ifdef HAVE_SYS_SYSCALL_H
24 #include <sys/syscall.h>
25 #endif
26
27 #include "cpuset.h"
28 #include "c.h"
29
30 static inline int val_to_char(int v)
31 {
32 if (v >= 0 && v < 10)
33 return '0' + v;
34 if (v >= 10 && v < 16)
35 return ('a' - 10) + v;
36 return -1;
37 }
38
39 static inline int char_to_val(int c)
40 {
41 int cl;
42
43 if (c >= '0' && c <= '9')
44 return c - '0';
45 cl = tolower(c);
46 if (cl >= 'a' && cl <= 'f')
47 return cl + (10 - 'a');
48 return -1;
49 }
50
51 static const char *nexttoken(const char *q, int sep)
52 {
53 if (q)
54 q = strchr(q, sep);
55 if (q)
56 q++;
57 return q;
58 }
59
60 /*
61 * Number of bits in a CPU bitmask on current system
62 */
63 int get_max_number_of_cpus(void)
64 {
65 #ifdef SYS_sched_getaffinity
66 int n, cpus = 2048;
67 size_t setsize;
68 cpu_set_t *set = cpuset_alloc(cpus, &setsize, NULL);
69
70 if (!set)
71 return -1; /* error */
72
73 for (;;) {
74 CPU_ZERO_S(setsize, set);
75
76 /* the library version does not return size of cpumask_t */
77 n = syscall(SYS_sched_getaffinity, 0, setsize, set);
78
79 if (n < 0 && errno == EINVAL && cpus < 1024 * 1024) {
80 cpuset_free(set);
81 cpus *= 2;
82 set = cpuset_alloc(cpus, &setsize, NULL);
83 if (!set)
84 return -1; /* error */
85 continue;
86 }
87 cpuset_free(set);
88 return n * 8;
89 }
90 #endif
91 return -1;
92 }
93
94 /*
95 * Allocates a new set for ncpus and returns size in bytes and size in bits
96 */
97 cpu_set_t *cpuset_alloc(int ncpus, size_t *setsize, size_t *nbits)
98 {
99 cpu_set_t *set = CPU_ALLOC(ncpus);
100
101 if (!set)
102 return NULL;
103 if (setsize)
104 *setsize = CPU_ALLOC_SIZE(ncpus);
105 if (nbits)
106 *nbits = cpuset_nbits(CPU_ALLOC_SIZE(ncpus));
107 return set;
108 }
109
110 void cpuset_free(cpu_set_t *set)
111 {
112 CPU_FREE(set);
113 }
114
115 #if !HAVE_DECL_CPU_ALLOC
116 /* Please, use CPU_COUNT_S() macro. This is fallback */
117 int __cpuset_count_s(size_t setsize, const cpu_set_t *set)
118 {
119 int s = 0;
120 const __cpu_mask *p = set->__bits;
121 const __cpu_mask *end = &set->__bits[setsize / sizeof (__cpu_mask)];
122
123 while (p < end) {
124 __cpu_mask l = *p++;
125
126 if (l == 0)
127 continue;
128 # if LONG_BIT > 32
129 l = (l & 0x5555555555555555ul) + ((l >> 1) & 0x5555555555555555ul);
130 l = (l & 0x3333333333333333ul) + ((l >> 2) & 0x3333333333333333ul);
131 l = (l & 0x0f0f0f0f0f0f0f0ful) + ((l >> 4) & 0x0f0f0f0f0f0f0f0ful);
132 l = (l & 0x00ff00ff00ff00fful) + ((l >> 8) & 0x00ff00ff00ff00fful);
133 l = (l & 0x0000ffff0000fffful) + ((l >> 16) & 0x0000ffff0000fffful);
134 l = (l & 0x00000000fffffffful) + ((l >> 32) & 0x00000000fffffffful);
135 # else
136 l = (l & 0x55555555ul) + ((l >> 1) & 0x55555555ul);
137 l = (l & 0x33333333ul) + ((l >> 2) & 0x33333333ul);
138 l = (l & 0x0f0f0f0ful) + ((l >> 4) & 0x0f0f0f0ful);
139 l = (l & 0x00ff00fful) + ((l >> 8) & 0x00ff00fful);
140 l = (l & 0x0000fffful) + ((l >> 16) & 0x0000fffful);
141 # endif
142 s += l;
143 }
144 return s;
145 }
146 #endif
147
148 /*
149 * Returns human readable representation of the cpuset. The output format is
150 * a list of CPUs with ranges (for example, "0,1,3-9").
151 */
152 char *cpulist_create(char *str, size_t len,
153 cpu_set_t *set, size_t setsize)
154 {
155 size_t i;
156 char *ptr = str;
157 int entry_made = 0;
158 size_t max = cpuset_nbits(setsize);
159
160 for (i = 0; i < max; i++) {
161 if (CPU_ISSET_S(i, setsize, set)) {
162 int rlen;
163 size_t j, run = 0;
164 entry_made = 1;
165 for (j = i + 1; j < max; j++) {
166 if (CPU_ISSET_S(j, setsize, set))
167 run++;
168 else
169 break;
170 }
171 if (!run)
172 rlen = snprintf(ptr, len, "%zu,", i);
173 else if (run == 1) {
174 rlen = snprintf(ptr, len, "%zu,%zu,", i, i + 1);
175 i++;
176 } else {
177 rlen = snprintf(ptr, len, "%zu-%zu,", i, i + run);
178 i += run;
179 }
180 if (rlen < 0 || (size_t) rlen >= len)
181 return NULL;
182 ptr += rlen;
183 len -= rlen;
184 }
185 }
186 ptr -= entry_made;
187 *ptr = '\0';
188
189 return str;
190 }
191
192 /*
193 * Returns string with CPU mask.
194 */
195 char *cpumask_create(char *str, size_t len,
196 cpu_set_t *set, size_t setsize)
197 {
198 char *ptr = str;
199 char *ret = NULL;
200 int cpu;
201
202 for (cpu = cpuset_nbits(setsize) - 4; cpu >= 0; cpu -= 4) {
203 char val = 0;
204
205 if (len == (size_t) (ptr - str))
206 break;
207
208 if (CPU_ISSET_S(cpu, setsize, set))
209 val |= 1;
210 if (CPU_ISSET_S(cpu + 1, setsize, set))
211 val |= 2;
212 if (CPU_ISSET_S(cpu + 2, setsize, set))
213 val |= 4;
214 if (CPU_ISSET_S(cpu + 3, setsize, set))
215 val |= 8;
216
217 if (!ret && val)
218 ret = ptr;
219 *ptr++ = val_to_char(val);
220 }
221 *ptr = '\0';
222 return ret ? ret : ptr - 1;
223 }
224
225 /*
226 * Parses string with CPUs mask.
227 */
228 int cpumask_parse(const char *str, cpu_set_t *set, size_t setsize)
229 {
230 int len = strlen(str);
231 const char *ptr = str + len - 1;
232 int cpu = 0;
233
234 /* skip 0x, it's all hex anyway */
235 if (len > 1 && !memcmp(str, "0x", 2L))
236 str += 2;
237
238 CPU_ZERO_S(setsize, set);
239
240 while (ptr >= str) {
241 char val;
242
243 /* cpu masks in /sys uses comma as a separator */
244 if (*ptr == ',')
245 ptr--;
246
247 val = char_to_val(*ptr);
248 if (val == (char) -1)
249 return -1;
250 if (val & 1)
251 CPU_SET_S(cpu, setsize, set);
252 if (val & 2)
253 CPU_SET_S(cpu + 1, setsize, set);
254 if (val & 4)
255 CPU_SET_S(cpu + 2, setsize, set);
256 if (val & 8)
257 CPU_SET_S(cpu + 3, setsize, set);
258 ptr--;
259 cpu += 4;
260 }
261
262 return 0;
263 }
264
265 static int nextnumber(const char *str, char **end, unsigned int *result)
266 {
267 errno = 0;
268 if (str == NULL || *str == '\0' || !isdigit(*str))
269 return -EINVAL;
270 *result = (unsigned int) strtoul(str, end, 10);
271 if (errno)
272 return -errno;
273 if (str == *end)
274 return -EINVAL;
275 return 0;
276 }
277
278 /*
279 * Parses string with list of CPU ranges.
280 * Returns 0 on success.
281 * Returns 1 on error.
282 * Returns 2 if fail is set and a cpu number passed in the list doesn't fit
283 * into the cpu_set. If fail is not set cpu numbers that do not fit are
284 * ignored and 0 is returned instead.
285 */
286 int cpulist_parse(const char *str, cpu_set_t *set, size_t setsize, int fail)
287 {
288 size_t max = cpuset_nbits(setsize);
289 const char *p, *q;
290 char *end = NULL;
291
292 q = str;
293 CPU_ZERO_S(setsize, set);
294
295 while (p = q, q = nexttoken(q, ','), p) {
296 unsigned int a; /* beginning of range */
297 unsigned int b; /* end of range */
298 unsigned int s; /* stride */
299 const char *c1, *c2;
300
301 if (nextnumber(p, &end, &a) != 0)
302 return 1;
303 b = a;
304 s = 1;
305 p = end;
306
307 c1 = nexttoken(p, '-');
308 c2 = nexttoken(p, ',');
309
310 if (c1 != NULL && (c2 == NULL || c1 < c2)) {
311 if (nextnumber(c1, &end, &b) != 0)
312 return 1;
313
314 c1 = end && *end ? nexttoken(end, ':') : NULL;
315
316 if (c1 != NULL && (c2 == NULL || c1 < c2)) {
317 if (nextnumber(c1, &end, &s) != 0)
318 return 1;
319 if (s == 0)
320 return 1;
321 }
322 }
323
324 if (!(a <= b))
325 return 1;
326 while (a <= b) {
327 if (fail && (a >= max))
328 return 2;
329 CPU_SET_S(a, setsize, set);
330 a += s;
331 }
332 }
333
334 if (end && *end)
335 return 1;
336 return 0;
337 }
338
339 #ifdef TEST_PROGRAM_CPUSET
340
341 #include <getopt.h>
342
343 int main(int argc, char *argv[])
344 {
345 cpu_set_t *set;
346 size_t setsize, buflen, nbits;
347 char *buf, *mask = NULL, *range = NULL;
348 int ncpus = 2048, rc, c;
349
350 static const struct option longopts[] = {
351 { "ncpus", 1, NULL, 'n' },
352 { "mask", 1, NULL, 'm' },
353 { "range", 1, NULL, 'r' },
354 { NULL, 0, NULL, 0 }
355 };
356
357 while ((c = getopt_long(argc, argv, "n:m:r:", longopts, NULL)) != -1) {
358 switch(c) {
359 case 'n':
360 ncpus = atoi(optarg);
361 break;
362 case 'm':
363 mask = strdup(optarg);
364 break;
365 case 'r':
366 range = strdup(optarg);
367 break;
368 default:
369 goto usage_err;
370 }
371 }
372
373 if (!mask && !range)
374 goto usage_err;
375
376 set = cpuset_alloc(ncpus, &setsize, &nbits);
377 if (!set)
378 err(EXIT_FAILURE, "failed to allocate cpu set");
379
380 /*
381 fprintf(stderr, "ncpus: %d, cpuset bits: %zd, cpuset bytes: %zd\n",
382 ncpus, nbits, setsize);
383 */
384
385 buflen = 7 * nbits;
386 buf = malloc(buflen);
387 if (!buf)
388 err(EXIT_FAILURE, "failed to allocate cpu set buffer");
389
390 if (mask)
391 rc = cpumask_parse(mask, set, setsize);
392 else
393 rc = cpulist_parse(range, set, setsize, 0);
394
395 if (rc)
396 errx(EXIT_FAILURE, "failed to parse string: %s", mask ? : range);
397
398 printf("%-15s = %15s ", mask ? : range,
399 cpumask_create(buf, buflen, set, setsize));
400 printf("[%s]\n", cpulist_create(buf, buflen, set, setsize));
401
402 free(buf);
403 free(mask);
404 free(range);
405 cpuset_free(set);
406
407 return EXIT_SUCCESS;
408
409 usage_err:
410 fprintf(stderr,
411 "usage: %s [--ncpus <num>] --mask <mask> | --range <list>\n",
412 program_invocation_short_name);
413 exit(EXIT_FAILURE);
414 }
415 #endif