fix #569: [FreeBSD] add support for process CPU affinity.

8 files changed, 149 insertions, 7 deletions

diff --git a/HISTORY.rst b/HISTORY.rst
index 8c1b4407..bd912f3c 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -7,6 +7,7 @@ Bug tracker at https://github.com/giampaolo/psutil/issues
 
 - #521: dropped support for Python 2.4 and 2.5.
 - #553: new examples/pstree.py script.
+- #569: [FreeBSD] added support for process CPU affinity.
 
 **Bug fixes**
 
diff --git a/docs/index.rst b/docs/index.rst
index 5e0914cb..0187bf74 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -847,7 +847,9 @@ Process class
       >>> p.cpu_affinity([0])  # set; from now on, process will run on CPU #0 only
       >>>
 
-     Availability: Linux, Windows
+     Availability: Linux, Windows, BSD
+
+     .. versionchanged:: 2.2.0 added support for FreeBSD
 
   .. method:: memory_info()
 
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 91c23ce5..2a77f224 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -672,13 +672,14 @@ class Process(object):
             else:
                 return self._proc.rlimit(resource, limits)
 
-    # Windows and Linux only
+    # Windows, Linux and BSD only
     if hasattr(_psplatform.Process, "cpu_affinity_get"):
 
         def cpu_affinity(self, cpus=None):
             """Get or set process CPU affinity.
             If specified 'cpus' must be a list of CPUs for which you
             want to set the affinity (e.g. [0, 1]).
+            (Windows, Linux and BSD only).
             """
             if cpus is None:
                 return self._proc.cpu_affinity_get()
diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py
index 83bfbd16..dfb22d23 100644
--- a/psutil/_psbsd.py
+++ b/psutil/_psbsd.py
@@ -387,3 +387,25 @@ class Process(object):
         proc_cwd = _not_implemented
         memory_maps = _not_implemented
         num_fds = _not_implemented
+
+    @wrap_exceptions
+    def cpu_affinity_get(self):
+        return cext.proc_cpu_affinity_get(self.pid)
+
+    @wrap_exceptions
+    def cpu_affinity_set(self, cpus):
+        try:
+            cext.proc_cpu_affinity_set(self.pid, cpus)
+        except OSError as err:
+            # 'man cpuset_setaffinity' about EDEADLK:
+            # <<the call would leave a thread without a valid CPU to run
+            # on because the set does not overlap with the thread's
+            # anonymous mask>>
+            if err.errno in (errno.EINVAL, errno.EDEADLK):
+                allcpus = tuple(range(len(per_cpu_times())))
+                for cpu in cpus:
+                    if cpu not in allcpus:
+                        raise ValueError("invalid CPU #%i (choose between %s)"
+                                         % (cpu, allcpus))
+            raise
+
diff --git a/psutil/_psutil_bsd.c b/psutil/_psutil_bsd.c
index fb5161bd..febe4004 100644
--- a/psutil/_psutil_bsd.c
+++ b/psutil/_psutil_bsd.c
@@ -21,6 +21,7 @@
 #include <sys/user.h>
 #include <sys/proc.h>
 #include <sys/file.h>
+#include <sys/cpuset.h>
 #include <net/route.h>
 
 #include <sys/socket.h>
@@ -2041,6 +2042,112 @@ error:
 
 
 /*
+ * Get process CPU affinity.
+ * Reference: http://sources.freebsd.org/RELENG_9/src/usr.bin/cpuset/cpuset.c
+ */
+static PyObject*
+psutil_proc_cpu_affinity_get(PyObject* self, PyObject* args)
+{
+    long pid;
+    int ret;
+    int i;
+    cpuset_t mask;
+    PyObject* py_retlist;
+    PyObject* py_cpu_num;
+
+    if (!PyArg_ParseTuple(args, "i", &pid)) {
+        return NULL;
+    }
+
+    ret = cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, pid,
+                             sizeof(mask), &mask);
+    if (ret != 0) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        return NULL;
+    }
+
+    py_retlist = PyList_New(0);
+    if (py_retlist == NULL)
+        return NULL;
+
+    for (i = 0; i < CPU_SETSIZE; i++) {
+        if (CPU_ISSET(i, &mask)) {
+            py_cpu_num = Py_BuildValue("i", i);
+            if (py_cpu_num == NULL)
+                goto error;
+            if (PyList_Append(py_retlist, py_cpu_num))
+                goto error;
+        }
+    }
+
+    return py_retlist;
+
+error:
+    Py_XDECREF(py_cpu_num);
+    Py_DECREF(py_retlist);
+    return NULL;
+}
+
+
+/*
+ * Set process CPU affinity.
+ * Reference: http://sources.freebsd.org/RELENG_9/src/usr.bin/cpuset/cpuset.c
+ */
+static PyObject *
+psutil_proc_cpu_affinity_set(PyObject *self, PyObject *args)
+{
+    long pid;
+    int i;
+    int seq_len;
+    int ret;
+    cpuset_t cpu_set;
+    PyObject *py_cpu_set;
+    PyObject *py_cpu_seq = NULL;
+
+    if (!PyArg_ParseTuple(args, "lO", &pid, &py_cpu_set)) {
+        goto error;
+    }
+
+    py_cpu_seq = PySequence_Fast(py_cpu_set, "expected a sequence or integer");
+    if (!py_cpu_seq) {
+        goto error;
+    }
+    seq_len = PySequence_Fast_GET_SIZE(py_cpu_seq);
+
+    // calculate the mask
+    CPU_ZERO(&cpu_set);
+    for (i = 0; i < seq_len; i++) {
+        PyObject *item = PySequence_Fast_GET_ITEM(py_cpu_seq, i);
+#if PY_MAJOR_VERSION >= 3
+        long value = PyLong_AsLong(item);
+#else
+        long value = PyInt_AsLong(item);
+#endif
+        if (value == -1 && PyErr_Occurred()) {
+            goto error;
+        }
+        CPU_SET(value, &cpu_set);
+    }
+
+    // set affinity
+    ret = cpuset_setaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, pid,
+                             sizeof(cpu_set), &cpu_set);
+    if (ret != 0) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+
+    Py_DECREF(py_cpu_seq);
+    Py_RETURN_NONE;
+
+error:
+    if (py_cpu_seq != NULL)
+        Py_DECREF(py_cpu_seq);
+    return NULL;
+}
+
+
+/*
  * define the psutil C module methods and initialize the module.
  */
 static PyMethodDef
@@ -2081,6 +2188,10 @@ PsutilMethods[] =
      "Return process IO counters"},
     {"proc_tty_nr", psutil_proc_tty_nr, METH_VARARGS,
      "Return process tty (terminal) number"},
+    {"proc_cpu_affinity_get", psutil_proc_cpu_affinity_get, METH_VARARGS,
+     "Return process CPU affinity."},
+    {"proc_cpu_affinity_set", psutil_proc_cpu_affinity_set, METH_VARARGS,
+     "Set process CPU affinity."},
 #if defined(__FreeBSD_version) && __FreeBSD_version >= 800000
     {"proc_open_files", psutil_proc_open_files, METH_VARARGS,
      "Return files opened by process as a list of (path, fd) tuples"},
diff --git a/psutil/_psutil_bsd.h b/psutil/_psutil_bsd.h
index 2bc7c701..803957da 100644
--- a/psutil/_psutil_bsd.h
+++ b/psutil/_psutil_bsd.h
@@ -26,6 +26,8 @@ static PyObject* psutil_proc_status(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_threads(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_tty_nr(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_uids(PyObject* self, PyObject* args);
+static PyObject* psutil_proc_cpu_affinity_get(PyObject* self, PyObject* args);
+static PyObject* psutil_proc_cpu_affinity_set(PyObject* self, PyObject* args);
 
 #if defined(__FreeBSD_version) && __FreeBSD_version >= 800000
 static PyObject* psutil_proc_open_files(PyObject* self, PyObject* args);
@@ -48,4 +50,4 @@ static PyObject* psutil_virtual_mem(PyObject* self, PyObject* args);
 
 #if defined(__FreeBSD_version) && __FreeBSD_version >= 800000
 static PyObject* psutil_per_cpu_times(PyObject* self, PyObject* args);
-#endif
\ No newline at end of file
+#endif
diff --git a/test/test_memory_leaks.py b/test/test_memory_leaks.py
index 963f9635..17eff121 100644
--- a/test/test_memory_leaks.py
+++ b/test/test_memory_leaks.py
@@ -26,7 +26,7 @@ import psutil
 import psutil._common
 
 from psutil._compat import callable, xrange
-from test_psutil import (WINDOWS, POSIX, OSX, LINUX, SUNOS, TESTFN,
+from test_psutil import (WINDOWS, POSIX, OSX, LINUX, SUNOS, BSD, TESTFN,
                          RLIMIT_SUPPORT)
 from test_psutil import (reap_children, supports_ipv6, safe_remove,
                          get_test_subprocess)
@@ -216,11 +216,13 @@ class TestProcessObjectLeaks(Base):
     def test_cwd(self):
         self.execute('cwd')
 
-    @unittest.skipUnless(WINDOWS or LINUX, "Windows or Linux only")
+    @unittest.skipUnless(WINDOWS or LINUX or BSD,
+                         "Windows or Linux or BSD only")
     def test_cpu_affinity_get(self):
         self.execute('cpu_affinity')
 
-    @unittest.skipUnless(WINDOWS or LINUX, "Windows or Linux only")
+    @unittest.skipUnless(WINDOWS or LINUX or BSD,
+                         "Windows or Linux or BSD only")
     def test_cpu_affinity_set(self):
         affinity = psutil.Process(os.getpid()).cpu_affinity()
         self.execute('cpu_affinity', affinity)
diff --git a/test/test_psutil.py b/test/test_psutil.py
index e62cccbb..605f1a96 100644
--- a/test/test_psutil.py
+++ b/test/test_psutil.py
@@ -1592,7 +1592,8 @@ class TestProcess(unittest.TestCase):
         p = psutil.Process(sproc.pid)
         call_until(p.cwd, "ret == os.path.dirname(os.getcwd())")
 
-    @unittest.skipUnless(WINDOWS or LINUX, 'not available on this platform')
+    @unittest.skipUnless(WINDOWS or LINUX or BSD,
+                         'not available on this platform')
     @unittest.skipIf(LINUX and TRAVIS, "unknown failure on travis")
     def test_cpu_affinity(self):
         p = psutil.Process()