fix #1126: cpu_affinity() segfaults on CentOS 5

Remove cpu_affinity() support for CentOS 5 (it's 8 years old anyway); remove the dual implementation. Recent manylinux versions should use CentOS 6.
giampaolo · Oct 10, 2019 · 79fa262 · 79fa262
1 parent 63fcc0e
commit 79fa262
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 78 deletions.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -12,6 +12,8 @@ XXXX-XX-XX
 
 **Bug fixes**
 
+- 1126_: [Linux] cpu_affinity() segfaults on CentOS 5 / manylinux.
+  cpu_affinity() support for CentOS 5 was removed.
 - 1528_: [AIX] compilation error on AIX 7.2 due to 32 vs 64 bit differences.
   (patch by Arnon Yaari)
 - 1535_: 'type' and 'family' fields returned by net_connections() are not

diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
@@ -73,6 +73,7 @@
 HAS_SMAPS = os.path.exists('/proc/%s/smaps' % os.getpid())
 HAS_PRLIMIT = hasattr(cext, "linux_prlimit")
 HAS_PROC_IO_PRIORITY = hasattr(cext, "proc_ioprio_get")
+HAS_CPU_AFFINITY = hasattr(cext, "proc_cpu_affinity_get")
 _DEFAULT = object()
 
 # RLIMIT_* constants, not guaranteed to be present on all kernels
@@ -1925,38 +1926,41 @@ def nice_get(self):
     def nice_set(self, value):
         return cext_posix.setpriority(self.pid, value)
 
-    @wrap_exceptions
-    def cpu_affinity_get(self):
-        return cext.proc_cpu_affinity_get(self.pid)
+    # starting from CentOS 6.
+    if HAS_CPU_AFFINITY:
 
-    def _get_eligible_cpus(
-            self, _re=re.compile(br"Cpus_allowed_list:\t(\d+)-(\d+)")):
-        # See: https://github.com/giampaolo/psutil/issues/956
-        data = self._read_status_file()
-        match = _re.findall(data)
-        if match:
-            return list(range(int(match[0][0]), int(match[0][1]) + 1))
-        else:
-            return list(range(len(per_cpu_times())))
+        @wrap_exceptions
+        def cpu_affinity_get(self):
+            return cext.proc_cpu_affinity_get(self.pid)
+
+        def _get_eligible_cpus(
+                self, _re=re.compile(br"Cpus_allowed_list:\t(\d+)-(\d+)")):
+            # See: https://github.com/giampaolo/psutil/issues/956
+            data = self._read_status_file()
+            match = _re.findall(data)
+            if match:
+                return list(range(int(match[0][0]), int(match[0][1]) + 1))
+            else:
+                return list(range(len(per_cpu_times())))
 
-    @wrap_exceptions
-    def cpu_affinity_set(self, cpus):
-        try:
-            cext.proc_cpu_affinity_set(self.pid, cpus)
-        except (OSError, ValueError) as err:
-            if isinstance(err, ValueError) or err.errno == errno.EINVAL:
-                eligible_cpus = self._get_eligible_cpus()
-                all_cpus = tuple(range(len(per_cpu_times())))
-                for cpu in cpus:
-                    if cpu not in all_cpus:
-                        raise ValueError(
-                            "invalid CPU number %r; choose between %s" % (
-                                cpu, eligible_cpus))
-                    if cpu not in eligible_cpus:
-                        raise ValueError(
-                            "CPU number %r is not eligible; choose "
-                            "between %s" % (cpu, eligible_cpus))
-            raise
+        @wrap_exceptions
+        def cpu_affinity_set(self, cpus):
+            try:
+                cext.proc_cpu_affinity_set(self.pid, cpus)
+            except (OSError, ValueError) as err:
+                if isinstance(err, ValueError) or err.errno == errno.EINVAL:
+                    eligible_cpus = self._get_eligible_cpus()
+                    all_cpus = tuple(range(len(per_cpu_times())))
+                    for cpu in cpus:
+                        if cpu not in all_cpus:
+                            raise ValueError(
+                                "invalid CPU number %r; choose between %s" % (
+                                    cpu, eligible_cpus))
+                        if cpu not in eligible_cpus:
+                            raise ValueError(
+                                "CPU number %r is not eligible; choose "
+                                "between %s" % (cpu, eligible_cpus))
+                raise
 
     # only starting from kernel 2.6.13
     if HAS_PROC_IO_PRIORITY:

diff --git a/psutil/_psutil_linux.c b/psutil/_psutil_linux.c
@@ -54,6 +54,11 @@ static const int NCPUS_START = sizeof(unsigned long) * CHAR_BIT;
     #include <sys/resource.h>
 #endif
 
+// Should exist starting from CentOS 6 (year 2011).
+#ifdef CPU_ALLOC
+    #define PSUTIL_HAVE_CPU_AFFINITY
+#endif
+
 #include "_psutil_common.h"
 #include "_psutil_posix.h"
 
@@ -279,11 +284,8 @@ psutil_linux_sysinfo(PyObject *self, PyObject *args) {
 
 /*
  * Return process CPU affinity as a Python list
- * The dual implementation exists because of:
- * https://github.com/giampaolo/psutil/issues/536
  */
-
-#ifdef CPU_ALLOC
+#ifdef PSUTIL_HAVE_CPU_AFFINITY
 
 static PyObject *
 psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args) {
@@ -347,50 +349,8 @@ psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args) {
     Py_XDECREF(py_list);
     return NULL;
 }
-#else
 
 
-/*
- * Alternative implementation in case CPU_ALLOC is not defined.
- */
-static PyObject *
-psutil_proc_cpu_affinity_get(PyObject *self, PyObject *args) {
-    cpu_set_t cpuset;
-    unsigned int len = sizeof(cpu_set_t);
-    long pid;
-    int i;
-    PyObject* py_retlist = NULL;
-    PyObject *py_cpu_num = NULL;
-
-    if (!PyArg_ParseTuple(args, "l", &pid))
-        return NULL;
-	CPU_ZERO(&cpuset);
-    if (sched_getaffinity(pid, len, &cpuset) < 0)
-        return PyErr_SetFromErrno(PyExc_OSError);
-
-    py_retlist = PyList_New(0);
-    if (py_retlist == NULL)
-        goto error;
-    for (i = 0; i < CPU_SETSIZE; ++i) {
-        if (CPU_ISSET(i, &cpuset)) {
-            py_cpu_num = Py_BuildValue("i", i);
-            if (py_cpu_num == NULL)
-                goto error;
-            if (PyList_Append(py_retlist, py_cpu_num))
-                goto error;
-            Py_DECREF(py_cpu_num);
-        }
-    }
-
-    return py_retlist;
-
-error:
-    Py_XDECREF(py_cpu_num);
-    Py_XDECREF(py_retlist);
-    return NULL;
-}
-#endif
-
 /*
  * Set process CPU affinity; expects a bitmask
  */
@@ -432,7 +392,6 @@ psutil_proc_cpu_affinity_set(PyObject *self, PyObject *args) {
         CPU_SET(value, &cpu_set);
     }
 
-
     len = sizeof(cpu_set);
     if (sched_setaffinity(pid, len, &cpu_set)) {
         PyErr_SetFromErrno(PyExc_OSError);
@@ -447,6 +406,7 @@ psutil_proc_cpu_affinity_set(PyObject *self, PyObject *args) {
         Py_DECREF(py_cpu_seq);
     return NULL;
 }
+#endif  /* PSUTIL_HAVE_CPU_AFFINITY */
 
 
 /*
@@ -583,16 +543,18 @@ static PyMethodDef
 PsutilMethods[] = {
     // --- per-process functions
 
-#if PSUTIL_HAVE_IOPRIO
+#ifdef PSUTIL_HAVE_IOPRIO
     {"proc_ioprio_get", psutil_proc_ioprio_get, METH_VARARGS,
      "Get process I/O priority"},
     {"proc_ioprio_set", psutil_proc_ioprio_set, METH_VARARGS,
      "Set process I/O priority"},
 #endif
+#ifdef PSUTIL_HAVE_CPU_AFFINITY
     {"proc_cpu_affinity_get", psutil_proc_cpu_affinity_get, METH_VARARGS,
      "Return process CPU affinity as a Python long (the bitmask)."},
     {"proc_cpu_affinity_set", psutil_proc_cpu_affinity_set, METH_VARARGS,
      "Set process CPU affinity; expects a bitmask."},
+#endif
 
     // --- system related functions