lib/cpuset.c

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