When setting CPU affinity, take into consideration the current CPU

affinity set. Use only CPU numbers, which exist in the current affinity
set. Maybe fixes https://github.com/google/nsjail/issues/200

cmdline.cc

@@ -484,7 +484,6 @@ std::unique_ptr<nsjconf_t> parseArgs(int argc, char* argv[]) {
 	nsjconf->forward_signals = false;
 	nsjconf->orig_uid = getuid();
 	nsjconf->orig_euid = geteuid();
-	nsjconf->num_cpus = sysconf(_SC_NPROCESSORS_ONLN);
 	nsjconf->seccomp_fprog.filter = NULL;
 	nsjconf->seccomp_fprog.len = 0;
 	nsjconf->seccomp_log = false;

cpu.cc

@@ -27,66 +27,83 @@
 #include <string.h>
 #include <unistd.h>
 
+#include <memory>
+
 #include "logs.h"
 #include "util.h"
 
 namespace cpu {
 
-static void setRandomCpu(cpu_set_t* mask, size_t mask_size, size_t cpu_num) {
-	if ((size_t)CPU_COUNT_S(mask_size, mask) >= cpu_num) {
-		LOG_F(
-		    "Number of CPUs in the mask '%d' is bigger than number of available CPUs '%zu'",
-		    CPU_COUNT(mask), cpu_num);
+static size_t getNthCpu(cpu_set_t* mask, size_t mask_size, size_t n, size_t cpus_left) {
+	for (size_t i = 0, j = 0; i < CPU_SETSIZE; i++) {
+		if (CPU_ISSET_S(i, mask_size, mask)) {
+			if (j == n) {
+				return i;
+			}
+			j++;
+		}
+	}
+	LOG_F("No CPU #%zu found, yet there should be %zu left", n,
+	    (size_t)CPU_COUNT_S(mask_size, mask));
+	return 0;
+}
+
+static void setRandomCpu(cpu_set_t* original_mask, cpu_set_t* new_mask, size_t mask_size) {
+	size_t cpus_left = CPU_COUNT_S(mask_size, original_mask);
+	if (cpus_left == 0) {
+		LOG_F("There are no more CPUs left to use");
 	}
 
-	for (;;) {
-		uint64_t n = util::rnd64() % cpu_num;
-		if (!CPU_ISSET_S(n, mask_size, mask)) {
-			LOG_D("Setting allowed CPU#:%" PRIu64 " of [0-%zu]", n, cpu_num - 1);
-			CPU_SET_S(n, mask_size, mask);
-			break;
-		}
-	}
+	size_t n = getNthCpu(original_mask, mask_size, util::rnd64() % cpus_left, cpus_left);
+	LOG_D("Setting allowed CPU#:%" PRIu64, n);
+	CPU_SET_S(n, mask_size, new_mask);
+	CPU_CLR_S(n, mask_size, original_mask);
 }
 
 bool initCpu(nsjconf_t* nsjconf) {
-	if (nsjconf->num_cpus < 0) {
-		PLOG_W("sysconf(_SC_NPROCESSORS_ONLN) returned %ld", nsjconf->num_cpus);
-		return false;
-	}
-	if (nsjconf->max_cpus > (size_t)nsjconf->num_cpus) {
-		LOG_W("Requested number of CPUs:%zu is bigger than CPUs online:%ld",
-		    nsjconf->max_cpus, nsjconf->num_cpus);
-		return true;
-	}
-	if (nsjconf->max_cpus == (size_t)nsjconf->num_cpus) {
-		LOG_D("All CPUs requested (%zu of %ld)", nsjconf->max_cpus, nsjconf->num_cpus);
-		return true;
-	}
 	if (nsjconf->max_cpus == 0) {
 		LOG_D("No max_cpus limit set");
 		return true;
 	}
 
-	cpu_set_t* mask = CPU_ALLOC(nsjconf->num_cpus);
-	if (mask == NULL) {
-		PLOG_W("Failure allocating cpu_set_t for %ld CPUs", nsjconf->num_cpus);
+	size_t mask_size = CPU_ALLOC_SIZE(CPU_SETSIZE);
+	std::unique_ptr<cpu_set_t> original_mask(CPU_ALLOC(CPU_SETSIZE));
+	if (original_mask.get() == NULL) {
+		PLOG_W("Failure allocating cpu_set_t for %d CPUs", CPU_SETSIZE);
 		return false;
 	}
+	if (sched_getaffinity(0, mask_size, original_mask.get()) == -1) {
+		PLOG_W("sched_getaffinity(0, mask_size=%zu)", mask_size);
+		return false;
+	}
+	size_t available_cpus = CPU_COUNT_S(mask_size, original_mask.get());
 
-	size_t mask_size = CPU_ALLOC_SIZE(nsjconf->num_cpus);
-	CPU_ZERO_S(mask_size, mask);
+	if (nsjconf->max_cpus > available_cpus) {
+		LOG_W(
+		    "Number of requested CPUs is bigger than number of available CPUs (%zu > %zu)",
+		    nsjconf->max_cpus, available_cpus);
+		return true;
+	}
+	if (nsjconf->max_cpus == available_cpus) {
+		LOG_D("All CPUs requested (%zu of %zu)", nsjconf->max_cpus, available_cpus);
+		return true;
+	}
+
+	std::unique_ptr<cpu_set_t> new_mask(CPU_ALLOC(CPU_SETSIZE));
+	if (new_mask.get() == NULL) {
+		PLOG_W("Failure allocating cpu_set_t for %d CPUs", CPU_SETSIZE);
+		return false;
+	}
+	CPU_ZERO_S(mask_size, new_mask.get());
 
 	for (size_t i = 0; i < nsjconf->max_cpus; i++) {
-		setRandomCpu(mask, mask_size, nsjconf->num_cpus);
+		setRandomCpu(original_mask.get(), new_mask.get(), mask_size);
 	}
 
-	if (sched_setaffinity(0, mask_size, mask) == -1) {
+	if (sched_setaffinity(0, mask_size, new_mask.get()) == -1) {
 		PLOG_W("sched_setaffinity(max_cpus=%zu) failed", nsjconf->max_cpus);
-		CPU_FREE(mask);
 		return false;
 	}
-	CPU_FREE(mask);
 
 	return true;
 }

nsjail.h

@@ -168,7 +168,6 @@ struct nsjconf_t {
 	struct sock_fprog seccomp_fprog;
 	bool seccomp_log;
 	int nice_level;
-	long num_cpus;
 	uid_t orig_uid;
 	uid_t orig_euid;
 	std::vector<mount_t> mountpts;