From e3bd69c7df12cb40894b5b0d30e6373fba687c3e Mon Sep 17 00:00:00 2001
From: Panu Matilainen <pmatilai@redhat.com>
Date: Tue, 15 Jan 2019 11:16:04 +0200
Subject: [PATCH] Add a built-in macro for fetching number of CPUs, affinity
 aware and all

There's an increasing number of placing wanting to know the number of
CPU's for parallel processing, and increasingly these things are running
in containers and such where the total host CPU count is not meaningful.

Merged-by: Vitaly Chikunov <vt@altlinux.org>
---
 configure.ac  |  1 +
 rpmio/macro.c | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/configure.ac b/configure.ac
index 9621dad..ed7edaa 100644
--- a/configure.ac
+++ b/configure.ac
@@ -542,6 +542,7 @@ dnl XXX AC_CHECK_FUNCS(gethostname mkdir mkfifo rmdir select uname)
 AC_CHECK_FUNCS(basename getcwd getwd inet_aton mtrace putenv realpath setenv)
 AC_CHECK_FUNCS(stpcpy stpncpy strcspn)
 AC_CHECK_FUNCS(strdup strerror strtol strtoul strspn strstr)
+AC_CHECK_FUNCS(sched_getaffinity, [], [], [#include <sched.h>])
 
 AC_CHECK_FUNCS(regcomp)
 
diff --git a/rpmio/macro.c b/rpmio/macro.c
index 3747c1a..cb8ed9f 100644
--- a/rpmio/macro.c
+++ b/rpmio/macro.c
@@ -9,6 +9,9 @@
 
 #include "system.h"
 #include <stdarg.h>
+#if HAVE_SCHED_GETAFFINITY
+#include <sched.h>
+#endif
 
 #if !defined(isblank)
 #define	isblank(_c)	((_c) == ' ' || (_c) == '\t')
@@ -542,6 +545,23 @@ doShellEscape(MacroBuf mb, const char * cmd, size_t clen)
     return 0;
 }
 
+static unsigned int getncpus(void)
+{
+    unsigned int ncpus = 0;
+#if HAVE_SCHED_GETAFFINITY
+    cpu_set_t set;
+    if (sched_getaffinity (0, sizeof(set), &set) == 0)
+	ncpus = CPU_COUNT(&set);
+#endif
+    /* Fallback to sysconf() if the above isn't supported or didn't work */
+    if (ncpus < 1)
+	ncpus = sysconf(_SC_NPROCESSORS_ONLN);
+    /* If all else fails, there's always the one we're running on... */
+    if (ncpus < 1)
+	ncpus = 1;
+    return ncpus;
+}
+
 /**
  * Parse (and execute) new macro definition.
  * @param mb		macro expansion state
@@ -1057,6 +1077,9 @@ doFoo(MacroBuf mb, int negate, const char * f, size_t fn,
 	    break;
 	}
 	b = be;
+    } else if (STREQ("getncpus", f, fn)) {
+	sprintf(buf, "%u", getncpus());
+	b = buf;
     } else if (STREQ("S", f, fn)) {
 	for (b = buf; (c = *b) && xisdigit(c);)
 	    b++;
@@ -1365,6 +1388,7 @@ expandMacro(MacroBuf mb)
 	    STREQ("getenv", f, fn) ||
 	    STREQ("_tmpdir", f, fn) ||
 	    STREQ("homedir", f, fn) ||
+	    STREQ("getncpus", f, fn) ||
 	    STREQ("S", f, fn) ||
 	    STREQ("P", f, fn) ||
 	    STREQ("F", f, fn)) {