AIX: implement num_ctx_switches (#1164)

* small changes

* AIX: implement num_ctx_switches

16 files changed, 185 insertions, 89 deletions

diff --git a/MANIFEST.in b/MANIFEST.in
index a07f16da..59a102a5 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -45,6 +45,8 @@ include psutil/arch/aix/ifaddrs.c
 include psutil/arch/aix/ifaddrs.h
 include psutil/arch/aix/net_connections.c
 include psutil/arch/aix/net_connections.h
+include psutil/arch/aix/common.c
+include psutil/arch/aix/common.h
 include psutil/arch/aix/net_kernel_structs.h
 include psutil/arch/freebsd/proc_socks.c
 include psutil/arch/freebsd/proc_socks.h
diff --git a/docs/index.rst b/docs/index.rst
index feb9fd78..ac591d30 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1376,7 +1376,7 @@ Process class
     The number voluntary and involuntary context switches performed by
     this process (cumulative).
 
-    Availability: all platforms except AIX
+    .. versionchanged:: 5.4.1 added AIX support
 
   .. method:: num_fds()
 
diff --git a/psutil/__init__.py b/psutil/__init__.py
index b8b94e7e..97d80940 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -899,13 +899,11 @@ class Process(object):
             """
             return self._proc.num_handles()
 
-    if hasattr(_psplatform.Process, "num_ctx_switches"):
-
-        def num_ctx_switches(self):
-            """Return the number of voluntary and involuntary context
-            switches performed by this process.
-            """
-            return self._proc.num_ctx_switches()
+    def num_ctx_switches(self):
+        """Return the number of voluntary and involuntary context
+        switches performed by this process.
+        """
+        return self._proc.num_ctx_switches()
 
     def num_threads(self):
         """Return the number of threads used by this process."""
diff --git a/psutil/_psaix.py b/psutil/_psaix.py
index 5cd088e5..c78922b0 100644
--- a/psutil/_psaix.py
+++ b/psutil/_psaix.py
@@ -555,6 +555,11 @@ class Process(object):
         return len(os.listdir("%s/%s/fd" % (self._procfs_path, self.pid)))
 
     @wrap_exceptions
+    def num_ctx_switches(self):
+        return _common.pctxsw(
+            *cext.proc_num_ctx_switches(self.pid))
+
+    @wrap_exceptions
     def wait(self, timeout=None):
         try:
             return _psposix.wait_pid(self.pid, timeout)
diff --git a/psutil/_psutil_aix.c b/psutil/_psutil_aix.c
index 8ffc8f47..0834726d 100644
--- a/psutil/_psutil_aix.c
+++ b/psutil/_psutil_aix.c
@@ -4,16 +4,16 @@
  * All rights reserved.
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
- /*
+ */
 
 /*
  * AIX support is experimental at this time.
  * The following functions and methods are unsupported on the AIX platform:
  * - psutil.Process.memory_maps
- * - psutil.Process.num_ctx_switches
  *
  * Known limitations:
  * - psutil.Process.io_counters read count is always 0
+ * - psutil.Process.threads may not be available on older AIX versions
  * - reading basic process info may fail or return incorrect values when
  *   process is starting (see IBM APAR IV58499 - fixed in newer AIX versions)
  * - sockets and pipes may not be counted in num_fds (fixed in newer AIX
@@ -49,6 +49,7 @@
 
 #include "arch/aix/ifaddrs.h"
 #include "arch/aix/net_connections.h"
+#include "arch/aix/common.h"
 #include "_psutil_common.h"
 #include "_psutil_posix.h"
 
@@ -308,6 +309,43 @@ psutil_proc_cred(PyObject *self, PyObject *args) {
 
 
 /*
+ * Return process voluntary and involuntary context switches as a Python tuple.
+ */
+static PyObject *
+psutil_proc_num_ctx_switches(PyObject *self, PyObject *args) {
+    PyObject *py_tuple = NULL;
+    pid32_t requested_pid;
+    pid32_t pid = 0;
+    int np = 0;
+    struct procentry64 *processes = (struct procentry64 *)NULL;
+    struct procentry64 *p;
+
+    if (! PyArg_ParseTuple(args, "i", &requested_pid))
+        return NULL;
+
+    processes = psutil_read_process_table(&np);
+    if (!processes)
+        return NULL;
+
+    /* Loop through processes */
+    for (p = processes; np > 0; np--, p++) {
+        pid = p->pi_pid;
+        if (requested_pid != pid)
+            continue;
+        py_tuple = Py_BuildValue("LL",
+            (long long) p->pi_ru.ru_nvcsw,    /* voluntary context switches */
+            (long long) p->pi_ru.ru_nivcsw);  /* involuntary */
+        free(processes);
+        return py_tuple;
+    }
+
+    /* finished iteration without finding requested pid */
+    free(processes);
+    return NoSuchProcess();
+}
+
+
+/*
  * Return users currently connected on the system.
  */
 static PyObject *
@@ -833,6 +871,8 @@ PsutilMethods[] =
 #endif
     {"proc_io_counters", psutil_proc_io_counters, METH_VARARGS,
      "Get process I/O counters."},
+    {"proc_num_ctx_switches", psutil_proc_num_ctx_switches, METH_VARARGS,
+     "Get process I/O counters."},
 
     // --- system-related functions
     {"users", psutil_users, METH_VARARGS,
diff --git a/psutil/_psutil_sunos.c b/psutil/_psutil_sunos.c
index 8f977342..083b78cb 100644
--- a/psutil/_psutil_sunos.c
+++ b/psutil/_psutil_sunos.c
@@ -360,7 +360,7 @@ psutil_proc_cred(PyObject *self, PyObject *args) {
 
 
 /*
- * Return process uids/gids as a Python tuple.
+ * Return process voluntary and involuntary context switches as a Python tuple.
  */
 static PyObject *
 psutil_proc_num_ctx_switches(PyObject *self, PyObject *args) {
diff --git a/psutil/arch/aix/common.c b/psutil/arch/aix/common.c
new file mode 100644
index 00000000..6115a15d
--- /dev/null
+++ b/psutil/arch/aix/common.c
@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) 2017, Arnon Yaari
+ * All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include <Python.h>
+#include <sys/core.h>
+#include <stdlib.h>
+#include "common.h"
+
+/* psutil_kread() - read from kernel memory */
+int
+psutil_kread(
+    int Kd,             /* kernel memory file descriptor */
+    KA_T addr,          /* kernel memory address */
+    char *buf,          /* buffer to receive data */
+    size_t len) {       /* length to read */
+    int br;
+
+    if (lseek64(Kd, (off64_t)addr, L_SET) == (off64_t)-1) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        return 1;
+    }
+    br = read(Kd, buf, len);
+    if (br == -1) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        return 1;
+    }
+    if (br != len) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "size mismatch when reading kernel memory fd");
+        return 1;
+    }
+    return 0;
+}
+
+struct procentry64 *
+psutil_read_process_table(int * num) {
+    size_t msz;
+    pid32_t pid = 0;
+    struct procentry64 *processes = (struct procentry64 *)NULL;
+    struct procentry64 *p;
+    int Np = 0;          /* number of processes allocated in 'processes' */
+    int np = 0;          /* number of processes read into 'processes' */
+    int i;               /* number of processes read in current iteration */
+
+    msz = (size_t)(PROCSIZE * PROCINFO_INCR);
+    processes = (struct procentry64 *)malloc(msz);
+    if (!processes) {
+        PyErr_NoMemory();
+        return NULL;
+    }
+    Np = PROCINFO_INCR;
+    p = processes;
+    while ((i = getprocs64(p, PROCSIZE, (struct fdsinfo64 *)NULL, 0, &pid,
+                 PROCINFO_INCR))
+    == PROCINFO_INCR) {
+        np += PROCINFO_INCR;
+        if (np >= Np) {
+            msz = (size_t)(PROCSIZE * (Np + PROCINFO_INCR));
+            processes = (struct procentry64 *)realloc((char *)processes, msz);
+            if (!processes) {
+                PyErr_NoMemory();
+                return NULL;
+            }
+            Np += PROCINFO_INCR;
+        }
+        p = (struct procentry64 *)((char *)processes + (np * PROCSIZE));
+    }
+
+    /* add the number of processes read in the last iteration */
+    if (i > 0)
+        np += i;
+
+    *num = np;
+    return processes;
+}
\ No newline at end of file
diff --git a/psutil/arch/aix/common.h b/psutil/arch/aix/common.h
new file mode 100644
index 00000000..b677d8c2
--- /dev/null
+++ b/psutil/arch/aix/common.h
@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2017, Arnon Yaari
+ * All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#ifndef __PSUTIL_AIX_COMMON_H__
+#define __PSUTIL_AIX_COMMON_H__
+
+#include <sys/core.h>
+
+#define PROCINFO_INCR   (256)
+#define PROCSIZE        (sizeof(struct procentry64))
+#define FDSINFOSIZE     (sizeof(struct fdsinfo64))
+#define KMEM            "/dev/kmem"
+
+typedef u_longlong_t    KA_T;
+
+/* psutil_kread() - read from kernel memory */
+int psutil_kread(int Kd,     /* kernel memory file descriptor */
+    KA_T addr,               /* kernel memory address */
+    char *buf,               /* buffer to receive data */
+    size_t len);             /* length to read */
+
+struct procentry64 *
+psutil_read_process_table(
+    int * num               /* out - number of processes read */
+);
+
+#endif  /* __PSUTIL_AIX_COMMON_H__ */
diff --git a/psutil/arch/aix/net_connections.c b/psutil/arch/aix/net_connections.c
index 364cd1b7..69b43892 100644
--- a/psutil/arch/aix/net_connections.c
+++ b/psutil/arch/aix/net_connections.c
@@ -13,57 +13,26 @@
  * - dialects/aix/dproc.c:get_kernel_access
 */
 
-#include "net_connections.h"
+#include <Python.h>
+#include <stdlib.h>
 #include <fcntl.h>
-#include <sys/types.h>
-#define _KERNEL 1
+#define _KERNEL
 #include <sys/file.h>
 #undef _KERNEL
-#include <stdlib.h>
+#include <sys/types.h>
 #include <sys/core.h>
 #include <sys/domain.h>
 #include <sys/un.h>
 #include <netinet/in_pcb.h>
 #include <arpa/inet.h>
-#include "net_kernel_structs.h"
-
 
+#include "../../_psutil_common.h"
+#include "net_kernel_structs.h"
+#include "net_connections.h"
+#include "common.h"
 
-#define PROCINFO_INCR   (256)
-#define PROCSIZE        (sizeof(struct procentry64))
-#define FDSINFOSIZE     (sizeof(struct fdsinfo64))
-#define KMEM            "/dev/kmem"
 #define NO_SOCKET       (PyObject *)(-1)
 
-typedef u_longlong_t    KA_T;
-static int PSUTIL_CONN_NONE = 128;
-
-/* psutil_kread() - read from kernel memory */
-static int
-psutil_kread(
-    int Kd,             /* kernel memory file descriptor */
-    KA_T addr,          /* kernel memory address */
-    char *buf,          /* buffer to receive data */
-    size_t len) {       /* length to read */
-    int br;
-
-    if (lseek64(Kd, (off64_t)addr, L_SET) == (off64_t)-1) {
-        PyErr_SetFromErrno(PyExc_OSError);
-        return 1;
-    }
-    br = read(Kd, buf, len);
-    if (br == -1) {
-        PyErr_SetFromErrno(PyExc_OSError);
-        return 1;
-    }
-    if (br != len) {
-        PyErr_SetString(PyExc_RuntimeError,
-                        "size mismatch when reading kernel memory fd");
-        return 1;
-    }
-    return 0;
-}
-
 static int
 read_unp_addr(
     int Kd,
@@ -244,10 +213,8 @@ psutil_net_connections(PyObject *self, PyObject *args) {
     int i, np;
     struct procentry64 *p;
     struct fdsinfo64 *fds = (struct fdsinfo64 *)NULL;
-    size_t msz;
     pid32_t requested_pid;
     pid32_t pid;
-    int Np = 0;          /* number of processes */
     struct procentry64 *processes = (struct procentry64 *)NULL;
                     /* the process table */
 
@@ -262,34 +229,9 @@ psutil_net_connections(PyObject *self, PyObject *args) {
         goto error;
     }
 
-    /* Read the process table */
-    msz = (size_t)(PROCSIZE * PROCINFO_INCR);
-    processes = (struct procentry64 *)malloc(msz);
-    if (!processes) {
-        PyErr_NoMemory();
+    processes = psutil_read_process_table(&np);
+    if (!processes)
         goto error;
-    }
-    Np = PROCINFO_INCR;
-    np = pid = 0;
-    p = processes;
-    while ((i = getprocs64(p, PROCSIZE, (struct fdsinfo64 *)NULL, 0, &pid,
-                 PROCINFO_INCR))
-    == PROCINFO_INCR) {
-        np += PROCINFO_INCR;
-        if (np >= Np) {
-            msz = (size_t)(PROCSIZE * (Np + PROCINFO_INCR));
-            processes = (struct procentry64 *)realloc((char *)processes, msz);
-            if (!processes) {
-                PyErr_NoMemory();
-                goto error;
-            }
-            Np += PROCINFO_INCR;
-        }
-        p = (struct procentry64 *)((char *)processes + (np * PROCSIZE));
-    }
-
-    if (i > 0)
-        np += i;
 
     /* Loop through processes */
     for (p = processes; np > 0; np--, p++) {
@@ -299,7 +241,6 @@ psutil_net_connections(PyObject *self, PyObject *args) {
         if (p->pi_state == 0 || p->pi_state == SZOMB)
             continue;
 
-
         if (!fds) {
             fds = (struct fdsinfo64 *)malloc((size_t)FDSINFOSIZE);
             if (!fds) {
diff --git a/psutil/arch/aix/net_connections.h b/psutil/arch/aix/net_connections.h
index f6a726cb..222bcaf3 100644
--- a/psutil/arch/aix/net_connections.h
+++ b/psutil/arch/aix/net_connections.h
@@ -5,6 +5,11 @@
  * found in the LICENSE file.
  */
 
+#ifndef __NET_CONNECTIONS_H__
+#define __NET_CONNECTIONS_H__
+
 #include <Python.h>
 
 PyObject* psutil_net_connections(PyObject *self, PyObject *args);
+
+#endif  /* __NET_CONNECTIONS_H__ */
\ No newline at end of file
diff --git a/psutil/arch/solaris/v10/ifaddrs.h b/psutil/arch/solaris/v10/ifaddrs.h
index d2771193..0953a9b9 100644
--- a/psutil/arch/solaris/v10/ifaddrs.h
+++ b/psutil/arch/solaris/v10/ifaddrs.h
@@ -1,8 +1,8 @@
 /* Reference: https://lists.samba.org/archive/samba-technical/2009-February/063079.html */
 
 
-#ifndef __IFADDRS_H___
-#define __IFADDRS_H___
+#ifndef __IFADDRS_H__
+#define __IFADDRS_H__
 
 #include <sys/socket.h>
 #include <net/if.h>
diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index cb46a463..14f1b53f 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -72,7 +72,6 @@ __all__ = [
     "HAS_IONICE", "HAS_MEMORY_MAPS", "HAS_PROC_CPU_NUM", "HAS_RLIMIT",
     "HAS_SENSORS_BATTERY", "HAS_BATTERY""HAS_SENSORS_FANS",
     "HAS_SENSORS_TEMPERATURES", "HAS_MEMORY_FULL_INFO",
-    "HAS_NUM_CTX_SWITCHES",
     # subprocesses
     'pyrun', 'reap_children', 'get_test_subprocess', 'create_zombie_proc',
     'create_proc_children_pair',
@@ -157,7 +156,6 @@ HAS_PROC_IO_COUNTERS = hasattr(psutil.Process, "io_counters")
 HAS_IONICE = hasattr(psutil.Process, "ionice")
 HAS_MEMORY_FULL_INFO = 'uss' in psutil.Process().memory_full_info()._fields
 HAS_MEMORY_MAPS = hasattr(psutil.Process, "memory_maps")
-HAS_NUM_CTX_SWITCHES = hasattr(psutil.Process, "num_ctx_switches")
 HAS_PROC_CPU_NUM = hasattr(psutil.Process, "cpu_num")
 HAS_RLIMIT = hasattr(psutil.Process, "rlimit")
 HAS_THREADS = hasattr(psutil.Process, "threads")
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index 13a737e8..d9633339 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -611,7 +611,7 @@ class TestFetchAllProcesses(unittest.TestCase):
     def num_ctx_switches(self, ret, proc):
         assert is_namedtuple(ret)
         for value in ret:
-            self.assertIsInstance(value, int)
+            self.assertIsInstance(value, (int, long))
             self.assertGreaterEqual(value, 0)
 
     def rlimit(self, ret, proc):
diff --git a/psutil/tests/test_memory_leaks.py b/psutil/tests/test_memory_leaks.py
index 76fab357..680fe780 100755
--- a/psutil/tests/test_memory_leaks.py
+++ b/psutil/tests/test_memory_leaks.py
@@ -38,7 +38,6 @@ from psutil.tests import HAS_CPU_FREQ
 from psutil.tests import HAS_ENVIRON
 from psutil.tests import HAS_IONICE
 from psutil.tests import HAS_MEMORY_MAPS
-from psutil.tests import HAS_NUM_CTX_SWITCHES
 from psutil.tests import HAS_PROC_CPU_NUM
 from psutil.tests import HAS_PROC_IO_COUNTERS
 from psutil.tests import HAS_RLIMIT
@@ -289,7 +288,6 @@ class TestProcessObjectLeaks(TestMemLeak):
         self.execute(self.proc.num_fds)
 
     @skip_if_linux()
-    @unittest.skipIf(not HAS_NUM_CTX_SWITCHES, "not supported")
     def test_num_ctx_switches(self):
         self.execute(self.proc.num_ctx_switches)
 
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 4d0a783c..1e01aea5 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -45,7 +45,6 @@ from psutil.tests import HAS_CPU_AFFINITY
 from psutil.tests import HAS_ENVIRON
 from psutil.tests import HAS_IONICE
 from psutil.tests import HAS_MEMORY_MAPS
-from psutil.tests import HAS_NUM_CTX_SWITCHES
 from psutil.tests import HAS_PROC_CPU_NUM
 from psutil.tests import HAS_PROC_IO_COUNTERS
 from psutil.tests import HAS_RLIMIT
@@ -1004,7 +1003,6 @@ class TestProcess(unittest.TestCase):
 
     @skip_on_not_implemented(only_if=LINUX)
     @unittest.skipIf(OPENBSD or NETBSD, "not reliable on OPENBSD & NETBSD")
-    @unittest.skipIf(not HAS_NUM_CTX_SWITCHES, "not supported")
     def test_num_ctx_switches(self):
         p = psutil.Process()
         before = sum(p.num_ctx_switches())
diff --git a/setup.py b/setup.py
index d11dd782..a170d2de 100755
--- a/setup.py
+++ b/setup.py
@@ -248,6 +248,7 @@ elif AIX:
         sources=sources + [
             'psutil/_psutil_aix.c',
             'psutil/arch/aix/net_connections.c',
+            'psutil/arch/aix/common.c',
             'psutil/arch/aix/ifaddrs.c'],
         libraries=['perfstat'],
         define_macros=macros)