Merge pull request #954 from giampaolo/357-proc-cpu-num

357 proc cpu num

12 files changed, 301 insertions, 20 deletions

diff --git a/docs/index.rst b/docs/index.rst
index 906a1e16..0131b56a 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -850,39 +850,41 @@ Process class
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
     | Linux                        | Windows                       | OSX                          | BSD                          | SunOS                    |
     +==============================+===============================+==============================+==============================+==========================+
-    | :meth:`~Process.cpu_percent` | :meth:`~Process.cpu_percent`  | :meth:`~Process.cpu_percent` | :meth:`~Process.cpu_percent` | :meth:`name`             |
+    | :meth:`cpu_num`              | :meth:`~Process.cpu_percent`  | :meth:`~Process.cpu_percent` | :meth:`cpu_num`              | :meth:`name`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`~Process.cpu_times`   | :meth:`~Process.cpu_times`    | :meth:`~Process.cpu_times`   | :meth:`~Process.cpu_times`   | :meth:`cmdline`          |
+    | :meth:`~Process.cpu_percent` | :meth:`~Process.cpu_times`    | :meth:`~Process.cpu_times`   | :meth:`~Process.cpu_percent` | :meth:`cmdline`          |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`create_time`          | :meth:`io_counters()`         | :meth:`memory_info`          | :meth:`create_time`          | :meth:`create_time`      |
+    | :meth:`~Process.cpu_times`   | :meth:`io_counters()`         | :meth:`memory_info`          | :meth:`~Process.cpu_times`   | :meth:`create_time`      |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`name`                 | :meth:`ionice`                | :meth:`memory_percent`       | :meth:`gids`                 |                          |
+    | :meth:`create_time`          | :meth:`ionice`                | :meth:`memory_percent`       | :meth:`create_time`          |                          |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`ppid`                 | :meth:`memory_info`           | :meth:`num_ctx_switches`     | :meth:`io_counters`          | :meth:`memory_info`      |
+    | :meth:`name`                 | :meth:`memory_info`           | :meth:`num_ctx_switches`     | :meth:`gids`                 | :meth:`memory_info`      |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`status`               | :meth:`nice`                  | :meth:`num_threads`          | :meth:`name`                 | :meth:`memory_percent`   |
+    | :meth:`ppid`                 | :meth:`nice`                  | :meth:`num_threads`          | :meth:`io_counters`          | :meth:`memory_percent`   |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`terminal`             | :meth:`memory_maps`           |                              | :meth:`memory_info`          | :meth:`nice`             |
+    | :meth:`status`               | :meth:`memory_maps`           |                              | :meth:`name`                 | :meth:`nice`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    |                              | :meth:`num_ctx_switches`      | :meth:`create_time`          | :meth:`memory_percent`       | :meth:`num_threads`      |
+    | :meth:`terminal`             | :meth:`num_ctx_switches`      | :meth:`create_time`          | :meth:`memory_info`          | :meth:`num_threads`      |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`gids`                 | :meth:`num_handles`           | :meth:`gids`                 | :meth:`num_ctx_switches`     | :meth:`ppid`             |
+    |                              | :meth:`num_handles`           | :meth:`gids`                 | :meth:`memory_percent`       | :meth:`ppid`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`num_ctx_switches`     | :meth:`num_threads`           | :meth:`name`                 | :meth:`ppid`                 | :meth:`status`           |
+    | :meth:`gids`                 | :meth:`num_threads`           | :meth:`name`                 | :meth:`num_ctx_switches`     | :meth:`status`           |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`num_threads`          | :meth:`username`              | :meth:`ppid`                 | :meth:`status`               | :meth:`terminal`         |
+    | :meth:`num_ctx_switches`     | :meth:`username`              | :meth:`ppid`                 | :meth:`ppid`                 | :meth:`terminal`         |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`uids`                 |                               | :meth:`status`               | :meth:`terminal`             |                          |
+    | :meth:`num_threads`          |                               | :meth:`status`               | :meth:`status`               |                          |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`username`             |                               | :meth:`terminal`             | :meth:`uids`                 | :meth:`gids`             |
+    | :meth:`uids`                 |                               | :meth:`terminal`             | :meth:`terminal`             | :meth:`gids`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    |                              |                               | :meth:`uids`                 | :meth:`username`             | :meth:`uids`             |
+    | :meth:`username`             |                               | :meth:`uids`                 | :meth:`uids`                 | :meth:`uids`             |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | :meth:`memory_full_info`     |                               | :meth:`username`             |                              | :meth:`username`         |
+    |                              |                               | :meth:`username`             | :meth:`username`             | :meth:`username`         |
+    +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
+    | :meth:`memory_full_info`     |                               |                              |                              |                          |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
     | :meth:`memory_maps`          |                               |                              |                              |                          |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
-    | *speedup: +2.5x*             | *speedup: +1.8x / +6.5x*      | *speedup: +1.9x*             | *speedup: +2.0x*             | *speedup: +1.3x*         |
+    | *speedup: +2.6x*             | *speedup: +1.8x / +6.5x*      | *speedup: +1.9x*             | *speedup: +2.0x*             | *speedup: +1.3x*         |
     +------------------------------+-------------------------------+------------------------------+------------------------------+--------------------------+
 
     .. versionadded:: 5.0.0
@@ -1239,6 +1241,18 @@ Process class
 
     .. versionchanged:: 2.2.0 added support for FreeBSD
 
+  .. method:: cpu_num()
+
+    Return what CPU this process is currently running on.
+    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, SunOS
+
+    .. versionadded:: 5.1.0
+
   .. method:: memory_info()
 
     Return a namedtuple with variable fields depending on the platform
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 77830230..b24b9822 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -837,6 +837,19 @@ 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):
+            """Return what CPU this process is currently running on.
+            The returned number should be <= psutil.cpu_count()
+            and <= len(psutil.cpu_percent(percpu=True)).
+            It may be used in conjunction with
+            psutil.cpu_percent(percpu=True) to observe the system
+            workload distributed across CPUs.
+            """
+            return self._proc.cpu_num()
+
     # Linux, OSX and Windows only
     if hasattr(_psplatform.Process, "environ"):
 
diff --git a/psutil/_psbsd.py b/psutil/_psbsd.py
index acf6c7c3..022f5758 100644
--- a/psutil/_psbsd.py
+++ b/psutil/_psbsd.py
@@ -122,7 +122,8 @@ kinfo_proc_map = dict(
     memtext=20,
     memdata=21,
     memstack=22,
-    name=23,
+    cpunum=23,
+    name=24,
 )
 
 
@@ -609,6 +610,11 @@ class Process(object):
             rawtuple[kinfo_proc_map['ch_user_time']],
             rawtuple[kinfo_proc_map['ch_sys_time']])
 
+    if FREEBSD:
+        @wrap_exceptions
+        def cpu_num(self):
+            return self.oneshot()[kinfo_proc_map['cpunum']]
+
     @wrap_exceptions
     def memory_info(self):
         rawtuple = self.oneshot()
diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
index 9154896a..ad95e39a 100644
--- a/psutil/_pslinux.py
+++ b/psutil/_pslinux.py
@@ -1299,6 +1299,11 @@ class Process(object):
         return _common.pcputimes(utime, stime, children_utime, children_stime)
 
     @wrap_exceptions
+    def cpu_num(self):
+        """What CPU the process is on."""
+        return int(self._parse_stat_file()[37])
+
+    @wrap_exceptions
     def wait(self, timeout=None):
         try:
             return _psposix.wait_pid(self.pid, timeout)
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_bsd.c b/psutil/_psutil_bsd.c
index afe3834e..adcedf79 100644
--- a/psutil/_psutil_bsd.c
+++ b/psutil/_psutil_bsd.c
@@ -198,6 +198,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];
@@ -252,9 +253,25 @@ psutil_proc_oneshot_info(PyObject *self, PyObject *args) {
         memstack = (long)kp.p_vm_ssize * pagesize;
 #endif
 
+#ifdef PSUTIL_FREEBSD
+    // 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;
+#else
+    // On Net/OpenBSD we have kp.p_cpuid but it appears it's always
+    // set to KI_NOCPU. Even if it's not, ki_lastcpu does not exist
+    // so there's no way to determine where "sleeping" processes
+    // were. Not supported.
+    oncpu = -1;
+#endif
+
     // Return a single big tuple with all process info.
     py_retlist = Py_BuildValue(
-        "(lillllllidllllddddlllllO)",
+        "(lillllllidllllddddlllllbO)",
 #ifdef PSUTIL_FREEBSD
         //
         (long)kp.ki_ppid,                // (long) ppid
@@ -287,6 +304,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
 #elif defined(PSUTIL_OPENBSD) || defined(PSUTIL_NETBSD)
         //
         (long)kp.p_ppid,                 // (long) ppid
@@ -321,6 +340,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/_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)
 
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index d25f4474..74750473 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -281,6 +281,16 @@ class TestProcess(unittest.TestCase):
         if (max([kernel_time, ktime]) - min([kernel_time, ktime])) > 0.1:
             self.fail("expected: %s, found: %s" % (ktime, kernel_time))
 
+    @unittest.skipUnless(hasattr(psutil.Process, "cpu_num"),
+                         "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):
         sproc = get_test_subprocess()
         now = time.time()
@@ -853,6 +863,10 @@ class TestProcess(unittest.TestCase):
             if hasattr(os, "sched_getaffinity"):
                 self.assertEqual(p.cpu_affinity(),
                                  list(os.sched_getaffinity(p.pid)))
+            # also test num_cpu()
+            if hasattr(p, "num_cpu"):
+                self.assertEqual(p.cpu_affinity()[0], p.num_cpu())
+
         #
         p.cpu_affinity(all_cpus)
         self.assertEqual(p.cpu_affinity(), all_cpus)
@@ -1723,6 +1737,12 @@ class TestFetchAllProcesses(unittest.TestCase):
         self.assertTrue(ret.user >= 0)
         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):
         for name in ret._fields:
             self.assertGreaterEqual(getattr(ret, name), 0)
@@ -1793,6 +1813,9 @@ class TestFetchAllProcesses(unittest.TestCase):
 
     def cpu_affinity(self, ret, proc):
         assert ret != [], ret
+        cpus = range(psutil.cpu_count())
+        for n in ret:
+            self.assertIn(n, cpus)
 
     def terminal(self, ret, proc):
         if ret is not None:
diff --git a/scripts/cpu_distribution.py b/scripts/cpu_distribution.py
new file mode 100755
index 00000000..31cdbb86
--- /dev/null
+++ b/scripts/cpu_distribution.py
@@ -0,0 +1,102 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""
+Shows CPU workload split across different CPUs.
+
+$ python scripts/cpu_workload.py
+CPU 0     CPU 1     CPU 2     CPU 3     CPU 4     CPU 5     CPU 6     CPU 7
+19.8      20.6      18.2      15.8      6.9       17.3      5.0       20.4
+gvfsd     pytho     kwork     chrom     unity     kwork     kwork     kwork
+chrom     chrom     indic     ibus-     whoop     nfsd      (sd-p     gvfsd
+ibus-     cat       at-sp     chrom     Modem     nfsd4     light     upsta
+ibus-     iprt-     ibus-     nacl_     cfg80     kwork     nfsd      bluet
+chrom     irqba     gpg-a     chrom     ext4-     biose     nfsd      dio/n
+chrom     acpid     bamfd     nvidi     kwork     scsi_     sshd      rpc.m
+upsta     rsysl     dbus-     nfsd      biose     scsi_     ext4-     polki
+rtkit     avahi     upowe     Netwo     scsi_     biose     UVM T     irq/9
+light     rpcbi     snapd     cron      ipv6_     biose     kwork     dbus-
+agett     kvm-i     avahi     kwork     biose     biose     scsi_     syste
+nfsd      syste     rpc.i     biose     biose     kbloc     kthro     UVM g
+nfsd      kwork     kwork     biose     vmsta     kwork     crypt     kaudi
+nfsd      scsi_     charg     biose     md        ksoft     kwork     kwork
+memca     biose     ksmd      ecryp     ksoft     watch     migra     nvme
+therm     biose     kcomp     kswap     migra     cpuhp     watch     biose
+syste     biose     kdevt     khuge     watch               cpuhp     biose
+led_w     devfr     kwork     write     cpuhp                         biose
+rpcio     oom_r     ksoft     kwork     syste                         biose
+kwork     kwork     watch     migra                                   acpi_
+biose     ksoft     cpuhp     watch                                   watch
+biose     migra               cpuhp                                   kinte
+biose     watch               rcu_s                                   netns
+biose     cpuhp               kthre                                   kwork
+cpuhp                                                                 ksoft
+watch                                                                 migra
+rcu_b                                                                 cpuhp
+kwork
+"""
+
+from __future__ import print_function
+import collections
+import os
+import sys
+import time
+
+import psutil
+
+
+if not hasattr(psutil.Process, "cpu_num"):
+    sys.exit("platform not supported")
+
+
+def clean_screen():
+    if psutil.POSIX:
+        os.system('clear')
+    else:
+        os.system('cls')
+
+
+def main():
+    total = psutil.cpu_count()
+    while True:
+        # header
+        clean_screen()
+        cpus_percent = psutil.cpu_percent(percpu=True)
+        for i in range(total):
+            print("CPU %-6i" % i, end="")
+        print()
+        for percent in cpus_percent:
+            print("%-10s" % percent, end="")
+        print()
+
+        # processes
+        procs = collections.defaultdict(list)
+        for p in psutil.process_iter():
+            try:
+                name = p.name()[:5]
+                cpunum = p.cpu_num()
+            except psutil.Error:
+                continue
+            else:
+                procs[cpunum].append(name)
+
+        end_marker = [[] for x in range(total)]
+        while True:
+            for num in range(total):
+                try:
+                    pname = procs[num].pop()
+                except IndexError:
+                    pname = ""
+                print("%-10s" % pname[:10], end="")
+            print()
+            if procs.values() == end_marker:
+                break
+
+        time.sleep(1)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/scripts/internal/bench_oneshot.py b/scripts/internal/bench_oneshot.py
index cf8497f8..639e9ad7 100755
--- a/scripts/internal/bench_oneshot.py
+++ b/scripts/internal/bench_oneshot.py
@@ -37,10 +37,11 @@ if psutil.POSIX:
 
 if psutil.LINUX:
     names += [
-        'cpu_times',
-        'gids',
         # 'memory_full_info',
         # 'memory_maps',
+        'cpu_num',
+        'cpu_times',
+        'gids',
         'name',
         'num_ctx_switches',
         'num_threads',
@@ -63,6 +64,8 @@ elif psutil.BSD:
         'terminal',
         'uids',
     ]
+    if psutil.FREEBSD:
+        names.append('cpu_num')
 elif psutil.SUNOS:
     names += [
         'cmdline',
diff --git a/scripts/procinfo.py b/scripts/procinfo.py
index 8dc34c45..d8625560 100755
--- a/scripts/procinfo.py
+++ b/scripts/procinfo.py
@@ -190,6 +190,8 @@ def run(pid, verbose=False):
     print_('cpu-times', str_ntuple(pinfo['cpu_times']))
     if hasattr(proc, "cpu_affinity"):
         print_("cpu-affinity", pinfo["cpu_affinity"])
+    if hasattr(proc, "cpu_num"):
+        print_("cpu-num", pinfo["cpu_num"])
 
     print_('memory', str_ntuple(pinfo['memory_info'], bytes2human=True))
     print_('memory %', round(pinfo['memory_percent'], 2))