From 501941aaf001bc069d530fa0edc1daa1cdc042d7 Mon Sep 17 00:00:00 2001
From: Yu Watanabe <watanabe.yu+github@gmail.com>
Date: Thu, 30 Nov 2017 23:16:58 +0900
Subject: [PATCH] core: merge multiple CPUAffinity= settings

---
 src/basic/cpu-set-util.h |  6 ++++++
 src/core/dbus-execute.c  | 40 ++++++++++++++++++++++++++++++----------
 src/core/load-fragment.c | 27 ++++++++++++++++++++-------
 3 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/src/basic/cpu-set-util.h b/src/basic/cpu-set-util.h
index 3ea6cdb..c98149e 100644
--- a/src/basic/cpu-set-util.h
+++ b/src/basic/cpu-set-util.h
@@ -25,6 +25,12 @@
 
 #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)
 
diff --git a/src/core/dbus-execute.c b/src/core/dbus-execute.c
index 197c1c3..196de56 100644
--- a/src/core/dbus-execute.c
+++ b/src/core/dbus-execute.c
@@ -29,6 +29,7 @@
 #include "bus-util.h"
 #include "cap-list.h"
 #include "capability-util.h"
+#include "cpu-set-util.h"
 #include "dbus-execute.h"
 #include "env-util.h"
 #include "errno-list.h"
@@ -1591,29 +1592,29 @@ int bus_exec_context_set_transient_property(
 
                 if (!UNIT_WRITE_FLAGS_NOOP(flags)) {
                         if (n == 0) {
-                                c->cpuset = mfree(c->cpuset);
+                                c->cpuset = cpu_set_mfree(c->cpuset);
+                                c->cpuset_ncpus = 0;
                                 unit_write_settingf(u, flags, name, "%s=", name);
                         } else {
                                 _cleanup_free_ char *str = NULL;
-                                uint8_t *l;
-                                size_t allocated = 0, len = 0, i;
+                                size_t allocated = 0, len = 0, i, ncpus;
 
-                                c->cpuset = (cpu_set_t*) memdup(a, sizeof(cpu_set_t) * n);
-                                if (c->cpuset)
-                                        return -ENOMEM;
+                                ncpus = CPU_SIZE_TO_NUM(n);
 
-                                l = (uint8_t*) a;
-                                for (i = 0; i < n; i++) {
+                                for (i = 0; i < ncpus; i++) {
                                         _cleanup_free_ char *p = NULL;
                                         size_t add;
 
-                                        r = asprintf(&p, "%hhi", l[i]);
+                                        if (!CPU_ISSET_S(i, n, (cpu_set_t*) a))
+                                                continue;
+
+                                        r = asprintf(&p, "%zu", i);
                                         if (r < 0)
                                                 return -ENOMEM;
 
                                         add = strlen(p);
 
-                                        if (GREEDY_REALLOC(str, allocated, len + add + 2))
+                                        if (!GREEDY_REALLOC(str, allocated, len + add + 2))
                                                 return -ENOMEM;
 
                                         strcpy(mempcpy(str + len, p, add), " ");
@@ -1623,6 +1624,25 @@ int bus_exec_context_set_transient_property(
                                 if (len != 0)
                                         str[len - 1] = '\0';
 
+                                if (!c->cpuset || c->cpuset_ncpus < ncpus) {
+                                        cpu_set_t *cpuset;
+
+                                        cpuset = CPU_ALLOC(ncpus);
+                                        if (!cpuset)
+                                                return -ENOMEM;
+
+                                        CPU_ZERO_S(n, cpuset);
+                                        if (c->cpuset) {
+                                                CPU_OR_S(CPU_ALLOC_SIZE(c->cpuset_ncpus), cpuset, c->cpuset, (cpu_set_t*) a);
+                                                CPU_FREE(c->cpuset);
+                                        } else
+                                                CPU_OR_S(n, cpuset, cpuset, (cpu_set_t*) a);
+
+                                        c->cpuset = cpuset;
+                                        c->cpuset_ncpus = ncpus;
+                                } else
+                                        CPU_OR_S(n, c->cpuset, c->cpuset, (cpu_set_t*) a);
+
                                 unit_write_settingf(u, flags, name, "%s=%s", name, str);
                         }
                 }
diff --git a/src/core/load-fragment.c b/src/core/load-fragment.c
index d99176a..d6c6165 100644
--- a/src/core/load-fragment.c
+++ b/src/core/load-fragment.c
@@ -1268,17 +1268,30 @@ int config_parse_exec_cpu_affinity(const char *unit,
         if (ncpus < 0)
                 return ncpus;
 
-        if (c->cpuset)
-                CPU_FREE(c->cpuset);
-
-        if (ncpus == 0)
+        if (ncpus == 0) {
                 /* An empty assignment resets the CPU list */
-                c->cpuset = NULL;
-        else {
+                c->cpuset = cpu_set_mfree(c->cpuset);
+                c->cpuset_ncpus = 0;
+                return 0;
+        }
+
+        if (!c->cpuset) {
+                c->cpuset = cpuset;
+                cpuset = NULL;
+                c->cpuset_ncpus = (unsigned) ncpus;
+                return 0;
+        }
+
+        if (c->cpuset_ncpus < (unsigned) ncpus) {
+                CPU_OR_S(CPU_ALLOC_SIZE(c->cpuset_ncpus), cpuset, c->cpuset, cpuset);
+                CPU_FREE(c->cpuset);
                 c->cpuset = cpuset;
                 cpuset = NULL;
+                c->cpuset_ncpus = (unsigned) ncpus;
+                return 0;
         }
-        c->cpuset_ncpus = ncpus;
+
+        CPU_OR_S(CPU_ALLOC_SIZE((unsigned) ncpus), c->cpuset, c->cpuset, cpuset);
 
         return 0;
 }
-- 
2.7.4