Dinamically size the cpu_set_t objects

To cope with mismatches on how glibc and the kernel defined the max
number of CPUs.

Signed-off-by: Arnaldo Carvalho de Melo <acme@redhat.com>

1 files changed, 100 insertions, 17 deletions

diff --git a/python-schedutils/schedutils.c b/python-schedutils/schedutils.c
index 82a1a3e..40072b9 100644
--- a/python-schedutils/schedutils.c
+++ b/python-schedutils/schedutils.c
@@ -1,6 +1,7 @@
 #include <Python.h>
 #include <sched.h>
 #include <errno.h>
+#include <syscall.h>
 
 #ifndef __unused
 #define __unused __attribute__ ((unused))
@@ -10,59 +11,137 @@
 #define SCHED_RESET_ON_FORK 0x40000000
 #endif
 
+/*
+ * The following bitmask declarations, bitmask_*() routines, and associated
+ * _setbit() and _getbit() routines are:
+ * Copyright (c) 2004 Silicon Graphics, Inc. (SGI) All rights reserved.
+ * SGI publishes it under the terms of the GNU General Public License, v2,
+ * as published by the Free Software Foundation.
+ */
+#define howmany(x,y) (((x)+((y)-1))/(y))
+#define bitsperlong (8 * sizeof(unsigned long))
+#define longsperbits(n) howmany(n, bitsperlong)
+#define bytesperbits(x) ((x+7)/8)
+
+/*
+ * Number of bits in a CPU bitmask on current system
+ */
+static int __get_max_number_of_cpus(void)
+{
+	int n;
+	int cpus = 2048;
+
+	for (;;) {
+		unsigned long buffer[longsperbits(cpus)];
+		memset(buffer, 0, sizeof(buffer));
+		/* the library version does not return size of cpumask_t */
+		n = syscall(SYS_sched_getaffinity, 0, bytesperbits(cpus), &buffer);
+		if (n < 0 && errno == EINVAL && cpus < 1024 * 1024) {
+			cpus *= 2;
+			continue;
+		}
+		return n * 8;
+	}
+	return -1;
+}
+
+static PyObject *get_max_number_of_cpus(PyObject *self __unused, PyObject *args __unused)
+{
+	int ret = __get_max_number_of_cpus();
+
+	if (ret < 0) {
+		PyErr_SetFromErrno(PyExc_OSError);
+		return NULL;
+	}
+
+	return Py_BuildValue("i", ret);
+}
+
 static PyObject *get_affinity(PyObject *self __unused, PyObject *args)
 {
-	PyObject *list;
-	cpu_set_t cpus;
+	PyObject *list = NULL;
+	cpu_set_t *cpus;
 	int pid, cpu;
+	size_t cpusetsize;
+	int max_cpus;
 
 	if (!PyArg_ParseTuple(args, "i", &pid))
-		return NULL;
+		goto out_error;
 
-	CPU_ZERO(&cpus);
+	max_cpus = __get_max_number_of_cpus();
+	if (max_cpus < 0)
+		goto out_error;
 
-	if (sched_getaffinity(pid, sizeof(cpus), &cpus) < 0) {
-		PyErr_SetFromErrno(PyExc_OSError);
-		return NULL;
-	}
+	cpus = CPU_ALLOC(max_cpus);
+	if (cpus == NULL)
+		goto out_error;
+
+	cpusetsize = CPU_ALLOC_SIZE(max_cpus);
+	CPU_ZERO_S(cpusetsize, cpus);
+
+	if (sched_getaffinity(pid, cpusetsize, cpus) < 0)
+		goto out_free;
 
 	list = PyList_New(0);
 	for (cpu = 0; cpu < CPU_SETSIZE; ++cpu)
-		if (CPU_ISSET(cpu, &cpus))
+		if (CPU_ISSET_S(cpu, cpusetsize, cpus))
 			PyList_Append(list, Py_BuildValue("i", cpu));
 
+	CPU_FREE(cpus);
+out:
 	return list;
+out_free:
+	CPU_FREE(cpus);
+out_error:
+	PyErr_SetFromErrno(PyExc_OSError);
+	goto out;
 }
 
 static PyObject *set_affinity(PyObject *self __unused, PyObject *args)
 {
-	int pid, nr_elements, i;
-	cpu_set_t cpus;
+	int pid, nr_elements, i, max_cpus;
+	cpu_set_t *cpus;
 	PyObject *list;
+	size_t cpusetsize;
 
 	if (!PyArg_ParseTuple(args, "iO", &pid, &list))
-		return NULL;
+		goto out_error;
+
+	max_cpus = __get_max_number_of_cpus();
+	if (max_cpus < 0)
+		goto out_error;
+
+	cpus = CPU_ALLOC(max_cpus);
+	if (cpus == NULL)
+		goto out_error;
 
-	CPU_ZERO(&cpus);
+	cpusetsize = CPU_ALLOC_SIZE(max_cpus);
+	CPU_ZERO_S(cpusetsize, cpus);
 
 	nr_elements = PyList_Size(list);
 	for (i = 0; i < nr_elements; ++i) {
 		int cpu = PyInt_AsLong(PyList_GetItem(list, i));
 
-		if (cpu >= CPU_SETSIZE) {
+		if (cpu >= max_cpus) {
 			PyErr_SetString(PyExc_OSError, "Invalid CPU");
-			return NULL;
+			goto out_free;
 		}
-		CPU_SET(cpu, &cpus);
+		CPU_SET_S(cpu, cpusetsize, cpus);
 	}
 
 	if (sched_setaffinity(pid, sizeof(cpus), &cpus) < 0) {
 		PyErr_SetFromErrno(PyExc_OSError);
 		return NULL;
 	}
-
+	CPU_FREE(cpus);
+out:
 	Py_INCREF(Py_None);
 	return Py_None;
+out_free:
+	CPU_FREE(cpus);
+out_error:
+	PyErr_SetFromErrno(PyExc_OSError);
+	goto out;
 }
 
 static PyObject *get_scheduler(PyObject *self __unused, PyObject *args)
@@ -246,6 +325,10 @@ static struct PyMethodDef PySchedutilsModuleMethods[] = {
 		.ml_meth = (PyCFunction)get_priority_max,
 		.ml_flags = METH_VARARGS,
 	},
+	{
+		.ml_name = "get_max_number_of_cpus",
+		.ml_meth = (PyCFunction)get_max_number_of_cpus,
+	},
 	{	.ml_name = NULL, },
 };