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 #include <sys/syscall.h>
24
25 #include "cpuset.h"
26 #include "c.h"
27
val_to_char(int v)28 static inline int val_to_char(int v)
29 {
30 if (v >= 0 && v < 10)
31 return '0' + v;
32 else if (v >= 10 && v < 16)
33 return ('a' - 10) + v;
34 else
35 return -1;
36 }
37
char_to_val(int c)38 static inline int char_to_val(int c)
39 {
40 int cl;
41
42 cl = tolower(c);
43 if (c >= '0' && c <= '9')
44 return c - '0';
45 else if (cl >= 'a' && cl <= 'f')
46 return cl + (10 - 'a');
47 else
48 return -1;
49 }
50
nexttoken(const char * q,int sep)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 */
get_max_number_of_cpus(void)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 */
cpuset_alloc(int ncpus,size_t * setsize,size_t * nbits)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
cpuset_free(cpu_set_t * set)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 */
__cpuset_count_s(size_t setsize,const cpu_set_t * set)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 */
cpulist_create(char * str,size_t len,cpu_set_t * set,size_t setsize)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, "%zd,", i);
173 else if (run == 1) {
174 rlen = snprintf(ptr, len, "%zd,%zd,", i, i + 1);
175 i++;
176 } else {
177 rlen = snprintf(ptr, len, "%zd-%zd,", i, i + run);
178 i += run;
179 }
180 if (rlen < 0 || (size_t) rlen + 1 > len)
181 return NULL;
182 ptr += rlen;
183 if (rlen > 0 && len > (size_t) rlen)
184 len -= rlen;
185 else
186 len = 0;
187 }
188 }
189 ptr -= entry_made;
190 *ptr = '\0';
191
192 return str;
193 }
194
195 /*
196 * Returns string with CPU mask.
197 */
cpumask_create(char * str,size_t len,cpu_set_t * set,size_t setsize)198 char *cpumask_create(char *str, size_t len,
199 cpu_set_t *set, size_t setsize)
200 {
201 char *ptr = str;
202 char *ret = NULL;
203 int cpu;
204
205 for (cpu = cpuset_nbits(setsize) - 4; cpu >= 0; cpu -= 4) {
206 char val = 0;
207
208 if (len == (size_t) (ptr - str))
209 break;
210
211 if (CPU_ISSET_S(cpu, setsize, set))
212 val |= 1;
213 if (CPU_ISSET_S(cpu + 1, setsize, set))
214 val |= 2;
215 if (CPU_ISSET_S(cpu + 2, setsize, set))
216 val |= 4;
217 if (CPU_ISSET_S(cpu + 3, setsize, set))
218 val |= 8;
219
220 if (!ret && val)
221 ret = ptr;
222 *ptr++ = val_to_char(val);
223 }
224 *ptr = '\0';
225 return ret ? ret : ptr - 1;
226 }
227
228 /*
229 * Parses string with CPUs mask.
230 */
cpumask_parse(const char * str,cpu_set_t * set,size_t setsize)231 int cpumask_parse(const char *str, cpu_set_t *set, size_t setsize)
232 {
233 int len = strlen(str);
234 const char *ptr = str + len - 1;
235 int cpu = 0;
236
237 /* skip 0x, it's all hex anyway */
238 if (len > 1 && !memcmp(str, "0x", 2L))
239 str += 2;
240
241 CPU_ZERO_S(setsize, set);
242
243 while (ptr >= str) {
244 char val;
245
246 /* cpu masks in /sys uses comma as a separator */
247 if (*ptr == ',')
248 ptr--;
249
250 val = char_to_val(*ptr);
251 if (val == (char) -1)
252 return -1;
253 if (val & 1)
254 CPU_SET_S(cpu, setsize, set);
255 if (val & 2)
256 CPU_SET_S(cpu + 1, setsize, set);
257 if (val & 4)
258 CPU_SET_S(cpu + 2, setsize, set);
259 if (val & 8)
260 CPU_SET_S(cpu + 3, setsize, set);
261 len--;
262 ptr--;
263 cpu += 4;
264 }
265
266 return 0;
267 }
268
269 /*
270 * Parses string with list of CPU ranges.
271 * Returns 0 on success.
272 * Returns 1 on error.
273 * Returns 2 if fail is set and a cpu number passed in the list doesn't fit
274 * into the cpu_set. If fail is not set cpu numbers that do not fit are
275 * ignored and 0 is returned instead.
276 */
cpulist_parse(const char * str,cpu_set_t * set,size_t setsize,int fail)277 int cpulist_parse(const char *str, cpu_set_t *set, size_t setsize, int fail)
278 {
279 size_t max = cpuset_nbits(setsize);
280 const char *p, *q;
281 int r = 0;
282
283 q = str;
284 CPU_ZERO_S(setsize, set);
285
286 while (p = q, q = nexttoken(q, ','), p) {
287 unsigned int a; /* beginning of range */
288 unsigned int b; /* end of range */
289 unsigned int s; /* stride */
290 const char *c1, *c2;
291 char c;
292
293 if ((r = sscanf(p, "%u%c", &a, &c)) < 1)
294 return 1;
295 b = a;
296 s = 1;
297
298 c1 = nexttoken(p, '-');
299 c2 = nexttoken(p, ',');
300 if (c1 != NULL && (c2 == NULL || c1 < c2)) {
301 if ((r = sscanf(c1, "%u%c", &b, &c)) < 1)
302 return 1;
303 c1 = nexttoken(c1, ':');
304 if (c1 != NULL && (c2 == NULL || c1 < c2)) {
305 if ((r = sscanf(c1, "%u%c", &s, &c)) < 1)
306 return 1;
307 if (s == 0)
308 return 1;
309 }
310 }
311
312 if (!(a <= b))
313 return 1;
314 while (a <= b) {
315 if (fail && (a >= max))
316 return 2;
317 CPU_SET_S(a, setsize, set);
318 a += s;
319 }
320 }
321
322 if (r == 2)
323 return 1;
324 return 0;
325 }
326
327 #ifdef TEST_PROGRAM
328
329 #include <getopt.h>
330
main(int argc,char * argv[])331 int main(int argc, char *argv[])
332 {
333 cpu_set_t *set;
334 size_t setsize, buflen, nbits;
335 char *buf, *mask = NULL, *range = NULL;
336 int ncpus = 2048, rc, c;
337
338 static const struct option longopts[] = {
339 { "ncpus", 1, 0, 'n' },
340 { "mask", 1, 0, 'm' },
341 { "range", 1, 0, 'r' },
342 { NULL, 0, 0, 0 }
343 };
344
345 while ((c = getopt_long(argc, argv, "n:m:r:", longopts, NULL)) != -1) {
346 switch(c) {
347 case 'n':
348 ncpus = atoi(optarg);
349 break;
350 case 'm':
351 mask = strdup(optarg);
352 break;
353 case 'r':
354 range = strdup(optarg);
355 break;
356 default:
357 goto usage_err;
358 }
359 }
360
361 if (!mask && !range)
362 goto usage_err;
363
364 set = cpuset_alloc(ncpus, &setsize, &nbits);
365 if (!set)
366 err(EXIT_FAILURE, "failed to allocate cpu set");
367
368 /*
369 fprintf(stderr, "ncpus: %d, cpuset bits: %zd, cpuset bytes: %zd\n",
370 ncpus, nbits, setsize);
371 */
372
373 buflen = 7 * nbits;
374 buf = malloc(buflen);
375 if (!buf)
376 err(EXIT_FAILURE, "failed to allocate cpu set buffer");
377
378 if (mask)
379 rc = cpumask_parse(mask, set, setsize);
380 else
381 rc = cpulist_parse(range, set, setsize, 0);
382
383 if (rc)
384 errx(EXIT_FAILURE, "failed to parse string: %s", mask ? : range);
385
386 printf("%-15s = %15s ", mask ? : range,
387 cpumask_create(buf, buflen, set, setsize));
388 printf("[%s]\n", cpulist_create(buf, buflen, set, setsize));
389
390 free(buf);
391 free(range);
392 cpuset_free(set);
393
394 return EXIT_SUCCESS;
395
396 usage_err:
397 fprintf(stderr,
398 "usage: %s [--ncpus <num>] --mask <mask> | --range <list>",
399 program_invocation_short_name);
400 exit(EXIT_FAILURE);
401 }
402 #endif
403