#357: implement cpu_num() on SunOS

5 files changed, 95 insertions, 3 deletions

diff --git a/docs/index.rst b/docs/index.rst
index f9ecf42f..5c25c785 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1242,13 +1242,12 @@ Process class
   .. method:: cpu_num()
 
     Return what CPU this process is currently running on.
-    The returned number should be ``<=`` :func:`psutil.cpu_count()` and
-    ``<= len(psutil.cpu_percent(percpu=True))``.
+    The returned number should be ``<=`` :func:`psutil.cpu_count()`.
     It may be used in conjunction with ``psutil.cpu_percent(percpu=True)`` to
     observe the system workload distributed across multiple CPUs as shown by
     `cpu_workload.py <https://github.com/giampaolo/psutil/blob/master/scripts/cpu_workload.py>`__ example script.
 
-    Availability: Linux, FreeBSD
+    Availability: Linux, FreeBSD, SunOS
 
     .. versionadded:: 5.1.0
 
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 68fdee8f..b24b9822 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -837,6 +837,7 @@ class Process(object):
             else:
                 self._proc.cpu_affinity_set(list(set(cpus)))
 
+    # Linux, FreeBSD, SunOS
     if hasattr(_psplatform.Process, "cpu_num"):
 
         def cpu_num(self):
diff --git a/psutil/_pssunos.py b/psutil/_pssunos.py
index a62e0bf5..e6796bf9 100644
--- a/psutil/_pssunos.py
+++ b/psutil/_pssunos.py
@@ -469,6 +469,10 @@ class Process(object):
         return _common.pcputimes(*times)
 
     @wrap_exceptions
+    def cpu_num(self):
+        return cext.proc_cpu_num(self.pid, self._procfs_path)
+
+    @wrap_exceptions
     def terminal(self):
         procfs_path = self._procfs_path
         hit_enoent = False
diff --git a/psutil/_psutil_sunos.c b/psutil/_psutil_sunos.c
index e98ff7f2..48767add 100644
--- a/psutil/_psutil_sunos.c
+++ b/psutil/_psutil_sunos.c
@@ -169,6 +169,86 @@ psutil_proc_cpu_times(PyObject *self, PyObject *args) {
 
 
 /*
+ * Return what CPU the process is running on.
+ */
+static PyObject *
+psutil_proc_cpu_num(PyObject *self, PyObject *args) {
+    int fd = NULL;
+    int pid;
+    char path[1000];
+    struct prheader header;
+    struct lwpsinfo *lwp;
+    char *lpsinfo = NULL;
+    char *ptr = NULL;
+    int nent;
+    int size;
+    int proc_num;
+    size_t nbytes;
+    const char *procfs_path;
+
+    if (! PyArg_ParseTuple(args, "is", &pid, &procfs_path))
+        return NULL;
+
+    sprintf(path, "%s/%i/lpsinfo", procfs_path, pid);
+    fd = open(path, O_RDONLY);
+    if (fd == -1) {
+        PyErr_SetFromErrnoWithFilename(PyExc_OSError, path);
+        return NULL;
+    }
+
+    // read header
+    nbytes = pread(fd, &header, sizeof(header), 0);
+    if (nbytes == -1) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+    if (nbytes != sizeof(header)) {
+        PyErr_SetString(
+            PyExc_RuntimeError, "read() file structure size mismatch");
+        goto error;
+    }
+
+    // malloc
+    nent = header.pr_nent;
+    size = header.pr_entsize * nent;
+    ptr = lpsinfo = malloc(size);
+    if (lpsinfo == NULL) {
+        PyErr_NoMemory();
+        goto error;
+    }
+
+    // read the rest
+    nbytes = pread(fd, lpsinfo, size, sizeof(header));
+    if (nbytes == -1) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+    if (nbytes != size) {
+        PyErr_SetString(
+            PyExc_RuntimeError, "read() file structure size mismatch");
+        goto error;
+    }
+
+    // done
+    lwp = (lwpsinfo_t *)ptr;
+    proc_num = lwp->pr_onpro;
+    close(fd);
+    free(ptr);
+    free(lpsinfo);
+    return Py_BuildValue("i", proc_num);
+
+error:
+    if (fd != NULL)
+        close(fd);
+    if (ptr != NULL)
+        free(ptr);
+    if (lpsinfo != NULL)
+        free(lpsinfo);
+    return NULL;
+}
+
+
+/*
  * Return process uids/gids as a Python tuple.
  */
 static PyObject *
@@ -1340,6 +1420,8 @@ PsutilMethods[] = {
      "Return process memory mappings"},
     {"proc_num_ctx_switches", psutil_proc_num_ctx_switches, METH_VARARGS,
      "Return the number of context switches performed by process"},
+    {"proc_cpu_num", psutil_proc_cpu_num, METH_VARARGS,
+     "Return what CPU the process is on"},
 
     // --- system-related functions
     {"swap_mem", psutil_swap_mem, METH_VARARGS,
diff --git a/psutil/tests/test_memory_leaks.py b/psutil/tests/test_memory_leaks.py
index f1a951f0..6f724339 100755
--- a/psutil/tests/test_memory_leaks.py
+++ b/psutil/tests/test_memory_leaks.py
@@ -267,6 +267,12 @@ class TestProcessObjectLeaks(TestMemLeak):
         self.execute(self.proc.cpu_times)
 
     @skip_if_linux()
+    @unittest.skipUnless(hasattr(psutil.Process, "cpu_num"),
+                         "platform not supported")
+    def test_cpu_num(self):
+        self.execute(self.proc.cpu_num)
+
+    @skip_if_linux()
     def test_memory_info(self):
         self.execute(self.proc.memory_info)