src/shared/cpu-set-util.c

   1 /* SPDX-License-Identifier: LGPL-2.1+ */
   2
   3 #include <errno.h>
   4 #include <stddef.h>
   5 #include <stdio.h>
   6 #include <syslog.h>
   7
   8 #include "alloc-util.h"
   9 #include "cpu-set-util.h"
  10 #include "extract-word.h"
  11 #include "log.h"
  12 #include "macro.h"
  13 #include "memory-util.h"
  14 #include "parse-util.h"
  15 #include "string-util.h"
  16
  17 char* cpu_set_to_string(const CPUSet *a) {
  18         _cleanup_free_ char *str = NULL;
  19         size_t allocated = 0, len = 0;
  20         int i, r;
  21
  22         for (i = 0; (size_t) i < a->allocated * 8; i++) {
  23                 if (!CPU_ISSET_S(i, a->allocated, a->set))
  24                         continue;
  25
  26                 if (!GREEDY_REALLOC(str, allocated, len + 1 + DECIMAL_STR_MAX(int)))
  27                         return NULL;
  28
  29                 r = sprintf(str + len, len > 0 ? " %d" : "%d", i);
  30                 assert_se(r > 0);
  31                 len += r;
  32         }
  33
  34         return TAKE_PTR(str) ?: strdup("");
  35 }
  36
  37 cpu_set_t* cpu_set_malloc(unsigned *ncpus) {
  38         cpu_set_t *c;
  39         unsigned n = 1024;
  40
  41         /* Allocates the cpuset in the right size */
  42
  43         for (;;) {
  44                 c = CPU_ALLOC(n);
  45                 if (!c)
  46                         return NULL;
  47
  48                 if (sched_getaffinity(0, CPU_ALLOC_SIZE(n), c) >= 0) {
  49                         CPU_ZERO_S(CPU_ALLOC_SIZE(n), c);
  50
  51                         if (ncpus)
  52                                 *ncpus = n;
  53
  54                         return c;
  55                 }
  56
  57                 CPU_FREE(c);
  58
  59                 if (errno != EINVAL)
  60                         return NULL;
  61
  62                 n *= 2;
  63         }
  64 }
  65
  66 static int cpu_set_realloc(CPUSet *cpu_set, unsigned ncpus) {
  67         size_t need;
  68
  69         assert(cpu_set);
  70
  71         need = CPU_ALLOC_SIZE(ncpus);
  72         if (need > cpu_set->allocated) {
  73                 cpu_set_t *t;
  74
  75                 t = realloc(cpu_set->set, need);
  76                 if (!t)
  77                         return -ENOMEM;
  78
  79                 memzero((uint8_t*) t + cpu_set->allocated, need - cpu_set->allocated);
  80
  81                 cpu_set->set = t;
  82                 cpu_set->allocated = need;
  83         }
  84
  85         return 0;
  86 }
  87
  88 static int cpu_set_add(CPUSet *cpu_set, unsigned cpu) {
  89         int r;
  90
  91         if (cpu >= 8192)
  92                 /* As of kernel 5.1, CONFIG_NR_CPUS can be set to 8192 on PowerPC */
  93                 return -ERANGE;
  94
  95         r = cpu_set_realloc(cpu_set, cpu + 1);
  96         if (r < 0)
  97                 return r;
  98
  99         CPU_SET_S(cpu, cpu_set->allocated, cpu_set->set);
 100         return 0;
 101 }
 102
 103 int cpu_set_add_all(CPUSet *a, const CPUSet *b) {
 104         int r;
 105
 106         /* Do this backwards, so if we fail, we fail before changing anything. */
 107         for (unsigned cpu_p1 = b->allocated * 8; cpu_p1 > 0; cpu_p1--)
 108                 if (CPU_ISSET_S(cpu_p1 - 1, b->allocated, b->set)) {
 109                         r = cpu_set_add(a, cpu_p1 - 1);
 110                         if (r < 0)
 111                                 return r;
 112                 }
 113
 114         return 0;
 115 }
 116
 117 int parse_cpu_set_full(
 118                 const char *rvalue,
 119                 CPUSet *cpu_set,
 120                 bool warn,
 121                 const char *unit,
 122                 const char *filename,
 123                 unsigned line,
 124                 const char *lvalue) {
 125
 126         _cleanup_(cpu_set_reset) CPUSet c = {};
 127         const char *p = rvalue;
 128
 129         assert(p);
 130
 131         for (;;) {
 132                 _cleanup_free_ char *word = NULL;
 133                 unsigned cpu_lower, cpu_upper;
 134                 int r;
 135
 136                 r = extract_first_word(&p, &word, WHITESPACE ",", EXTRACT_QUOTES);
 137                 if (r == -ENOMEM)
 138                         return warn ? log_oom() : -ENOMEM;
 139                 if (r < 0)
 140                         return warn ? log_syntax(unit, LOG_ERR, filename, line, r, "Invalid value for %s: %s", lvalue, rvalue) : r;
 141                 if (r == 0)
 142                         break;
 143
 144                 r = parse_range(word, &cpu_lower, &cpu_upper);
 145                 if (r < 0)
 146                         return warn ? log_syntax(unit, LOG_ERR, filename, line, r, "Failed to parse CPU affinity '%s'", word) : r;
 147
 148                 if (cpu_lower > cpu_upper) {
 149                         if (warn)
 150                                 log_syntax(unit, LOG_WARNING, filename, line, 0, "Range '%s' is invalid, %u > %u, ignoring.",
 151                                            word, cpu_lower, cpu_upper);
 152
 153                         /* Make sure something is allocated, to distinguish this from the empty case */
 154                         r = cpu_set_realloc(&c, 1);
 155                         if (r < 0)
 156                                 return r;
 157                 }
 158
 159                 for (unsigned cpu_p1 = MIN(cpu_upper, UINT_MAX-1) + 1; cpu_p1 > cpu_lower; cpu_p1--) {
 160                         r = cpu_set_add(&c, cpu_p1 - 1);
 161                         if (r < 0)
 162                                 return warn ? log_syntax(unit, LOG_ERR, filename, line, r,
 163                                                          "Cannot add CPU %u to set: %m", cpu_p1 - 1) : r;
 164                 }
 165         }
 166
 167         /* On success, transfer ownership to the output variable */
 168         *cpu_set = c;
 169         c = (CPUSet) {};
 170
 171         return 0;
 172 }
 173
 174 int parse_cpu_set_extend(
 175                 const char *rvalue,
 176                 CPUSet *old,
 177                 bool warn,
 178                 const char *unit,
 179                 const char *filename,
 180                 unsigned line,
 181                 const char *lvalue) {
 182
 183         _cleanup_(cpu_set_reset) CPUSet cpuset = {};
 184         int r;
 185
 186         r = parse_cpu_set_full(rvalue, &cpuset, true, unit, filename, line, lvalue);
 187         if (r < 0)
 188                 return r;
 189
 190         if (!cpuset.set) {
 191                 /* An empty assignment resets the CPU list */
 192                 cpu_set_reset(old);
 193                 return 0;
 194         }
 195
 196         if (!old->set) {
 197                 *old = cpuset;
 198                 cpuset = (CPUSet) {};
 199                 return 0;
 200         }
 201
 202         return cpu_set_add_all(old, &cpuset);
 203 }
 204
 205 int cpus_in_affinity_mask(void) {
 206         size_t n = 16;
 207         int r;
 208
 209         for (;;) {
 210                 cpu_set_t *c;
 211
 212                 c = CPU_ALLOC(n);
 213                 if (!c)
 214                         return -ENOMEM;
 215
 216                 if (sched_getaffinity(0, CPU_ALLOC_SIZE(n), c) >= 0) {
 217                         int k;
 218
 219                         k = CPU_COUNT_S(CPU_ALLOC_SIZE(n), c);
 220                         CPU_FREE(c);
 221
 222                         if (k <= 0)
 223                                 return -EINVAL;
 224
 225                         return k;
 226                 }
 227
 228                 r = -errno;
 229                 CPU_FREE(c);
 230
 231                 if (r != -EINVAL)
 232                         return r;
 233                 if (n > SIZE_MAX/2)
 234                         return -ENOMEM;
 235                 n *= 2;
 236         }
 237 }