process-util: add new helper cpus_in_affinity_mask()

Let's add a new helper for determining how many CPUs are configured for
us

2 files changed, 36 insertions, 0 deletions

diff --git a/src/basic/process-util.c b/src/basic/process-util.c
index c41a2aa5c9..1f228be5d3 100644
--- a/src/basic/process-util.c
+++ b/src/basic/process-util.c
@@ -1535,6 +1535,40 @@ int set_oom_score_adjust(int value) {
                                  WRITE_STRING_FILE_VERIFY_ON_FAILURE|WRITE_STRING_FILE_DISABLE_BUFFER);
 }
 
+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;
+        }
+}
+
 static const char *const ioprio_class_table[] = {
         [IOPRIO_CLASS_NONE] = "none",
         [IOPRIO_CLASS_RT] = "realtime",
diff --git a/src/basic/process-util.h b/src/basic/process-util.h
index 9950723996..f6eb82d6b6 100644
--- a/src/basic/process-util.h
+++ b/src/basic/process-util.h
@@ -193,3 +193,5 @@ assert_cc(TASKS_MAX <= (unsigned long) PID_T_MAX)
                 (pid) = 0;                      \
                 _pid_;                          \
         })
+
+int cpus_in_affinity_mask(void);