Add optional total/proc/thread arguments to %{getncpus} macro

"total" equals calling with no arguments, "proc" and "thread" consider
further constraints, what is implemented here is heuristics based on
available physical memory and address-space and %_smp_tasksize_proc /
%_smp_tasksize_thread tunables.

Change the previous %getncpus related tests to use %getconfdir instead,
they are testing unexpected arguments behavior for this type of macro,
not %getncpus itself. Add a test for the actual functionality: if nproc
is available, test that our total matches with that, and that defining
tasksize to total memory only allocates one thread. Optimally we'd
test separately for 32bit address space limitations but that gets tough
when we have no idea where this will be executed.

4 files changed, 118 insertions, 9 deletions

diff --git a/docs/manual/macros.md b/docs/manual/macros.md
index 0a892af09..6a2f77e3a 100644
--- a/docs/manual/macros.md
+++ b/docs/manual/macros.md
@@ -68,7 +68,10 @@ to perform useful operations. The current list is
 	%trace		toggle print of debugging information before/after
 			expansion
 	%dump		print the active (i.e. non-covered) macro table
-	%getncpus	return the number of CPUs
+	%getncpus	expand to the number of available CPUs
+	%{getncpus:<total|proc|thread>} expand to the number of available CPUs,
+                "proc" and "thread" additionally accounting for available
+                memory (eg address space limitations for threads)
 	%getconfdir	expand to rpm "home" directory (typically /usr/lib/rpm)
 	%dnl		discard to next line (without expanding)
 	%verbose	expand to 1 if rpm is in verbose mode, 0 if not
diff --git a/macros.in b/macros.in
index 064d9248c..796278b97 100644
--- a/macros.in
+++ b/macros.in
@@ -727,6 +727,11 @@ Supplements:   (%{name} = %{version}-%{release} and langpacks-%{1})\
 
 %_smp_build_nthreads %{_smp_build_ncpus}
 
+# Assumed task size of processes and threads in megabytes.
+# Used to limit the amount of parallelism based on available memory.
+%_smp_tasksize_proc 512
+%_smp_tasksize_thread %{_smp_tasksize_proc}
+
 #==============================================================================
 # ---- Scriptlet template templates.
 #	Global defaults used for building scriptlet templates.
diff --git a/rpmio/macro.c b/rpmio/macro.c
index d86b84608..5dc926213 100644
--- a/rpmio/macro.c
+++ b/rpmio/macro.c
@@ -9,6 +9,9 @@
 #ifdef HAVE_SCHED_GETAFFINITY
 #include <sched.h>
 #endif
+#if defined(__linux__)
+#include <sys/personality.h>
+#endif
 
 #if !defined(isblank)
 #define	isblank(_c)	((_c) == ' ' || (_c) == '\t')
@@ -1174,6 +1177,89 @@ static void doShescape(MacroBuf mb, rpmMacroEntry me, ARGV_t argv, size_t *parse
     mbAppend(mb, '\'');
 }
 
+static unsigned long getmem_total(void)
+{
+    unsigned long mem = 0;
+    long int pagesize = sysconf(_SC_PAGESIZE);
+    long int pages = sysconf(_SC_PHYS_PAGES);
+
+    if (pagesize < 0)
+	pagesize = 4096;
+    if (pages > 0)
+	mem = pages * pagesize;
+
+    return mem;
+}
+
+static unsigned long getmem_proc(int thread)
+{
+    unsigned long mem = getmem_total();
+    /*
+     * Conservative estimates for thread use on 32bit systems where address
+     * space is an issue: 2GB for bare metal, 3GB for a 32bit process
+     * on a 64bit system.
+     */
+    if (thread) {
+	unsigned long vmem = mem;
+#if __WORDSIZE == 32
+	vmem = UINT32_MAX / 2;
+#else
+#if defined(__linux__)
+	if ((personality(0xffffffff) & PER_MASK) == PER_LINUX32)
+	    vmem = (UINT32_MAX / 4) * 3;
+#endif
+#endif
+	if (vmem < mem)
+	    mem = vmem;
+    }
+    /* Fixup to get nice even numbers */
+    mem = mem / (1024*1024) + 1;
+
+    return mem;
+}
+
+static void doGetncpus(MacroBuf mb, rpmMacroEntry me, ARGV_t argv, size_t *parsed)
+{
+    const char *sizemacro = NULL;
+    const char *arg = argv[1] ? argv[1] : "total";
+    char buf[32];
+    unsigned int ncpus = getncpus();
+    unsigned long mem = 0;
+
+    if (rstreq(arg, "total")) {
+	/* nothing */
+    } else if (rstreq(arg, "proc")) {
+	mem = getmem_proc(0);
+	sizemacro = "%{?_smp_tasksize_proc}";
+    } else if (rstreq(arg, "thread")) {
+	mem = getmem_proc(1);
+	sizemacro = "%{?_smp_tasksize_thread}";
+    } else {
+	mbErr(mb, 1, _("invalid argument: %s\n"), arg);
+	return;
+    }
+
+    if (sizemacro) {
+	unsigned int mcpus;
+	unsigned long tasksize = rpmExpandNumeric(sizemacro);
+
+	if (tasksize == 0)
+	    tasksize = 512;
+
+	if (mem == 0) {
+	    mbErr(mb, 1, _("failed to get available memory for %s\n"), arg);
+	    return;
+	}
+
+	mcpus = mem / tasksize;
+	if (mcpus < ncpus)
+	    ncpus = mcpus;
+    }
+
+    sprintf(buf, "%u", ncpus);
+    mbAppendStr(mb, buf);
+}
+
 static void doFoo(MacroBuf mb, rpmMacroEntry me, ARGV_t argv, size_t *parsed)
 {
     char *buf = NULL;
@@ -1227,9 +1313,6 @@ static void doFoo(MacroBuf mb, rpmMacroEntry me, ARGV_t argv, size_t *parsed)
 	b = getenv(argv[1]);
     } else if (rstreq("getconfdir", me->name)) {
 	b = (char *)rpmConfigDir();
-    } else if (rstreq("getncpus", me->name)) {
-	b = buf = xmalloc(MACROBUFSIZ);
-	sprintf(buf, "%u", getncpus());
     } else if (rstreq("exists", me->name)) {
 	b = (access(argv[1], F_OK) == 0) ? "1" : "0";
     }
@@ -1276,7 +1359,7 @@ static struct builtins_s {
     { "expr",		doFoo,		1,	0 },
     { "getconfdir",	doFoo,		0,	0 },
     { "getenv",		doFoo,		1,	0 },
-    { "getncpus",	doFoo,		0,	0 },
+    { "getncpus",	doGetncpus,	-1,	0 },
     { "global",		doGlobal,	1,	ME_PARSE },
     { "gsub",		doString,	1,	0 },
     { "len",		doString,	1,	0 },
diff --git a/tests/rpmmacro.at b/tests/rpmmacro.at
index 03da6db40..9fe9cc259 100644
--- a/tests/rpmmacro.at
+++ b/tests/rpmmacro.at
@@ -263,8 +263,26 @@ runroot_other ${RPM_CONFIGDIR}/rpmuncompress "/tmp/some%%ath"
 [0],
 [xxxxxxxxxxxxxxxxxxxxxxxxx
 ])
+AT_CLEANUP
 
+AT_SETUP([getncpus macro])
+AT_KEYWORDS([macros])
+# skip if nproc not available
+AT_SKIP_IF([test -z "$(nproc 2>/dev/null)"])
+AT_CHECK([
+mem=$(expr $(getconf PAGESIZE) \* $(getconf _PHYS_PAGES) / 1024 / 1024)
+expr $(runroot rpm --eval "%{getncpus}") = $(nproc)
+expr $(runroot rpm --define "_smp_tasksize_thread ${mem}" --eval "%{getncpus:thread}") = 1
+expr $(runroot rpm --define "_smp_tasksize_proc ${mem}" --eval "%{getncpus:proc}") = 1
+],
+[ignore],
+[1
+1
+1
+],
+[])
 AT_CLEANUP
+
 AT_SETUP([basename macro])
 AT_KEYWORDS([macros])
 AT_CHECK([
@@ -338,8 +356,8 @@ runroot rpm --eval "%dirname dir"
 runroot rpm --define '%xxx /hello/%%%%/world' --eval '%{dirname:%xxx}'
 runroot rpm --eval "%{uncompress}"
 runroot rpm --eval "%{uncompress:}"
-runroot rpm --eval "%{getncpus:}"
-runroot rpm --eval "%{getncpus:5}"
+runroot rpm --eval "%{getconfdir:}"
+runroot rpm --eval "%{getconfdir:5}"
 runroot rpm --eval "%{define:}"
 runroot rpm --eval "%{define:foo}"
 runroot rpm --eval "%{define:foo bar}%{foo}"
@@ -368,8 +386,8 @@ bar baz\baz
 [error: %dirname: argument expected
 error: %dirname: argument expected
 error: %uncompress: argument expected
-error: %getncpus: unexpected argument
-error: %getncpus: unexpected argument
+error: %getconfdir: unexpected argument
+error: %getconfdir: unexpected argument
 error: Macro % has illegal name (%define)
 error: Macro %foo has empty body
 error: Macro %foo has empty body