gh-150114: Log the memory usage in regrtest on Windows

Author: vstinner

Lib/test/libregrtest/utils.py

@@ -13,6 +13,10 @@
 import tempfile
 import textwrap
 from collections.abc import Callable
+try:
+    from _testcapi import get_process_memory_usage as _get_process_memory_usage
+except AttributeError:
+    _get_process_memory_usage = None
 
 from test import support
 from test.support import os_helper
@@ -756,8 +760,13 @@ def display_title(title):
 
 def get_process_memory_usage(pid: int) -> int | None:
     """
-    Read the private memory in bytes from /proc/pid/smaps.
+    Get process memory usage in bytes.
     """
+    if _get_process_memory_usage is not None:
+        return _get_process_memory_usage(pid)
+
+    # Linux implementation: read the private memory in bytes from
+    # /proc/pid/smaps.
     try:
         fp = open(f"/proc/{pid}/smaps", "rb")
     except OSError:

Modules/_testcapi/mem.c

@@ -1,5 +1,9 @@
 #include "parts.h"
 
+#ifdef MS_WINDOWS
+#  include <windows.h>
+#  include <psapi.h>              // GetProcessMemoryInfo()
+#endif
 #include <stddef.h>
 
 
@@ -684,6 +688,34 @@ tracemalloc_track_race(PyObject *self, PyObject *args)
 }
 
 
+#ifdef MS_WINDOWS
+// Get process memory usage in bytes.
+static PyObject *
+get_process_memory_usage(PyObject *self, PyObject *args)
+{
+    int pid;
+    if (!PyArg_ParseTuple(args, "i", &pid)) {
+        return NULL;
+    }
+
+    HANDLE handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, (DWORD)pid);
+    if (handle == NULL) {
+        return PyErr_SetFromWindowsErr(0);
+    }
+
+    PROCESS_MEMORY_COUNTERS pmc;
+    if (!GetProcessMemoryInfo(handle, &pmc, sizeof(pmc))) {
+        CloseHandle(handle);
+        return PyErr_SetFromWindowsErr(0);
+    }
+    CloseHandle(handle);
+
+    size_t size = (pmc.WorkingSetSize + pmc.PagefileUsage);
+    return PyLong_FromSize_t(size);
+}
+#endif
+
+
 static PyMethodDef test_methods[] = {
     {"pymem_api_misuse",              pymem_api_misuse,              METH_NOARGS},
     {"pymem_buffer_overflow",         pymem_buffer_overflow,         METH_NOARGS},
@@ -698,6 +730,9 @@ static PyMethodDef test_methods[] = {
     {"test_pymem_setrawallocators",   test_pymem_setrawallocators,   METH_NOARGS},
     {"test_pyobject_new",             test_pyobject_new,             METH_NOARGS},
     {"test_pyobject_setallocators",   test_pyobject_setallocators,   METH_NOARGS},
+#ifdef MS_WINDOWS
+    {"get_process_memory_usage",      get_process_memory_usage,      METH_VARARGS},
+#endif
 
     // Tracemalloc tests
     {"tracemalloc_track",             tracemalloc_track,             METH_VARARGS},