#357: implement process cpu_num() on FreeBSD

3 files changed, 27 insertions, 4 deletions

diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py
index a3301504..7cb9f2cd 100644
--- a/psutil/_psbsd.py
+++ b/psutil/_psbsd.py
@@ -121,7 +121,8 @@ kinfo_proc_map = dict(
     memtext=20,
     memdata=21,
     memstack=22,
-    name=23,
+    cpunum=23,
+    name=24,
 )
 
 
@@ -590,6 +591,10 @@ class Process(object):
             rawtuple[kinfo_proc_map['ch_sys_time']])
 
     @wrap_exceptions
+    def cpu_num(self):
+        return self.oneshot()[kinfo_proc_map['cpunum']]
+
+    @wrap_exceptions
     def memory_info(self):
         rawtuple = self.oneshot()
         return pmem(
diff --git a/psutil/_psutil_bsd.c b/psutil/_psutil_bsd.c
index 33448ccf..48391279 100644
--- a/psutil/_psutil_bsd.c
+++ b/psutil/_psutil_bsd.c
@@ -203,6 +203,7 @@ psutil_proc_oneshot_info(PyObject *self, PyObject *args) {
     long memtext;
     long memdata;
     long memstack;
+    unsigned char oncpu;
     kinfo_proc kp;
     long pagesize = sysconf(_SC_PAGESIZE);
     char str[1000];
@@ -257,11 +258,18 @@ psutil_proc_oneshot_info(PyObject *self, PyObject *args) {
         memstack = (long)kp.p_vm_ssize * pagesize;
 #endif
 
+    // what CPU we're on; top was used as an example:
+    // https://svnweb.freebsd.org/base/head/usr.bin/top/machine.c?
+    //     view=markup&pathrev=273835
+    if (kp.ki_stat == SRUN && kp.ki_oncpu != NOCPU)
+        oncpu = kp.ki_oncpu;
+    else
+        oncpu = kp.ki_lastcpu;
+
     // Return a single big tuple with all process info.
     py_retlist = Py_BuildValue(
-        "(lillllllidllllddddlllllO)",
+        "(lillllllidllllddddlllllbO)",
 #ifdef __FreeBSD__
-        //
         (long)kp.ki_ppid,                // (long) ppid
         (int)kp.ki_stat,                 // (int) status
         // UIDs
@@ -292,8 +300,9 @@ psutil_proc_oneshot_info(PyObject *self, PyObject *args) {
         memtext,                          // (long) mem text
         memdata,                          // (long) mem data
         memstack,                         // (long) mem stack
+        // others
+        oncpu,                            // (unsigned char) the CPU we are on
 #elif defined(__OpenBSD__) || defined(__NetBSD__)
-        //
         (long)kp.p_ppid,                 // (long) ppid
         (int)kp.p_stat,                  // (int) status
         // UIDs
@@ -326,6 +335,8 @@ psutil_proc_oneshot_info(PyObject *self, PyObject *args) {
         memtext,                          // (long) mem text
         memdata,                          // (long) mem data
         memstack,                         // (long) mem stack
+        // others
+        oncpu,                            // (unsigned char) the CPU we are on
 #endif
         py_name                           // (pystr) name
     );
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 49aa4d86..afbdc690 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -285,6 +285,10 @@ class TestProcess(unittest.TestCase):
                          "platform not supported")
     def test_cpu_num(self):
         p = psutil.Process()
+        num = p.cpu_num()
+        self.assertGreaterEqual(num, 0)
+        if psutil.cpu_count() == 1:
+            self.assertEqual(num, 0)
         self.assertIn(p.cpu_num(), range(psutil.cpu_count()))
 
     def test_create_time(self):
@@ -1735,6 +1739,9 @@ class TestFetchAllProcesses(unittest.TestCase):
         self.assertTrue(ret.system >= 0)
 
     def cpu_num(self, ret, proc):
+        self.assertGreaterEqual(ret, 0)
+        if psutil.cpu_count() == 1:
+            self.assertEqual(ret, 0)
         self.assertIn(ret, range(psutil.cpu_count()))
 
     def memory_info(self, ret, proc):