lib/cpuset.c

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