Merge pull request #527 from ddaeschler/master

Fix get_affinity for linux
giampaolo · Sep 15, 2014 · 1e5dc37 · 1e5dc37
2 parents 45ceea9 + 4be391d
commit 1e5dc37
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 9 deletions.
diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
@@ -20,7 +20,7 @@
 from psutil import _common
 from psutil import _psposix
 from psutil._common import (isfile_strict, usage_percent, deprecated)
-from psutil._compat import PY3, xrange, namedtuple, wraps, b, defaultdict
+from psutil._compat import PY3, namedtuple, wraps, b, defaultdict
 import _psutil_linux as cext
 import _psutil_posix
 
@@ -1058,9 +1058,7 @@ def nice_set(self, value):
 
     @wrap_exceptions
     def cpu_affinity_get(self):
-        from_bitmask = lambda x: [i for i in xrange(64) if (1 << i) & x]
-        bitmask = cext.proc_cpu_affinity_get(self.pid)
-        return from_bitmask(bitmask)
+        return cext.proc_cpu_affinity_get(self.pid)
 
     @wrap_exceptions
     def cpu_affinity_set(self, cpus):

diff --git a/psutil/_psutil_linux.c b/psutil/_psutil_linux.c
@@ -22,6 +22,8 @@
 
 #include "_psutil_linux.h"
 
+/* The minimum number of CPUs allocated in a cpu_set_t */
+static const int NCPUS_START = sizeof(unsigned long) * CHAR_BIT;
 
 // Linux >= 2.6.13
 #define PSUTIL_HAVE_IOPRIO defined(__NR_ioprio_get) && defined(__NR_ioprio_set)
@@ -255,17 +257,65 @@ psutil_linux_sysinfo(PyObject *self, PyObject *args)
 static PyObject *
 psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args)
 {
-    unsigned long mask;
-    unsigned int len = sizeof(mask);
+    int cpu, ncpus, count, cpucount_s;
     long pid;
+    size_t setsize;
+    cpu_set_t *mask = NULL;
+    PyObject *res = NULL;
 
     if (!PyArg_ParseTuple(args, "i", &pid)) {
         return NULL;
     }
-    if (sched_getaffinity(pid, len, (cpu_set_t *)&mask) < 0) {
-        return PyErr_SetFromErrno(PyExc_OSError);
+
+    ncpus = NCPUS_START;
+    while (1) {
+        setsize = CPU_ALLOC_SIZE(ncpus);
+        mask = CPU_ALLOC(ncpus);
+        if (mask == NULL)
+            return PyErr_NoMemory();
+        if (sched_getaffinity(pid, setsize, mask) == 0)
+            break;
+        CPU_FREE(mask);
+        if (errno != EINVAL)
+            return PyErr_SetFromErrno(PyExc_OSError);
+        if (ncpus > INT_MAX / 2) {
+            PyErr_SetString(PyExc_OverflowError, "could not allocate "
+                            "a large enough CPU set");
+            return NULL;
+        }
+        ncpus = ncpus * 2;
     }
-    return Py_BuildValue("l", mask);
+
+    res = PyList_New(0);
+    if (res == NULL)
+        goto error;
+
+    cpucount_s = CPU_COUNT_S(setsize, mask);
+    for (cpu = 0, count = cpucount_s; count; cpu++) {
+        if (CPU_ISSET_S(cpu, setsize, mask)) {
+#if PY_MAJOR_VERSION >= 3
+            PyObject *cpu_num = PyLong_FromLong(cpu);
+#else
+            PyObject *cpu_num = PyInt_FromLong(cpu);
+#endif
+            --count;
+            if (cpu_num == NULL)
+                goto error;
+            if (PyList_Append(res, cpu_num)) {
+                Py_DECREF(cpu_num);
+                goto error;
+            }
+            Py_DECREF(cpu_num);
+        }
+    }
+    CPU_FREE(mask);
+    return res;
+
+error:
+    if (mask)
+        CPU_FREE(mask);
+    Py_XDECREF(res);
+    return NULL;
 }