gh-158552: Wait for Windows threads to suspend before blocking sampling (#158802)

* gh-158552: Wait for Windows threads to suspend before blocking sampling

* Use a named Windows thread enumeration status constant

Author: pablogsal

Lib/test/test_profiling/test_sampling_profiler/test_blocking.py

@@ -5,6 +5,7 @@
 import subprocess
 import sys
 import textwrap
+import time
 import unittest
 from unittest import mock
 
@@ -213,3 +214,59 @@ def test_run_blocking_exits_after_target_process_exits(self):
                 replay.returncode, 0,
                 f"stdout:\n{replay.stdout}\nstderr:\n{replay.stderr}",
             )
+
+
+@requires_remote_subprocess_debugging()
+@unittest.skipUnless(sys.platform == "win32", "Windows only")
+class TestBlockingModeSuspension(unittest.TestCase):
+    def test_all_threads_stop_before_pause_returns(self):
+        import mmap
+
+        tag = f"cpython_blocking_{os.getpid()}_{id(self)}"
+        script = textwrap.dedent(f'''
+            import mmap
+            import struct
+            import threading
+
+            memory = mmap.mmap(-1, 16, tagname={tag!r})
+
+            def worker(offset):
+                counter = 0
+                while True:
+                    counter += 1
+                    struct.pack_into("q", memory, offset, counter)
+
+            for offset in (0, 8):
+                threading.Thread(target=worker, args=(offset,), daemon=True).start()
+            _test_sock.sendall(b"working")
+            _test_sock.recv(1)
+        ''')
+        with mmap.mmap(-1, 16, tagname=tag) as memory:
+            with test_subprocess(script, wait_for_working=True) as subproc:
+                unwinder = _remote_debugging.RemoteUnwinder(
+                    subproc.process.pid, all_threads=True)
+                deadline = time.monotonic() + SHORT_TIMEOUT
+                while not all(memory[offset:offset + 8] != bytes(8)
+                              for offset in (0, 8)):
+                    self.assertLess(time.monotonic(), deadline,
+                                    "Worker threads did not start")
+                    time.sleep(0.001)
+                for _ in range(100):
+                    self.assertTrue(unwinder.pause_threads())
+                    try:
+                        before = memory[:]
+                        self.assertFalse(unwinder.pause_threads())
+                        unwinder.get_stack_trace()
+                        time.sleep(0.001)
+                        self.assertEqual(memory[:], before,
+                                         "Target memory changed while paused")
+                    finally:
+                        unwinder.resume_threads()
+                    self.assertFalse(unwinder.resume_threads())
+                before = memory[:]
+                deadline = time.monotonic() + SHORT_TIMEOUT
+                while any(memory[offset:offset + 8] == before[offset:offset + 8]
+                          for offset in (0, 8)):
+                    self.assertLess(time.monotonic(), deadline,
+                                    "Worker threads did not resume")
+                    time.sleep(0.001)

Misc/NEWS.d/next/Library/2026-10-04-20-00-00.gh-issue-158552.windows.rst

@@ -0,0 +1,3 @@
+Wait for Windows threads to finish suspending before reading target memory
+in the sampling profiler's blocking mode. Resume the target if waiting for
+suspension fails.

Modules/_remote_debugging/_remote_debugging.h

@@ -105,6 +105,7 @@ typedef struct {
 #ifdef MS_WINDOWS
 #define STATUS_SUCCESS                   ((NTSTATUS)0x00000000L)
 #define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)
+#define STATUS_NO_MORE_ENTRIES           ((NTSTATUS)0x8000001AL)
 typedef enum _WIN32_THREADSTATE {
     WIN32_THREADSTATE_INITIALIZED = 0,
     WIN32_THREADSTATE_READY       = 1,

Modules/_remote_debugging/threads.c

@@ -828,6 +828,89 @@ _Py_RemoteDebug_ResumeAllThreads(RemoteUnwinderObject *unwinder, _Py_RemoteDebug
 
 #elif defined(MS_WINDOWS)
 
+static int
+wait_for_threads_to_stop(RemoteUnwinderObject *unwinder)
+{
+    typedef NTSTATUS (NTAPI *NtGetNextThreadFunc)(
+        HANDLE, HANDLE, ACCESS_MASK, ULONG, ULONG, PHANDLE);
+    static NtGetNextThreadFunc pNtGetNextThread = NULL;
+    static int tried_load = 0;
+
+    if (!tried_load) {
+        HMODULE hNtdll = GetModuleHandleW(L"ntdll.dll");
+        if (hNtdll) {
+            pNtGetNextThread = (NtGetNextThreadFunc)GetProcAddress(
+                hNtdll, "NtGetNextThread");
+        }
+        tried_load = 1;
+    }
+    if (pNtGetNextThread == NULL) {
+        PyErr_SetString(PyExc_RuntimeError, "NtGetNextThread not available");
+        return -1;
+    }
+
+    HANDLE previous = NULL;
+    int result = -1;
+    for (;;) {
+        HANDLE next = NULL;
+        // Enumerate with the available access, then obtain context access
+        // separately so threads that deny it are not silently skipped.
+        NTSTATUS status = pNtGetNextThread(unwinder->handle.hProcess,
+                                           previous, MAXIMUM_ALLOWED, 0, 0, &next);
+        if (previous != NULL) {
+            CloseHandle(previous);
+        }
+        previous = next;
+        if (status == STATUS_NO_MORE_ENTRIES) {
+            break;
+        }
+        if (status < 0) {
+            if (!is_process_alive(unwinder->handle.hProcess)) {
+                PyErr_Format(PyExc_ProcessLookupError,
+                             "Process %d has terminated", unwinder->handle.pid);
+            }
+            else {
+                PyErr_Format(PyExc_RuntimeError,
+                             "NtGetNextThread failed: 0x%lx", status);
+            }
+            goto done;
+        }
+        HANDLE thread;
+        if (!DuplicateHandle(GetCurrentProcess(), next, GetCurrentProcess(),
+                             &thread, THREAD_GET_CONTEXT | SYNCHRONIZE,
+                             FALSE, 0)) {
+            PyErr_SetFromWindowsErr(GetLastError());
+            goto done;
+        }
+        // Suspension is asynchronous. Reading the context waits for the
+        // thread to stop before we start reading the target's memory.
+        CONTEXT context = {.ContextFlags = CONTEXT_CONTROL};
+        if (!GetThreadContext(thread, &context)) {
+            DWORD error = GetLastError();
+            int exited = WaitForSingleObject(thread, 0) == WAIT_OBJECT_0;
+            CloseHandle(thread);
+            if (exited) {
+                continue;
+            }
+            PyErr_SetFromWindowsErr(error);
+            goto done;
+        }
+        CloseHandle(thread);
+    }
+    if (!is_process_alive(unwinder->handle.hProcess)) {
+        PyErr_Format(PyExc_ProcessLookupError,
+                     "Process %d has terminated", unwinder->handle.pid);
+        goto done;
+    }
+    result = 0;
+
+done:
+    if (previous != NULL) {
+        CloseHandle(previous);
+    }
+    return result;
+}
+
 void
 _Py_RemoteDebug_InitThreadsState(RemoteUnwinderObject *unwinder, _Py_RemoteDebug_ThreadsState *st)
 {
@@ -858,6 +941,12 @@ _Py_RemoteDebug_StopAllThreads(RemoteUnwinderObject *unwinder, _Py_RemoteDebug_T
     if (status >= 0) {
         st->hProcess = unwinder->handle.hProcess;
         st->suspended = 1;
+        if (wait_for_threads_to_stop(unwinder) < 0) {
+            // pause_threads() has not yet set threads_stopped, so its caller
+            // will not resume the process when we return an error.
+            _Py_RemoteDebug_ResumeAllThreads(unwinder, st);
+            return -1;
+        }
         _Py_RemoteDebug_ClearCache(&unwinder->handle);
         return 0;
     }