Implement getloadavg on Windows. Fixes #604 and #1484 (#1485) (patch by Ammar Askar)

15 files changed, 229 insertions, 12 deletions

diff --git a/.ci/travis/install.sh b/.ci/travis/install.sh
index b0f28c66..1e37c39b 100755
--- a/.ci/travis/install.sh
+++ b/.ci/travis/install.sh
@@ -21,8 +21,8 @@ if [[ "$(uname -s)" == 'Darwin' ]]; then
         #     pyenv virtualenv 2.6.9 psutil
         #     ;;
         py27)
-            pyenv install 2.7.10
-            pyenv virtualenv 2.7.10 psutil
+            pyenv install 2.7.16
+            pyenv virtualenv 2.7.16 psutil
             ;;
         py36)
             pyenv install 3.6.6
diff --git a/CREDITS b/CREDITS
index 33c62ea6..488ce3e3 100644
--- a/CREDITS
+++ b/CREDITS
@@ -607,3 +607,8 @@ I: 1471
 N: Samer Masterson
 W: https://github.com/samertm
 I: 1480
+
+N: Ammar Askar
+E: ammar@ammaraskar.com
+W: http://ammaraskar.com/
+I: 604, 1484
\ No newline at end of file
diff --git a/HISTORY.rst b/HISTORY.rst
index 5741e247..5427ef4e 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -14,6 +14,8 @@
   new constants: IOPRIO_VERYLOW, IOPRIO_LOW, IOPRIO_NORMAL, IOPRIO_HIGH.
   Also it was not possible to set high I/O priority (not it is).
 - 1478_: add make command to re-run tests failed on last run.
+- 604: [UNIX, Windows] add new psutil.getloadavg() returning system load 
+  average calculation. (patch by Ammar Askar)
 
 **Bug fixes**
 
diff --git a/docs/DEVNOTES b/docs/DEVNOTES
index c54fa938..abd2e368 100644
--- a/docs/DEVNOTES
+++ b/docs/DEVNOTES
@@ -69,8 +69,6 @@ FEATURES
 
 - #613: thread names.
 
-- #604: emulate os.getloadavg() on Windows
-
 - scripts/taskmgr-gui.py (using tk).
 
 - system-wide number of open file descriptors:
diff --git a/docs/index.rst b/docs/index.rst
index 85429eb6..1469e61b 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -238,6 +238,30 @@ CPU
 
     .. versionchanged:: 5.5.1 added FreeBSD support.
 
+.. function:: getloadavg()
+
+    Returns the average load on the system over the last 1, 5 and 15 minutes
+    respectively as a tuple. The load represents how many processes are waiting
+    to be run by the operating system.
+
+    On UNIX systems this relies on `os.getloadavg`_. On Windows, this is
+    emulated by using a Windows API call that spawns a thread which updates the
+    average every 5 seconds mimicking the UNIX behavior. Thus, the first time
+    this is called and up until 5 seconds it returns a meaningless
+    ``(0.0, 0.0, 0.0)`` tuple.
+
+    Example:
+
+    .. code-block:: python
+
+       >>> import psutil
+       >>> psutil.getloadavg()
+       (3.14, 3.89, 4.67)
+
+    Availability: Unix, Windows
+
+    .. versionadded:: 5.6.2
+
 Memory
 ------
 
@@ -2556,12 +2580,6 @@ FAQs
   the Python script as a Windows service (this is the trick used by tools
   such as ProcessHacker).
 
-----
-
-* Q: What about load average?
-* A: psutil does not expose any load average function as it's already available
-  in python as `os.getloadavg`_.
-
 Running tests
 =============
 
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 5d2b8d3c..07f1104b 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -2017,6 +2017,17 @@ if hasattr(_psplatform, "cpu_freq"):
     __all__.append("cpu_freq")
 
 
+if hasattr(os, "getloadavg") or hasattr(_psplatform, "getloadavg"):
+    # Perform this hasattr check once on import time to either use the
+    # platform based code or proxy straight from the os module.
+    if hasattr(os, "getloadavg"):
+        getloadavg = os.getloadavg
+    else:
+        getloadavg = _psplatform.getloadavg
+
+    __all__.append("getloadavg")
+
+
 # =====================================================================
 # --- system memory related functions
 # =====================================================================
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 53617c58..f071648f 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -40,6 +40,7 @@
 #include "arch/windows/process_handles.h"
 #include "arch/windows/inet_ntop.h"
 #include "arch/windows/services.h"
+#include "arch/windows/wmi.h"
 #include "_psutil_common.h"
 
 
@@ -3508,6 +3509,12 @@ PsutilMethods[] = {
      "Return NICs stats."},
     {"cpu_freq", psutil_cpu_freq, METH_VARARGS,
      "Return CPU frequency."},
+#if (_WIN32_WINNT >= 0x0600)  // Windows Vista
+    {"init_loadavg_counter", psutil_init_loadavg_counter, METH_VARARGS,
+     "Initializes the emulated load average calculator."},
+    {"getloadavg", psutil_get_loadavg, METH_VARARGS,
+     "Returns the emulated POSIX-like load average."},
+#endif
     {"sensors_battery", psutil_sensors_battery, METH_VARARGS,
      "Return battery metrics usage."},
     {"getpagesize", psutil_getpagesize, METH_VARARGS,
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 929e27d7..3f131980 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -80,6 +80,7 @@ __extra__all__ = [
 
 CONN_DELETE_TCB = "DELETE_TCB"
 HAS_PROC_IO_PRIORITY = hasattr(cext, "proc_io_priority_get")
+HAS_GETLOADAVG = hasattr(cext, "getloadavg")
 
 
 if enum is None:
@@ -353,6 +354,23 @@ def cpu_freq():
     return [_common.scpufreq(float(curr), min_, float(max_))]
 
 
+if HAS_GETLOADAVG:
+    _loadavg_inititialized = False
+
+    def getloadavg():
+        """Return the number of processes in the system run queue averaged
+        over the last 1, 5, and 15 minutes respectively as a tuple"""
+        global _loadavg_inititialized
+
+        if not _loadavg_inititialized:
+            cext.init_loadavg_counter()
+            _loadavg_inititialized = True
+
+        # Drop to 2 decimal points which is what Linux does
+        raw_loads = cext.getloadavg()
+        return tuple([round(load, 2) for load in raw_loads])
+
+
 # =====================================================================
 # --- network
 # =====================================================================
diff --git a/psutil/arch/windows/wmi.c b/psutil/arch/windows/wmi.c
new file mode 100644
index 00000000..5858a9e0
--- /dev/null
+++ b/psutil/arch/windows/wmi.c
@@ -0,0 +1,115 @@
+/*
+ * 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.
+ *
+ * Functions related to the Windows Management Instrumentation API.
+ */
+
+#include <Python.h>
+#include <windows.h>
+#include <pdh.h>
+
+#include "../../_psutil_common.h"
+
+
+// We use an exponentially weighted moving average, just like Unix systems do
+// https://en.wikipedia.org/wiki/Load_(computing)#Unix-style_load_calculation
+//
+// These constants serve as the damping factor and are calculated with
+// 1 / exp(sampling interval in seconds / window size in seconds)
+//
+// This formula comes from linux's include/linux/sched/loadavg.h
+// https://github.com/torvalds/linux/blob/345671ea0f9258f410eb057b9ced9cefbbe5dc78/include/linux/sched/loadavg.h#L20-L23
+#define LOADAVG_FACTOR_1F  0.9200444146293232478931553241
+#define LOADAVG_FACTOR_5F  0.6592406302004437462547604110
+#define LOADAVG_FACTOR_15F 0.2865047968601901003248854266
+// The time interval in seconds between taking load counts, same as Linux
+#define SAMPLING_INTERVAL 5
+
+double load_avg_1m = 0;
+double load_avg_5m = 0;
+double load_avg_15m = 0;
+
+
+VOID CALLBACK LoadAvgCallback(PVOID hCounter) {
+    PDH_FMT_COUNTERVALUE displayValue;
+    double currentLoad;
+    PDH_STATUS err;
+
+    err = PdhGetFormattedCounterValue(
+        (PDH_HCOUNTER)hCounter, PDH_FMT_DOUBLE, 0, &displayValue);
+    // Skip updating the load if we can't get the value successfully
+    if (err != ERROR_SUCCESS) {
+        return;
+    }
+    currentLoad = displayValue.doubleValue;
+
+    load_avg_1m = load_avg_1m * LOADAVG_FACTOR_1F + currentLoad * \
+        (1.0 - LOADAVG_FACTOR_1F);
+    load_avg_5m = load_avg_5m * LOADAVG_FACTOR_5F + currentLoad * \
+        (1.0 - LOADAVG_FACTOR_5F);
+    load_avg_15m = load_avg_15m * LOADAVG_FACTOR_15F + currentLoad * \
+        (1.0 - LOADAVG_FACTOR_15F);
+}
+
+
+PyObject *
+psutil_init_loadavg_counter(PyObject *self, PyObject *args) {
+    WCHAR *szCounterPath = L"\\System\\Processor Queue Length";
+    PDH_STATUS s;
+    BOOL ret;
+    HQUERY hQuery;
+    HCOUNTER hCounter;
+    HANDLE event;
+    HANDLE waitHandle;
+
+    if ((PdhOpenQueryW(NULL, 0, &hQuery)) != ERROR_SUCCESS)
+        goto error;
+
+    s = PdhAddEnglishCounterW(hQuery, szCounterPath, 0, &hCounter);
+    if (s != ERROR_SUCCESS)
+        goto error;
+
+    event = CreateEventW(NULL, FALSE, FALSE, L"LoadUpdateEvent");
+    if (event == NULL) {
+        PyErr_SetFromWindowsErr(GetLastError());
+        return NULL;
+    }
+
+    s = PdhCollectQueryDataEx(hQuery, SAMPLING_INTERVAL, event);
+    if (s != ERROR_SUCCESS)
+        goto error;
+
+    ret = RegisterWaitForSingleObject(
+        &waitHandle,
+        event,
+        (WAITORTIMERCALLBACK)LoadAvgCallback,
+        (PVOID)
+        hCounter,
+        INFINITE,
+        WT_EXECUTEDEFAULT);
+
+    if (ret == 0) {
+        PyErr_SetFromWindowsErr(GetLastError());
+        return NULL;
+    }
+
+    Py_RETURN_NONE;
+
+error:
+    PyErr_SetExcFromWindowsErr(PyExc_OSError, 0);
+    return NULL;
+}
+
+
+/*
+ * Gets the emulated 1 minute, 5 minute and 15 minute load averages 
+ * (processor queue length) for the system.
+ * `init_loadavg_counter` must be called before this function to engage the
+ * mechanism that records load values.
+ */
+PyObject *
+psutil_get_loadavg(PyObject *self, PyObject *args) {
+    return Py_BuildValue("(ddd)", load_avg_1m, load_avg_5m, load_avg_15m);
+}
\ No newline at end of file
diff --git a/psutil/arch/windows/wmi.h b/psutil/arch/windows/wmi.h
new file mode 100644
index 00000000..0210f2d6
--- /dev/null
+++ b/psutil/arch/windows/wmi.h
@@ -0,0 +1,12 @@
+/*
+ * Copyright (c) 2009, Jay Loden, 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.
+ */
+
+
+
+#include <Python.h>
+
+PyObject* psutil_init_loadavg_counter();
+PyObject* psutil_get_loadavg();
diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index 3aebbcbd..79681719 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -159,6 +159,7 @@ HERE = os.path.abspath(os.path.dirname(__file__))
 HAS_CONNECTIONS_UNIX = POSIX and not SUNOS
 HAS_CPU_AFFINITY = hasattr(psutil.Process, "cpu_affinity")
 HAS_CPU_FREQ = hasattr(psutil, "cpu_freq")
+HAS_GETLOADAVG = hasattr(psutil, "getloadavg")
 HAS_ENVIRON = hasattr(psutil.Process, "environ")
 HAS_IONICE = hasattr(psutil.Process, "ionice")
 HAS_MEMORY_MAPS = hasattr(psutil.Process, "memory_maps")
diff --git a/psutil/tests/test_linux.py b/psutil/tests/test_linux.py
index 85cf33c4..e8745a75 100755
--- a/psutil/tests/test_linux.py
+++ b/psutil/tests/test_linux.py
@@ -29,6 +29,7 @@ from psutil._compat import u
 from psutil.tests import call_until
 from psutil.tests import HAS_BATTERY
 from psutil.tests import HAS_CPU_FREQ
+from psutil.tests import HAS_GETLOADAVG
 from psutil.tests import HAS_RLIMIT
 from psutil.tests import MEMORY_TOLERANCE
 from psutil.tests import mock
@@ -866,6 +867,20 @@ class TestSystemCPUStats(unittest.TestCase):
         self.assertAlmostEqual(vmstat_value, psutil_value, delta=500)
 
 
+@unittest.skipIf(not LINUX, "LINUX only")
+class TestLoadAvg(unittest.TestCase):
+
+    @unittest.skipIf(not HAS_GETLOADAVG, "not supported")
+    def test_getloadavg(self):
+        psutil_value = psutil.getloadavg()
+        with open("/proc/loadavg", "r") as f:
+            proc_value = f.read().split()
+
+        self.assertAlmostEqual(float(proc_value[0]), psutil_value[0], delta=1)
+        self.assertAlmostEqual(float(proc_value[1]), psutil_value[1], delta=1)
+        self.assertAlmostEqual(float(proc_value[2]), psutil_value[2], delta=1)
+
+
 # =====================================================================
 # --- system network
 # =====================================================================
diff --git a/psutil/tests/test_memory_leaks.py b/psutil/tests/test_memory_leaks.py
index 170a1d0c..dde50a57 100755
--- a/psutil/tests/test_memory_leaks.py
+++ b/psutil/tests/test_memory_leaks.py
@@ -36,6 +36,7 @@ from psutil.tests import create_sockets
 from psutil.tests import get_test_subprocess
 from psutil.tests import HAS_CPU_AFFINITY
 from psutil.tests import HAS_CPU_FREQ
+from psutil.tests import HAS_GETLOADAVG
 from psutil.tests import HAS_ENVIRON
 from psutil.tests import HAS_IONICE
 from psutil.tests import HAS_MEMORY_MAPS
@@ -474,6 +475,10 @@ class TestModuleFunctionsLeaks(TestMemLeak):
     def test_cpu_freq(self):
         self.execute(psutil.cpu_freq)
 
+    @unittest.skipIf(not HAS_GETLOADAVG, "not supported")
+    def test_getloadavg(self):
+        self.execute(psutil.getloadavg)
+
     # --- mem
 
     def test_virtual_memory(self):
diff --git a/psutil/tests/test_system.py b/psutil/tests/test_system.py
index 75e65b0f..c4c503a2 100755
--- a/psutil/tests/test_system.py
+++ b/psutil/tests/test_system.py
@@ -38,6 +38,7 @@ from psutil.tests import enum
 from psutil.tests import get_test_subprocess
 from psutil.tests import HAS_BATTERY
 from psutil.tests import HAS_CPU_FREQ
+from psutil.tests import HAS_GETLOADAVG
 from psutil.tests import HAS_NET_IO_COUNTERS
 from psutil.tests import HAS_SENSORS_BATTERY
 from psutil.tests import HAS_SENSORS_FANS
@@ -778,6 +779,15 @@ class TestSystemAPIs(unittest.TestCase):
         if LINUX:
             self.assertEqual(len(ls), psutil.cpu_count())
 
+    @unittest.skipIf(not HAS_GETLOADAVG, "not supported")
+    def test_getloadavg(self):
+        loadavg = psutil.getloadavg()
+        assert len(loadavg) == 3
+
+        for load in loadavg:
+            self.assertIsInstance(load, float)
+            self.assertGreaterEqual(load, 0.0)
+
     def test_os_constants(self):
         names = ["POSIX", "WINDOWS", "LINUX", "MACOS", "FREEBSD", "OPENBSD",
                  "NETBSD", "BSD", "SUNOS"]
diff --git a/setup.py b/setup.py
index 465a9b9e..2c3d9b36 100755
--- a/setup.py
+++ b/setup.py
@@ -135,12 +135,12 @@ if WINDOWS:
             'psutil/arch/windows/inet_ntop.c',
             'psutil/arch/windows/services.c',
             'psutil/arch/windows/global.c',
-            # 'psutil/arch/windows/connections.c',
+            'psutil/arch/windows/wmi.c',
         ],
         define_macros=macros,
         libraries=[
             "psapi", "kernel32", "advapi32", "shell32", "netapi32",
-            "wtsapi32", "ws2_32", "PowrProf",
+            "wtsapi32", "ws2_32", "PowrProf", "pdh",
         ],
         # extra_compile_args=["/Z7"],
         # extra_link_args=["/DEBUG"]