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