gh-150114: Log the memory usage in regrtest on Windows (#150267)

Add _winapi.GetProcessMemoryInfo() function.

Co-authored-by: Cody Maloney <cmaloney@users.noreply.github.com>

diff --git a/Lib/test/libregrtest/utils.py b/Lib/test/libregrtest/utils.py
index 1b4cb96406d6f604375a1da8b64885299f6a3418..21b84f7555b7713475f8ff63f9968b107d277428 100644 (file)
--- a/Lib/test/libregrtest/utils.py
+++ b/Lib/test/libregrtest/utils.py
@@ -12,7 +12,13 @@ import sys
 import sysconfig
 import tempfile
 import textwrap
+import types
 from collections.abc import Callable
+_winapi: types.ModuleType | None
+try:
+    import _winapi
+except ImportError:
+    _winapi = None
 
 from test import support
 from test.support import os_helper
@@ -754,10 +760,9 @@ def display_title(title):
     print(flush=True)
 
 
-def get_process_memory_usage(pid: int) -> int | None:
-    """
-    Read the private memory in bytes from /proc/pid/smaps.
-    """
+def _get_process_memory_usage_linux(pid: int) -> int | None:
+    # Linux implementation: read the private memory in bytes from
+    # /proc/pid/smaps.
     try:
         fp = open(f"/proc/{pid}/smaps", "rb")
     except OSError:
@@ -775,3 +780,26 @@ def get_process_memory_usage(pid: int) -> int | None:
         return total
     except ProcessLookupError:
         return None
+
+
+def _get_process_memory_usage_windows(pid: int) -> int | None:
+    assert _winapi is not None  # to make mypy happy
+    handle = _winapi.OpenProcess(_winapi.PROCESS_QUERY_LIMITED_INFORMATION,
+                                 False, pid)
+    try:
+        mem_info = _winapi.GetProcessMemoryInfo(handle)
+    finally:
+        _winapi.CloseHandle(handle)
+    return mem_info['WorkingSetSize']
+
+
+if _winapi is not None:
+    get_process_memory_usage = _get_process_memory_usage_windows
+elif sys.platform == 'linux':
+    get_process_memory_usage = _get_process_memory_usage_linux
+else:
+    def get_process_memory_usage(pid: int) -> int | None:
+        """
+        Get process memory usage in bytes.
+        """
+        return None

diff --git a/Modules/_winapi.c b/Modules/_winapi.c
index ffa407b2f21f733b8a7db7cd4ef14ced50530ec5..fc2c0890468a6b96ba43212299836491eab60b77 100644 (file)
--- a/Modules/_winapi.c
+++ b/Modules/_winapi.c
@@ -49,6 +49,13 @@
 #include <crtdbg.h>
 #include "winreparse.h"
 
+// PSAPI_VERSION=2 redirects GetProcessMemoryInfo() to
+// K32GetProcessMemoryInfo() in kernel32.dll, so we don't need to link
+// psapi.lib. See:
+// https://learn.microsoft.com/windows/win32/api/psapi/nf-psapi-getprocessmemoryinfo
+#define PSAPI_VERSION 2
+#include <psapi.h>                // GetProcessMemoryInfo()
+
 #if defined(MS_WIN32) && !defined(MS_WIN64)
 #define HANDLE_TO_PYNUM(handle) \
     PyLong_FromUnsignedLong((unsigned long) handle)
@@ -3080,6 +3087,61 @@ _winapi_ReportEvent_impl(PyObject *module, HANDLE handle,
 }
 
 
+/*[clinic input]
+_winapi.GetProcessMemoryInfo
+    handle: HANDLE
+    /
+
+Return the memory usage of the given process handle as a dict.
+[clinic start generated code]*/
+
+static PyObject *
+_winapi_GetProcessMemoryInfo_impl(PyObject *module, HANDLE handle)
+/*[clinic end generated code: output=00a5d09732e84120 input=5b90ad61cdc68d2a]*/
+{
+    PROCESS_MEMORY_COUNTERS pmc;
+    if (!GetProcessMemoryInfo(handle, &pmc, sizeof(pmc))) {
+        return PyErr_SetFromWindowsErr(0);
+    }
+
+    PyObject *result = PyDict_New();
+    if (result == NULL) {
+        return NULL;
+    }
+
+#define ADD(ATTR) \
+    do { \
+        PyObject *obj = PyLong_FromSize_t(pmc.ATTR); \
+        if (obj == NULL) { \
+            goto error; \
+        } \
+        if (PyDict_SetItemString(result, #ATTR, obj) < 0) { \
+            Py_DECREF(obj); \
+            goto error; \
+        } \
+        Py_DECREF(obj); \
+    } while (0)
+
+    ADD(PageFaultCount);
+    ADD(PeakWorkingSetSize);
+    ADD(WorkingSetSize);
+    ADD(QuotaPeakPagedPoolUsage);
+    ADD(QuotaPagedPoolUsage);
+    ADD(QuotaPeakNonPagedPoolUsage);
+    ADD(QuotaNonPagedPoolUsage);
+    ADD(PagefileUsage);
+    ADD(PeakPagefileUsage);
+
+#undef ADD
+
+    return result;
+
+error:
+    Py_DECREF(result);
+    return NULL;
+}
+
+
 static PyMethodDef winapi_functions[] = {
     _WINAPI_CLOSEHANDLE_METHODDEF
     _WINAPI_CONNECTNAMEDPIPE_METHODDEF
@@ -3130,6 +3192,7 @@ static PyMethodDef winapi_functions[] = {
     _WINAPI__MIMETYPES_READ_WINDOWS_REGISTRY_METHODDEF
     _WINAPI_NEEDCURRENTDIRECTORYFOREXEPATH_METHODDEF
     _WINAPI_COPYFILE2_METHODDEF
+    _WINAPI_GETPROCESSMEMORYINFO_METHODDEF
     {NULL, NULL}
 };
 
@@ -3226,6 +3289,7 @@ static int winapi_exec(PyObject *m)
     WINAPI_CONSTANT(F_DWORD, PROCESS_ALL_ACCESS);
     WINAPI_CONSTANT(F_DWORD, SYNCHRONIZE);
     WINAPI_CONSTANT(F_DWORD, PROCESS_DUP_HANDLE);
+    WINAPI_CONSTANT(F_DWORD, PROCESS_QUERY_LIMITED_INFORMATION);
     WINAPI_CONSTANT(F_DWORD, SEC_COMMIT);
     WINAPI_CONSTANT(F_DWORD, SEC_IMAGE);
     WINAPI_CONSTANT(F_DWORD, SEC_LARGE_PAGES);

diff --git a/Modules/clinic/_winapi.c.h b/Modules/clinic/_winapi.c.h
index 00cce91dca43b1c302cf22337a50a4f6a1dc87df..dd9dbffaa9ac23c55c64727d0517e66ca308fe98 100644 (file)
--- a/Modules/clinic/_winapi.c.h
+++ b/Modules/clinic/_winapi.c.h
@@ -2331,7 +2331,34 @@ exit:
     return return_value;
 }
 
+PyDoc_STRVAR(_winapi_GetProcessMemoryInfo__doc__,
+"GetProcessMemoryInfo($module, handle, /)\n"
+"--\n"
+"\n"
+"Return the memory usage of the given process handle as a dict.");
+
+#define _WINAPI_GETPROCESSMEMORYINFO_METHODDEF    \
+    {"GetProcessMemoryInfo", (PyCFunction)_winapi_GetProcessMemoryInfo, METH_O, _winapi_GetProcessMemoryInfo__doc__},
+
+static PyObject *
+_winapi_GetProcessMemoryInfo_impl(PyObject *module, HANDLE handle);
+
+static PyObject *
+_winapi_GetProcessMemoryInfo(PyObject *module, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    HANDLE handle;
+
+    if (!PyArg_Parse(arg, "" F_HANDLE ":GetProcessMemoryInfo", &handle)) {
+        goto exit;
+    }
+    return_value = _winapi_GetProcessMemoryInfo_impl(module, handle);
+
+exit:
+    return return_value;
+}
+
 #ifndef _WINAPI_GETSHORTPATHNAME_METHODDEF
     #define _WINAPI_GETSHORTPATHNAME_METHODDEF
 #endif /* !defined(_WINAPI_GETSHORTPATHNAME_METHODDEF) */
-/*[clinic end generated code: output=4ab94eaee93a0a90 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=07dfd4bbacaed4a8 input=a9049054013a1b77]*/