Merge pull request #12628 from keszybz/dbus-execute

Rework cpu affinity parsing

4 files changed, 276 insertions, 64 deletions

diff --git a/src/shared/bus-unit-util.c b/src/shared/bus-unit-util.c
index 2b425efc9c..fb86391975 100644
--- a/src/shared/bus-unit-util.c
+++ b/src/shared/bus-unit-util.c
@@ -992,13 +992,19 @@ static int bus_append_execute_property(sd_bus_message *m, const char *field, con
         }
 
         if (streq(field, "CPUAffinity")) {
-                _cleanup_cpu_free_ cpu_set_t *cpuset = NULL;
+                _cleanup_(cpu_set_reset) CPUSet cpuset = {};
+                _cleanup_free_ uint8_t *array = NULL;
+                size_t allocated;
 
                 r = parse_cpu_set(eq, &cpuset);
                 if (r < 0)
                         return log_error_errno(r, "Failed to parse %s value: %s", field, eq);
 
-                return bus_append_byte_array(m, field, cpuset, CPU_ALLOC_SIZE(r));
+                r = cpu_set_to_dbus(&cpuset, &array, &allocated);
+                if (r < 0)
+                        return log_error_errno(r, "Failed to serialize CPUAffinity: %m");
+
+                return bus_append_byte_array(m, field, array, allocated);
         }
 
         if (STR_IN_SET(field, "RestrictAddressFamilies", "SystemCallFilter")) {
diff --git a/src/shared/condition.c b/src/shared/condition.c
index d1ac9de756..2d521bc8c6 100644
--- a/src/shared/condition.c
+++ b/src/shared/condition.c
@@ -21,6 +21,7 @@
 #include "cap-list.h"
 #include "cgroup-util.h"
 #include "condition.h"
+#include "cpu-set-util.h"
 #include "efivars.h"
 #include "env-file.h"
 #include "extract-word.h"
diff --git a/src/shared/cpu-set-util.c b/src/shared/cpu-set-util.c
index 9a789ae756..7c8f4d42d2 100644
--- a/src/shared/cpu-set-util.c
+++ b/src/shared/cpu-set-util.c
@@ -2,6 +2,7 @@
 
 #include <errno.h>
 #include <stddef.h>
+#include <stdio.h>
 #include <syslog.h>
 
 #include "alloc-util.h"
@@ -9,56 +10,141 @@
 #include "extract-word.h"
 #include "log.h"
 #include "macro.h"
+#include "memory-util.h"
 #include "parse-util.h"
 #include "string-util.h"
 
-cpu_set_t* cpu_set_malloc(unsigned *ncpus) {
-        cpu_set_t *c;
-        unsigned n = 1024;
+char* cpu_set_to_string(const CPUSet *a) {
+        _cleanup_free_ char *str = NULL;
+        size_t allocated = 0, len = 0;
+        int i, r;
 
-        /* Allocates the cpuset in the right size */
+        for (i = 0; (size_t) i < a->allocated * 8; i++) {
+                if (!CPU_ISSET_S(i, a->allocated, a->set))
+                        continue;
 
-        for (;;) {
-                c = CPU_ALLOC(n);
-                if (!c)
+                if (!GREEDY_REALLOC(str, allocated, len + 1 + DECIMAL_STR_MAX(int)))
                         return NULL;
 
-                if (sched_getaffinity(0, CPU_ALLOC_SIZE(n), c) >= 0) {
-                        CPU_ZERO_S(CPU_ALLOC_SIZE(n), c);
+                r = sprintf(str + len, len > 0 ? " %d" : "%d", i);
+                assert_se(r > 0);
+                len += r;
+        }
 
-                        if (ncpus)
-                                *ncpus = n;
+        return TAKE_PTR(str) ?: strdup("");
+}
 
-                        return c;
-                }
+char *cpu_set_to_range_string(const CPUSet *set) {
+        unsigned range_start = 0, range_end;
+        _cleanup_free_ char *str = NULL;
+        size_t allocated = 0, len = 0;
+        bool in_range = false;
+        int r;
 
-                CPU_FREE(c);
+        for (unsigned i = 0; i < set->allocated * 8; i++)
+                if (CPU_ISSET_S(i, set->allocated, set->set)) {
+                        if (in_range)
+                                range_end++;
+                        else {
+                                range_start = range_end = i;
+                                in_range = true;
+                        }
+                } else if (in_range) {
+                        in_range = false;
 
-                if (errno != EINVAL)
+                        if (!GREEDY_REALLOC(str, allocated, len + 2 + 2 * DECIMAL_STR_MAX(unsigned)))
+                                return NULL;
+
+                        if (range_end > range_start || len == 0)
+                                r = sprintf(str + len, len > 0 ? " %d-%d" : "%d-%d", range_start, range_end);
+                        else
+                                r = sprintf(str + len, len > 0 ? " %d" : "%d", range_start);
+                        assert_se(r > 0);
+                        len += r;
+                }
+
+        if (in_range) {
+                if (!GREEDY_REALLOC(str, allocated, len + 2 + 2 * DECIMAL_STR_MAX(int)))
                         return NULL;
 
-                n *= 2;
+                if (range_end > range_start || len == 0)
+                        r = sprintf(str + len, len > 0 ? " %d-%d" : "%d-%d", range_start, range_end);
+                else
+                        r = sprintf(str + len, len > 0 ? " %d" : "%d", range_start);
+                assert_se(r > 0);
+        }
+
+        return TAKE_PTR(str) ?: strdup("");
+}
+
+int cpu_set_realloc(CPUSet *cpu_set, unsigned ncpus) {
+        size_t need;
+
+        assert(cpu_set);
+
+        need = CPU_ALLOC_SIZE(ncpus);
+        if (need > cpu_set->allocated) {
+                cpu_set_t *t;
+
+                t = realloc(cpu_set->set, need);
+                if (!t)
+                        return -ENOMEM;
+
+                memzero((uint8_t*) t + cpu_set->allocated, need - cpu_set->allocated);
+
+                cpu_set->set = t;
+                cpu_set->allocated = need;
         }
+
+        return 0;
 }
 
-int parse_cpu_set_internal(
+static int cpu_set_add(CPUSet *cpu_set, unsigned cpu) {
+        int r;
+
+        if (cpu >= 8192)
+                /* As of kernel 5.1, CONFIG_NR_CPUS can be set to 8192 on PowerPC */
+                return -ERANGE;
+
+        r = cpu_set_realloc(cpu_set, cpu + 1);
+        if (r < 0)
+                return r;
+
+        CPU_SET_S(cpu, cpu_set->allocated, cpu_set->set);
+        return 0;
+}
+
+int cpu_set_add_all(CPUSet *a, const CPUSet *b) {
+        int r;
+
+        /* Do this backwards, so if we fail, we fail before changing anything. */
+        for (unsigned cpu_p1 = b->allocated * 8; cpu_p1 > 0; cpu_p1--)
+                if (CPU_ISSET_S(cpu_p1 - 1, b->allocated, b->set)) {
+                        r = cpu_set_add(a, cpu_p1 - 1);
+                        if (r < 0)
+                                return r;
+                }
+
+        return 0;
+}
+
+int parse_cpu_set_full(
                 const char *rvalue,
-                cpu_set_t **cpu_set,
+                CPUSet *cpu_set,
                 bool warn,
                 const char *unit,
                 const char *filename,
                 unsigned line,
                 const char *lvalue) {
 
-        _cleanup_cpu_free_ cpu_set_t *c = NULL;
+        _cleanup_(cpu_set_reset) CPUSet c = {};
         const char *p = rvalue;
-        unsigned ncpus = 0;
 
-        assert(rvalue);
+        assert(p);
 
         for (;;) {
                 _cleanup_free_ char *word = NULL;
-                unsigned cpu, cpu_lower, cpu_upper;
+                unsigned cpu_lower, cpu_upper;
                 int r;
 
                 r = extract_first_word(&p, &word, WHITESPACE ",", EXTRACT_QUOTES);
@@ -69,31 +155,135 @@ int parse_cpu_set_internal(
                 if (r == 0)
                         break;
 
-                if (!c) {
-                        c = cpu_set_malloc(&ncpus);
-                        if (!c)
-                                return warn ? log_oom() : -ENOMEM;
-                }
-
                 r = parse_range(word, &cpu_lower, &cpu_upper);
                 if (r < 0)
                         return warn ? log_syntax(unit, LOG_ERR, filename, line, r, "Failed to parse CPU affinity '%s'", word) : r;
-                if (cpu_lower >= ncpus || cpu_upper >= ncpus)
-                        return warn ? log_syntax(unit, LOG_ERR, filename, line, EINVAL, "CPU out of range '%s' ncpus is %u", word, ncpus) : -EINVAL;
 
                 if (cpu_lower > cpu_upper) {
                         if (warn)
-                                log_syntax(unit, LOG_WARNING, filename, line, 0, "Range '%s' is invalid, %u > %u, ignoring", word, cpu_lower, cpu_upper);
-                        continue;
+                                log_syntax(unit, LOG_WARNING, filename, line, 0, "Range '%s' is invalid, %u > %u, ignoring.",
+                                           word, cpu_lower, cpu_upper);
+
+                        /* Make sure something is allocated, to distinguish this from the empty case */
+                        r = cpu_set_realloc(&c, 1);
+                        if (r < 0)
+                                return r;
                 }
 
-                for (cpu = cpu_lower; cpu <= cpu_upper; cpu++)
-                        CPU_SET_S(cpu, CPU_ALLOC_SIZE(ncpus), c);
+                for (unsigned cpu_p1 = MIN(cpu_upper, UINT_MAX-1) + 1; cpu_p1 > cpu_lower; cpu_p1--) {
+                        r = cpu_set_add(&c, cpu_p1 - 1);
+                        if (r < 0)
+                                return warn ? log_syntax(unit, LOG_ERR, filename, line, r,
+                                                         "Cannot add CPU %u to set: %m", cpu_p1 - 1) : r;
+                }
         }
 
-        /* On success, sets *cpu_set and returns ncpus for the system. */
-        if (c)
-                *cpu_set = TAKE_PTR(c);
+        /* On success, transfer ownership to the output variable */
+        *cpu_set = c;
+        c = (CPUSet) {};
+
+        return 0;
+}
+
+int parse_cpu_set_extend(
+                const char *rvalue,
+                CPUSet *old,
+                bool warn,
+                const char *unit,
+                const char *filename,
+                unsigned line,
+                const char *lvalue) {
+
+        _cleanup_(cpu_set_reset) CPUSet cpuset = {};
+        int r;
+
+        r = parse_cpu_set_full(rvalue, &cpuset, true, unit, filename, line, lvalue);
+        if (r < 0)
+                return r;
+
+        if (!cpuset.set) {
+                /* An empty assignment resets the CPU list */
+                cpu_set_reset(old);
+                return 0;
+        }
+
+        if (!old->set) {
+                *old = cpuset;
+                cpuset = (CPUSet) {};
+                return 0;
+        }
+
+        return cpu_set_add_all(old, &cpuset);
+}
+
+int cpus_in_affinity_mask(void) {
+        size_t n = 16;
+        int r;
+
+        for (;;) {
+                cpu_set_t *c;
+
+                c = CPU_ALLOC(n);
+                if (!c)
+                        return -ENOMEM;
+
+                if (sched_getaffinity(0, CPU_ALLOC_SIZE(n), c) >= 0) {
+                        int k;
+
+                        k = CPU_COUNT_S(CPU_ALLOC_SIZE(n), c);
+                        CPU_FREE(c);
+
+                        if (k <= 0)
+                                return -EINVAL;
+
+                        return k;
+                }
+
+                r = -errno;
+                CPU_FREE(c);
+
+                if (r != -EINVAL)
+                        return r;
+                if (n > SIZE_MAX/2)
+                        return -ENOMEM;
+                n *= 2;
+        }
+}
+
+int cpu_set_to_dbus(const CPUSet *set, uint8_t **ret, size_t *allocated) {
+        uint8_t *out;
+
+        assert(set);
+        assert(ret);
+
+        out = new0(uint8_t, set->allocated);
+        if (!out)
+                return -ENOMEM;
+
+        for (unsigned cpu = 0; cpu < set->allocated * 8; cpu++)
+                if (CPU_ISSET_S(cpu, set->allocated, set->set))
+                        out[cpu / 8] |= 1u << (cpu % 8);
+
+        *ret = out;
+        *allocated = set->allocated;
+        return 0;
+}
+
+int cpu_set_from_dbus(const uint8_t *bits, size_t size, CPUSet *set) {
+        _cleanup_(cpu_set_reset) CPUSet s = {};
+        int r;
+
+        assert(bits);
+        assert(set);
+
+        for (unsigned cpu = size * 8; cpu > 0; cpu--)
+                if (bits[(cpu - 1) / 8] & (1u << ((cpu - 1) % 8))) {
+                        r = cpu_set_add(&s, cpu - 1);
+                        if (r < 0)
+                                return r;
+                }
 
-        return (int) ncpus;
+        *set = s;
+        s = (CPUSet) {};
+        return 0;
 }
diff --git a/src/shared/cpu-set-util.h b/src/shared/cpu-set-util.h
index 1b6bd35b1c..fd6a15f446 100644
--- a/src/shared/cpu-set-util.h
+++ b/src/shared/cpu-set-util.h
@@ -5,31 +5,46 @@
 
 #include "macro.h"
 
-#ifdef __NCPUBITS
-#define CPU_SIZE_TO_NUM(n) ((n) * __NCPUBITS)
-#else
-#define CPU_SIZE_TO_NUM(n) ((n) * sizeof(cpu_set_t) * 8)
-#endif
-
-DEFINE_TRIVIAL_CLEANUP_FUNC(cpu_set_t*, CPU_FREE);
-#define _cleanup_cpu_free_ _cleanup_(CPU_FREEp)
-
-static inline cpu_set_t* cpu_set_mfree(cpu_set_t *p) {
-        if (p)
-                CPU_FREE(p);
-        return NULL;
+/* This wraps the libc interface with a variable to keep the allocated size. */
+typedef struct CPUSet {
+        cpu_set_t *set;
+        size_t allocated; /* in bytes */
+} CPUSet;
+
+static inline void cpu_set_reset(CPUSet *a) {
+        assert((a->allocated > 0) == !!a->set);
+        if (a->set)
+                CPU_FREE(a->set);
+        *a = (CPUSet) {};
 }
 
-cpu_set_t* cpu_set_malloc(unsigned *ncpus);
-
-int parse_cpu_set_internal(const char *rvalue, cpu_set_t **cpu_set, bool warn, const char *unit, const char *filename, unsigned line, const char *lvalue);
-
-static inline int parse_cpu_set_and_warn(const char *rvalue, cpu_set_t **cpu_set, const char *unit, const char *filename, unsigned line, const char *lvalue) {
-        assert(lvalue);
-
-        return parse_cpu_set_internal(rvalue, cpu_set, true, unit, filename, line, lvalue);
+int cpu_set_add_all(CPUSet *a, const CPUSet *b);
+
+char* cpu_set_to_string(const CPUSet *a);
+char *cpu_set_to_range_string(const CPUSet *a);
+int cpu_set_realloc(CPUSet *cpu_set, unsigned ncpus);
+
+int parse_cpu_set_full(
+                const char *rvalue,
+                CPUSet *cpu_set,
+                bool warn,
+                const char *unit,
+                const char *filename, unsigned line,
+                const char *lvalue);
+int parse_cpu_set_extend(
+                const char *rvalue,
+                CPUSet *old,
+                bool warn,
+                const char *unit,
+                const char *filename,
+                unsigned line,
+                const char *lvalue);
+
+static inline int parse_cpu_set(const char *rvalue, CPUSet *cpu_set){
+        return parse_cpu_set_full(rvalue, cpu_set, false, NULL, NULL, 0, NULL);
 }
 
-static inline int parse_cpu_set(const char *rvalue, cpu_set_t **cpu_set){
-        return parse_cpu_set_internal(rvalue, cpu_set, false, NULL, NULL, 0, NULL);
-}
+int cpu_set_to_dbus(const CPUSet *set, uint8_t **ret, size_t *allocated);
+int cpu_set_from_dbus(const uint8_t *bits, size_t size, CPUSet *set);
+
+int cpus_in_affinity_mask(void);