expose load counter code on winwin-getloadavg

5 files changed, 131 insertions, 2 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index 5d2b8d3c..dc1a2c75 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -2017,6 +2017,20 @@ if hasattr(_psplatform, "cpu_freq"):
     __all__.append("cpu_freq")
 
 
+_loadcounter_init = False
+
+
+def getloadavg():
+    global _loadcounter_init
+    if POSIX:
+        return os.getloadavg()
+    else:
+        if not _loadcounter_init:
+            _psplatform.cext.init_loadcounter()
+            _loadcounter_init = True
+        return _psplatform.cext.get_loadcounter()
+
+
 # =====================================================================
 # --- system memory related functions
 # =====================================================================
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 53617c58..aa181781 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"
 
 
@@ -3512,6 +3513,10 @@ PsutilMethods[] = {
      "Return battery metrics usage."},
     {"getpagesize", psutil_getpagesize, METH_VARARGS,
      "Return system memory page size."},
+    {"init_loadcounter", psutil_init_loadcounter, METH_VARARGS,
+     "XXX"},
+    {"get_loadcounter", psutil_get_loadcounter, METH_VARARGS,
+     "XXX"},
 
     // --- windows services
     {"winservice_enumerate", psutil_winservice_enumerate, METH_VARARGS,
diff --git a/psutil/arch/windows/wmi.c b/psutil/arch/windows/wmi.c
new file mode 100644
index 00000000..18f0d5c7
--- /dev/null
+++ b/psutil/arch/windows/wmi.c
@@ -0,0 +1,97 @@
+/*
+ * 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.
+ */
+
+#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
+// This constant serves as the damping factor.
+#define LOADAVG_FACTOR_1F 0.9200444146293232478931553241
+// The time interval in seconds between taking load counts
+#define SAMPLING_INTERVAL 5
+
+
+double load_1m = 0;
+
+
+VOID CALLBACK LoadCallback(PVOID hCounter) {
+    PDH_FMT_COUNTERVALUE displayValue;
+    double currentLoad = displayValue.doubleValue;
+    double newLoad;
+
+    PdhGetFormattedCounterValue(
+        (PDH_HCOUNTER)hCounter, PDH_FMT_DOUBLE, 0, &displayValue);
+
+    newLoad = load_1m * LOADAVG_FACTOR_1F + currentLoad * \
+        (1.0 - LOADAVG_FACTOR_1F);
+    load_1m = newLoad;
+}
+
+
+PyObject *
+psutil_init_loadcounter(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)LoadCallback,
+        (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 1 minute load average (processor queue length) for the system.
+ * InitializeLoadCounter must be called before this function to engage the
+ * mechanism that records load values.
+ */
+PyObject *
+psutil_get_loadcounter(PyObject *self, PyObject *args) {
+    PyObject* load = PyFloat_FromDouble(load_1m);
+    return load;
+}
diff --git a/psutil/arch/windows/wmi.h b/psutil/arch/windows/wmi.h
new file mode 100644
index 00000000..db574a4f
--- /dev/null
+++ b/psutil/arch/windows/wmi.h
@@ -0,0 +1,13 @@
+/*
+ * 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_loadcounter();
+PyObject* psutil_get_loadcounter();
+
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"]