bigbiff bigbiff | e60683a | 2013-02-22 20:55:50 -0500 | [diff] [blame] | 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 | |
| 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 | |
| 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 | |
| 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, "%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 | */ |
| 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 | */ |
| 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 | */ |
| 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 | |
| 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); |
bigbiff | 7b4c7a6 | 2015-01-01 19:44:14 -0500 | [diff] [blame^] | 391 | free(mask); |
bigbiff bigbiff | e60683a | 2013-02-22 20:55:50 -0500 | [diff] [blame] | 392 | free(range); |
| 393 | cpuset_free(set); |
| 394 | |
| 395 | return EXIT_SUCCESS; |
| 396 | |
| 397 | usage_err: |
| 398 | fprintf(stderr, |
| 399 | "usage: %s [--ncpus <num>] --mask <mask> | --range <list>", |
| 400 | program_invocation_short_name); |
| 401 | exit(EXIT_FAILURE); |
| 402 | } |
| 403 | #endif |