gh-138122: Make sampling profiler integration tests more resilient

pablogsal · pablogsal · commit 56f0d937d18d · 2025-12-07T17:09:28.000Z
The tests were flaky on slow machines because subprocesses could finish
before enough samples were collected. This adds synchronization similar
to test_external_inspection: test scripts now signal when they start
working, and the profiler waits for this signal before sampling.

Test scripts now run in infinite loops until killed rather than for
fixed iterations, ensuring the profiler always has active work to
sample regardless of machine speed.
diff --git a/Lib/test/test_profiling/test_sampling_profiler/helpers.py b/Lib/test/test_profiling/test_sampling_profiler/helpers.py
@@ -4,6 +4,7 @@
 import socket
 import subprocess
 import sys
+import time
 import unittest
 from collections import namedtuple
 
@@ -14,6 +15,9 @@
 
 PROCESS_VM_READV_SUPPORTED = False
 
+# Maximum number of retry attempts for operations that may fail transiently
+MAX_TRIES = 10
+
 try:
     from _remote_debugging import PROCESS_VM_READV_SUPPORTED  # noqa: F401
     import _remote_debugging  # noqa: F401
@@ -38,12 +42,66 @@
 SubprocessInfo = namedtuple("SubprocessInfo", ["process", "socket"])
 
 
+def _wait_for_signal(sock, expected_signals, timeout=SHORT_TIMEOUT):
+    """
+    Wait for expected signal(s) from a socket with proper timeout and EOF handling.
+
+    Args:
+        sock: Connected socket to read from
+        expected_signals: Single bytes object or list of bytes objects to wait for
+        timeout: Socket timeout in seconds
+
+    Returns:
+        bytes: Complete accumulated response buffer
+
+    Raises:
+        RuntimeError: If connection closed before signal received or timeout
+    """
+    if isinstance(expected_signals, bytes):
+        expected_signals = [expected_signals]
+
+    sock.settimeout(timeout)
+    buffer = b""
+
+    while True:
+        # Check if all expected signals are in buffer
+        if all(sig in buffer for sig in expected_signals):
+            return buffer
+
+        try:
+            chunk = sock.recv(4096)
+            if not chunk:
+                # EOF - connection closed
+                raise RuntimeError(
+                    f"Connection closed before receiving expected signals. "
+                    f"Expected: {expected_signals}, Got: {buffer[-200:]!r}"
+                )
+            buffer += chunk
+        except socket.timeout:
+            raise RuntimeError(
+                f"Timeout waiting for signals. "
+                f"Expected: {expected_signals}, Got: {buffer[-200:]!r}"
+            )
+
+
+def _cleanup_sockets(*sockets):
+    """Safely close multiple sockets, ignoring errors."""
+    for sock in sockets:
+        if sock is not None:
+            try:
+                sock.close()
+            except OSError:
+                pass
+
+
 @contextlib.contextmanager
-def test_subprocess(script):
+def test_subprocess(script, wait_for_working=True):
     """Context manager to create a test subprocess with socket synchronization.
 
     Args:
-        script: Python code to execute in the subprocess
+        script: Python code to execute in the subprocess. Should send b"working"
+                to signal when work has started.
+        wait_for_working: If True, wait for both "ready" and "working" signals
 
     Yields:
         SubprocessInfo: Named tuple with process and socket objects
@@ -80,19 +138,27 @@ def test_subprocess(script):
         # Wait for process to connect and send ready signal
         client_socket, _ = server_socket.accept()
         server_socket.close()
-        response = client_socket.recv(1024)
-        if response != b"ready":
-            raise RuntimeError(
-                f"Unexpected response from subprocess: {response!r}"
-            )
+        server_socket = None
+
+        # Wait for ready signal, and optionally working signal
+        if wait_for_working:
+            _wait_for_signal(client_socket, [b"ready", b"working"])
+        else:
+            _wait_for_signal(client_socket, b"ready")
 
         yield SubprocessInfo(proc, client_socket)
     finally:
-        if client_socket is not None:
-            client_socket.close()
+        _cleanup_sockets(client_socket, server_socket)
         if proc.poll() is None:
-            proc.kill()
-        proc.wait()
+            proc.terminate()
+            try:
+                proc.wait(timeout=SHORT_TIMEOUT)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                try:
+                    proc.wait(timeout=SHORT_TIMEOUT)
+                except subprocess.TimeoutExpired:
+                    pass  # Process refuses to die
 
 
 def close_and_unlink(file):
diff --git a/Lib/test/test_profiling/test_sampling_profiler/test_integration.py b/Lib/test/test_profiling/test_sampling_profiler/test_integration.py
@@ -33,12 +33,16 @@
     close_and_unlink,
     skip_if_not_supported,
     PROCESS_VM_READV_SUPPORTED,
+    MAX_TRIES,
 )
 from .mocks import MockFrameInfo, MockThreadInfo, MockInterpreterInfo
 
 # Duration for profiling tests - long enough for process to complete naturally
 PROFILING_TIMEOUT = str(int(SHORT_TIMEOUT))
 
+# Duration for profiling in tests - short enough to complete quickly
+PROFILING_DURATION_SEC = 2
+
 
 @skip_if_not_supported
 @unittest.skipIf(
@@ -364,9 +368,44 @@ def total_occurrences(func):
 class TestSampleProfilerIntegration(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
+        # Test script that runs in an infinite loop until killed.
+        # Sends "working" signal after starting work to synchronize with profiler.
+        # Used with test_subprocess() which injects socket code.
         cls.test_script = '''
+def slow_fibonacci(n):
+    """Recursive fibonacci - should show up prominently in profiler."""
+    if n <= 1:
+        return n
+    return slow_fibonacci(n-1) + slow_fibonacci(n-2)
+
+def cpu_intensive_work():
+    """CPU intensive work that should show in profiler."""
+    result = 0
+    for i in range(10000):
+        result += i * i
+        if i % 100 == 0:
+            result = result % 1000000
+    return result
+
+def main_loop():
+    """Main test loop - runs indefinitely until killed."""
+    # Signal that we're starting work
+    _test_sock.sendall(b"working")
+
+    iteration = 0
+    while True:
+        if iteration % 2 == 0:
+            result = slow_fibonacci(15)
+        else:
+            result = cpu_intensive_work()
+        iteration += 1
+
+if __name__ == "__main__":
+    main_loop()
+'''
+        # CLI test script - runs for a fixed duration (used for CLI tests without socket sync)
+        cls.cli_test_script = '''
 import time
-import os
 
 def slow_fibonacci(n):
     """Recursive fibonacci - should show up prominently in profiler."""
@@ -384,14 +423,15 @@ def cpu_intensive_work():
     return result
 
 def main_loop():
-    """Main test loop."""
-    max_iterations = 200
-
-    for iteration in range(max_iterations):
+    """Main test loop - runs for a fixed duration."""
+    end_time = time.time() + 30  # Run for 30 seconds max
+    iteration = 0
+    while time.time() < end_time:
         if iteration % 2 == 0:
             result = slow_fibonacci(15)
         else:
             result = cpu_intensive_work()
+        iteration += 1
 
 if __name__ == "__main__":
     main_loop()
@@ -404,12 +444,11 @@ def test_sampling_basic_functionality(self):
             mock.patch("sys.stdout", captured_output),
         ):
             try:
-                # Sample for up to SHORT_TIMEOUT seconds, but process exits after fixed iterations
                 collector = PstatsCollector(sample_interval_usec=1000, skip_idle=False)
                 profiling.sampling.sample.sample(
                     subproc.process.pid,
                     collector,
-                    duration_sec=SHORT_TIMEOUT,
+                    duration_sec=PROFILING_DURATION_SEC,
                 )
                 collector.print_stats(show_summary=False)
             except PermissionError:
@@ -442,7 +481,7 @@ def test_sampling_with_pstats_export(self):
                     profiling.sampling.sample.sample(
                         subproc.process.pid,
                         collector,
-                        duration_sec=1,
+                        duration_sec=PROFILING_DURATION_SEC,
                     )
                     collector.export(pstats_out.name)
                 except PermissionError:
@@ -488,7 +527,7 @@ def test_sampling_with_collapsed_export(self):
                     profiling.sampling.sample.sample(
                         subproc.process.pid,
                         collector,
-                        duration_sec=1,
+                        duration_sec=PROFILING_DURATION_SEC,
                     )
                     collector.export(collapsed_file.name)
                 except PermissionError:
@@ -536,7 +575,7 @@ def test_sampling_all_threads(self):
                 profiling.sampling.sample.sample(
                     subproc.process.pid,
                     collector,
-                    duration_sec=1,
+                    duration_sec=PROFILING_DURATION_SEC,
                     all_threads=True,
                 )
                 collector.print_stats(show_summary=False)
@@ -548,12 +587,16 @@ def test_sampling_all_threads(self):
 
     def test_sample_target_script(self):
         script_file = tempfile.NamedTemporaryFile(delete=False)
-        script_file.write(self.test_script.encode("utf-8"))
+        script_file.write(self.cli_test_script.encode("utf-8"))
         script_file.flush()
         self.addCleanup(close_and_unlink, script_file)
 
-        # Sample for up to SHORT_TIMEOUT seconds, but process exits after fixed iterations
-        test_args = ["profiling.sampling.sample", "run", "-d", PROFILING_TIMEOUT, script_file.name]
+        # Sample for PROFILING_DURATION_SEC seconds
+        test_args = [
+            "profiling.sampling.sample", "run",
+            "-d", str(PROFILING_DURATION_SEC),
+            script_file.name
+        ]
 
         with (
             mock.patch("sys.argv", test_args),
@@ -583,13 +626,13 @@ def test_sample_target_module(self):
         module_path = os.path.join(tempdir.name, "test_module.py")
 
         with open(module_path, "w") as f:
-            f.write(self.test_script)
+            f.write(self.cli_test_script)
 
         test_args = [
             "profiling.sampling.cli",
             "run",
             "-d",
-            PROFILING_TIMEOUT,
+            str(PROFILING_DURATION_SEC),
             "-m",
             "test_module",
         ]
@@ -630,8 +673,10 @@ def test_invalid_pid(self):
             profiling.sampling.sample.sample(-1, collector, duration_sec=1)
 
     def test_process_dies_during_sampling(self):
+        # Use wait_for_working=False since this simple script doesn't send "working"
         with test_subprocess(
-            "import time; time.sleep(0.5); exit()"
+            "import time; time.sleep(0.5); exit()",
+            wait_for_working=False
         ) as subproc:
             with (
                 io.StringIO() as captured_output,
@@ -654,7 +699,11 @@ def test_process_dies_during_sampling(self):
             self.assertIn("Error rate", output)
 
     def test_is_process_running(self):
-        with test_subprocess("import time; time.sleep(1000)") as subproc:
+        # Use wait_for_working=False since this simple script doesn't send "working"
+        with test_subprocess(
+            "import time; time.sleep(1000)",
+            wait_for_working=False
+        ) as subproc:
             try:
                 profiler = SampleProfiler(
                     pid=subproc.process.pid,
@@ -681,7 +730,11 @@ def test_is_process_running(self):
 
     @unittest.skipUnless(sys.platform == "linux", "Only valid on Linux")
     def test_esrch_signal_handling(self):
-        with test_subprocess("import time; time.sleep(1000)") as subproc:
+        # Use wait_for_working=False since this simple script doesn't send "working"
+        with test_subprocess(
+            "import time; time.sleep(1000)",
+            wait_for_working=False
+        ) as subproc:
             try:
                 unwinder = _remote_debugging.RemoteUnwinder(
                     subproc.process.pid
@@ -793,25 +846,29 @@ class TestAsyncAwareProfilingIntegration(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
+        # Async test script that runs indefinitely until killed.
+        # Sends "working" signal after starting work to synchronize with profiler.
         cls.async_script = '''
 import asyncio
 
 async def sleeping_leaf():
     """Leaf task that just sleeps - visible in 'all' mode."""
-    for _ in range(50):
+    while True:
         await asyncio.sleep(0.02)
 
 async def cpu_leaf():
     """Leaf task that does CPU work - visible in both modes."""
     total = 0
-    for _ in range(200):
+    while True:
         for i in range(10000):
             total += i * i
         await asyncio.sleep(0)
-    return total
 
 async def supervisor():
     """Middle layer that spawns leaf tasks."""
+    # Signal that we're starting work
+    _test_sock.sendall(b"working")
+
     tasks = [
         asyncio.create_task(sleeping_leaf(), name="Sleeper-0"),
         asyncio.create_task(sleeping_leaf(), name="Sleeper-1"),
@@ -838,7 +895,7 @@ def _collect_async_samples(self, async_aware_mode):
                 profiling.sampling.sample.sample(
                     subproc.process.pid,
                     collector,
-                    duration_sec=SHORT_TIMEOUT,
+                    duration_sec=PROFILING_DURATION_SEC,
                     async_aware=async_aware_mode,
                 )
             except PermissionError:
