diff --git a/Cargo.lock b/Cargo.lock index 51b77ecb5..4cbb8189b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3859,11 +3859,13 @@ dependencies = [ "libffi", "libloading 0.9.0", "mac_address", + "mach2", "memchr", "memmap2", "nix", "num-traits", "num_cpus", + "object", "parking_lot", "paste", "pkcs8 0.11.0", diff --git a/Cargo.toml b/Cargo.toml index 95b57c139..638a86a72 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -259,6 +259,7 @@ log = "0.4.30" lz4_flex = "0.14" nix = { version = "0.31", features = ["fs", "user", "process", "term", "time", "signal", "ioctl", "socket", "sched", "zerocopy", "dir", "hostname", "net", "poll"] } mac_address = "1.1.3" +mach2 = "0.4.3" malachite-bigint = "0.14.0" malachite-q = "0.14.0" malachite-base = "0.14.0" @@ -271,6 +272,7 @@ num-integer = "0.1.46" num-traits = "0.2" num_cpus = "1.17.0" num_enum = { version = "0.7", default-features = false } +object = { version = "0.37.3", default-features = false, features = ["read_core", "elf", "macho", "pe", "std", "unaligned"] } openssl = "0.10.80" openssl-sys = "0.9.110" optional = "0.5" diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py index 9a7d41679..e7761cdde 100644 --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -22,6 +22,7 @@ import warnings __all__ = [ + "has_remote_subprocess_debugging", "requires_remote_subprocess_debugging", "requires_root_user", "requires_non_root_user", # globals "PIPE_MAX_SIZE", "verbose", "max_memuse", "use_resources", "failfast", @@ -657,6 +658,92 @@ def requires_working_socket(*, module=False): else: return unittest.skipUnless(has_socket_support, msg) +@functools.cache +def has_remote_subprocess_debugging(): + """Check if we have permissions to debug subprocesses remotely. + + Returns True if we have permissions, False if we don't. + Checks for: + - Platform support (Linux, macOS, Windows only) + - On Linux: process_vm_readv support + - _remote_debugging module availability + - Actual subprocess debugging permissions (e.g., macOS entitlements) + Result is cached. + """ + # Check platform support + if sys.platform not in ("linux", "darwin", "win32"): + return False + + try: + import _remote_debugging + except ImportError: + return False + + # On Linux, check for process_vm_readv support + if sys.platform == "linux": + if not getattr(_remote_debugging, "PROCESS_VM_READV_SUPPORTED", False): + return False + + # First check if we can read our own process + if not _remote_debugging.is_python_process(os.getpid()): + return False + + # Check subprocess access - debugging child processes may require + # additional permissions depending on platform security settings + import socket + import subprocess + + # Create a socket for child to signal readiness + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind(("127.0.0.1", 0)) + server.listen(1) + port = server.getsockname()[1] + + # Child connects to signal it's ready, then waits for parent to close + child_code = f""" +import socket +s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) +s.connect(("127.0.0.1", {port})) +s.recv(1) # Wait for parent to signal done +""" + proc = subprocess.Popen( + [sys.executable, "-c", child_code], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + try: + server.settimeout(5.0) + conn, _ = server.accept() + # Child is ready, test if we can probe it + result = _remote_debugging.is_python_process(proc.pid) + # Check if subprocess is still alive after probing + if proc.poll() is not None: + return False + conn.close() # Signal child to exit + return result + except (socket.timeout, OSError): + return False + finally: + server.close() + proc.kill() + proc.wait() + + +def requires_remote_subprocess_debugging(): + """Skip tests that require remote subprocess debugging permissions. + + This also implies subprocess support, so no need to use both + @requires_subprocess() and @requires_remote_subprocess_debugging(). + """ + if not has_subprocess_support: + return unittest.skip("requires subprocess support") + return unittest.skipUnless( + has_remote_subprocess_debugging(), + "requires remote subprocess debugging permissions" + ) + + # Does strftime() support glibc extension like '%4Y'? has_strftime_extensions = False if sys.platform != "win32": diff --git a/Lib/test/test_external_inspection.py b/Lib/test/test_external_inspection.py index 08779bdb0..797b13af5 100644 --- a/Lib/test/test_external_inspection.py +++ b/Lib/test/test_external_inspection.py @@ -1,11 +1,14 @@ +import asyncio import unittest import os import textwrap +import contextlib import importlib import sys import socket import threading import time +from contextlib import contextmanager from asyncio import staggered, taskgroups, base_events, tasks from unittest.mock import ANY from test.support import ( @@ -13,6 +16,7 @@ from test.support import ( SHORT_TIMEOUT, busy_retry, requires_gil_enabled, + requires_remote_subprocess_debugging, ) from test.support.import_helper import import_module from test.support.script_helper import make_script @@ -20,6 +24,31 @@ from test.support.socket_helper import find_unused_port import subprocess +# Profiling mode constants +PROFILING_MODE_WALL = 0 +PROFILING_MODE_CPU = 1 +PROFILING_MODE_GIL = 2 +PROFILING_MODE_ALL = 3 +PROFILING_MODE_EXCEPTION = 4 + +# Thread status flags +THREAD_STATUS_HAS_GIL = 1 << 0 +THREAD_STATUS_ON_CPU = 1 << 1 +THREAD_STATUS_UNKNOWN = 1 << 2 +THREAD_STATUS_HAS_EXCEPTION = 1 << 4 + +# Maximum number of retry attempts for operations that may fail transiently +MAX_TRIES = 10 +RETRY_DELAY = 0.1 + +# Exceptions that can occur transiently when reading from a live process +TRANSIENT_ERRORS = (OSError, RuntimeError, UnicodeDecodeError) + +try: + from concurrent import interpreters +except ImportError: + interpreters = None + PROCESS_VM_READV_SUPPORTED = False try: @@ -32,12 +61,149 @@ except ImportError: ) +# ============================================================================ +# Module-level helper functions +# ============================================================================ + + def _make_test_script(script_dir, script_basename, source): to_return = make_script(script_dir, script_basename, source) importlib.invalidate_caches() return to_return +def _create_server_socket(port, backlog=1): + """Create and configure a server socket for test communication.""" + server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server_socket.bind(("localhost", port)) + server_socket.settimeout(SHORT_TIMEOUT) + server_socket.listen(backlog) + return server_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 _wait_for_n_signals(sock, signal_pattern, count, timeout=SHORT_TIMEOUT): + """ + Wait for N occurrences of a signal pattern. + + Args: + sock: Connected socket to read from + signal_pattern: bytes pattern to count (e.g., b"ready") + count: Number of occurrences expected + timeout: Socket timeout in seconds + + Returns: + bytes: Complete accumulated response buffer + + Raises: + RuntimeError: If connection closed or timeout before receiving all signals + """ + sock.settimeout(timeout) + buffer = b"" + found_count = 0 + + while found_count < count: + try: + chunk = sock.recv(4096) + if not chunk: + raise RuntimeError( + f"Connection closed after {found_count}/{count} signals. " + f"Last 200 bytes: {buffer[-200:]!r}" + ) + buffer += chunk + # Count occurrences in entire buffer + found_count = buffer.count(signal_pattern) + except socket.timeout: + raise RuntimeError( + f"Timeout waiting for {count} signals (found {found_count}). " + f"Last 200 bytes: {buffer[-200:]!r}" + ) + + return buffer + + +@contextmanager +def _managed_subprocess(args, timeout=SHORT_TIMEOUT): + """ + Context manager for subprocess lifecycle management. + + Ensures process is properly terminated and cleaned up even on exceptions. + Uses graceful termination first, then forceful kill if needed. + """ + p = subprocess.Popen(args) + try: + yield p + finally: + try: + p.terminate() + try: + p.wait(timeout=timeout) + except subprocess.TimeoutExpired: + p.kill() + try: + p.wait(timeout=timeout) + except subprocess.TimeoutExpired: + pass # Process refuses to die, nothing more we can do + except OSError: + pass # Process already dead + + +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 + + +# ============================================================================ +# Decorators and skip conditions +# ============================================================================ + skip_if_not_supported = unittest.skipIf( ( sys.platform != "darwin" @@ -48,36 +214,311 @@ skip_if_not_supported = unittest.skipIf( ) +def requires_subinterpreters(meth): + """Decorator to skip a test if subinterpreters are not supported.""" + return unittest.skipIf(interpreters is None, "subinterpreters required")( + meth + ) + + +# ============================================================================ +# Simple wrapper functions for RemoteUnwinder +# ============================================================================ + def get_stack_trace(pid): - unwinder = RemoteUnwinder(pid, all_threads=True, debug=True) - return unwinder.get_stack_trace() + for _ in busy_retry(SHORT_TIMEOUT): + try: + unwinder = RemoteUnwinder(pid, all_threads=True, debug=True) + return unwinder.get_stack_trace() + except RuntimeError as e: + continue + raise RuntimeError("Failed to get stack trace after retries") def get_async_stack_trace(pid): - unwinder = RemoteUnwinder(pid, debug=True) - return unwinder.get_async_stack_trace() + for _ in busy_retry(SHORT_TIMEOUT): + try: + unwinder = RemoteUnwinder(pid, debug=True) + return unwinder.get_async_stack_trace() + except RuntimeError as e: + continue + raise RuntimeError("Failed to get async stack trace after retries") def get_all_awaited_by(pid): - unwinder = RemoteUnwinder(pid, debug=True) - return unwinder.get_all_awaited_by() + for _ in busy_retry(SHORT_TIMEOUT): + try: + unwinder = RemoteUnwinder(pid, debug=True) + return unwinder.get_all_awaited_by() + except RuntimeError as e: + continue + raise RuntimeError("Failed to get all awaited_by after retries") -class TestGetStackTrace(unittest.TestCase): +def _get_stack_trace_with_retry(unwinder, timeout=SHORT_TIMEOUT, condition=None): + """Get stack trace from an existing unwinder with retry for transient errors. + + This handles the case where we want to reuse an existing RemoteUnwinder + instance but still handle transient failures like "Failed to parse initial + frame in chain" that can occur when sampling at an inopportune moment. + If condition is provided, keeps retrying until condition(traces) is True. + """ + last_error = None + for _ in busy_retry(timeout): + try: + traces = unwinder.get_stack_trace() + if condition is None or condition(traces): + return traces + # Condition not met yet, keep retrying + except TRANSIENT_ERRORS as e: + last_error = e + continue + if last_error: + raise RuntimeError( + f"Failed to get stack trace after retries: {last_error}" + ) + raise RuntimeError("Condition never satisfied within timeout") + + +# ============================================================================ +# Base test class with shared infrastructure +# ============================================================================ + + +class RemoteInspectionTestBase(unittest.TestCase): + """Base class for remote inspection tests with common helpers.""" + maxDiff = None + def _run_script_and_get_trace( + self, + script, + trace_func, + wait_for_signals=None, + port=None, + backlog=1, + ): + """ + Common pattern: run a script, wait for signals, get trace. + + Args: + script: Script content (will be formatted with port if {port} present) + trace_func: Function to call with pid to get trace (e.g., get_stack_trace) + wait_for_signals: Signal(s) to wait for before getting trace + port: Port to use (auto-selected if None) + backlog: Socket listen backlog + + Returns: + tuple: (trace_result, script_name) + """ + if port is None: + port = find_unused_port() + + # Format script with port if needed + if "{port}" in script or "{{port}}" in script: + script = script.replace("{{port}}", "{port}").format(port=port) + + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port, backlog) + script_name = _make_test_script(script_dir, "script", script) + client_socket = None + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + if wait_for_signals: + _wait_for_signal(client_socket, wait_for_signals) + + trace = trace_func(p.pid) + return trace, script_name + finally: + _cleanup_sockets(client_socket, server_socket) + + def _find_frame_in_trace(self, stack_trace, predicate): + """ + Find a frame matching predicate in stack trace. + + Args: + stack_trace: List of InterpreterInfo objects + predicate: Function(frame) -> bool + + Returns: + FrameInfo or None + """ + for interpreter_info in stack_trace: + for thread_info in interpreter_info.threads: + for frame in thread_info.frame_info: + if predicate(frame): + return frame + return None + + def _find_thread_by_id(self, stack_trace, thread_id): + """Find a thread by its native thread ID.""" + for interpreter_info in stack_trace: + for thread_info in interpreter_info.threads: + if thread_info.thread_id == thread_id: + return thread_info + return None + + def _find_thread_with_frame(self, stack_trace, frame_predicate): + """Find a thread containing a frame matching predicate.""" + for interpreter_info in stack_trace: + for thread_info in interpreter_info.threads: + for frame in thread_info.frame_info: + if frame_predicate(frame): + return thread_info + return None + + def _get_thread_statuses(self, stack_trace): + """Extract thread_id -> status mapping from stack trace.""" + statuses = {} + for interpreter_info in stack_trace: + for thread_info in interpreter_info.threads: + statuses[thread_info.thread_id] = thread_info.status + return statuses + + def _get_task_id_map(self, stack_trace): + """Create task_id -> task mapping from async stack trace.""" + return {task.task_id: task for task in stack_trace[0].awaited_by} + + def _get_awaited_by_relationships(self, stack_trace): + """Extract task name to awaited_by set mapping.""" + id_to_task = self._get_task_id_map(stack_trace) + return { + task.task_name: set( + id_to_task[awaited.task_name].task_name + for awaited in task.awaited_by + ) + for task in stack_trace[0].awaited_by + } + + def _extract_coroutine_stacks(self, stack_trace): + """Extract and format coroutine stacks from tasks.""" + return { + task.task_name: sorted( + tuple(tuple(frame) for frame in coro.call_stack) + for coro in task.coroutine_stack + ) + for task in stack_trace[0].awaited_by + } + + @staticmethod + def _frame_to_lineno_tuple(frame): + """Convert frame to (filename, lineno, funcname, opcode) tuple. + + This extracts just the line number from the location, ignoring column + offsets which can vary due to sampling timing (e.g., when two statements + are on the same line, the sample might catch either one). + """ + filename, location, funcname, opcode = frame + return (filename, location.lineno, funcname, opcode) + + def _extract_coroutine_stacks_lineno_only(self, stack_trace): + """Extract coroutine stacks with line numbers only (no column offsets). + + Use this for tests where sampling timing can cause column offset + variations (e.g., 'expr1; expr2' on the same line). + """ + return { + task.task_name: sorted( + tuple(self._frame_to_lineno_tuple(frame) for frame in coro.call_stack) + for coro in task.coroutine_stack + ) + for task in stack_trace[0].awaited_by + } + + +# ============================================================================ +# Test classes +# ============================================================================ + + +@requires_remote_subprocess_debugging() +class TestSelfStackTrace(RemoteInspectionTestBase): + @skip_if_not_supported + def test_long_task_name_is_truncated(self): + # gh-157788 + async def main(): + asyncio.create_task(asyncio.sleep(10_000), name="x" * 300) + await asyncio.sleep(0) + names = [ + task.task_name + for info in RemoteUnwinder(os.getpid()).get_all_awaited_by() + for task in info.awaited_by + ] + return asyncio.current_task().get_name(), names + + main_name, names = asyncio.run(main()) + self.assertIn(main_name, names) + self.assertEqual([len(n) for n in names if n.startswith("x")], [255]) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_self_trace_with_large_linetable(self): + script = textwrap.dedent("""\ + import os + import _remote_debugging + + assignments = "\\n".join( + f"value_{i} = {i}" for i in range(1000) + ) + expected_lineno = len(assignments.splitlines()) + 1 + source = ( + f"{assignments}\\n" + "stack_trace = " + "_remote_debugging.RemoteUnwinder(os.getpid()).get_stack_trace()\\n" + ) + code = compile(source, "large_linetable.py", "exec") + assert len(code.co_linetable) > 4096, len(code.co_linetable) + namespace = {"os": os, "_remote_debugging": _remote_debugging} + exec(code, namespace) + large_linetable_frames = [ + frame + for interpreter in namespace["stack_trace"] + for thread in interpreter.threads + for frame in thread.frame_info + if frame.filename == "large_linetable.py" + ] + assert len(large_linetable_frames) == 1, large_linetable_frames + assert large_linetable_frames[0].location.lineno == expected_lineno, ( + large_linetable_frames[0] + ) + """) + + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=SHORT_TIMEOUT, + ) + self.assertEqual( + result.returncode, 0, + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + + +@requires_remote_subprocess_debugging() +class TestGetStackTrace(RemoteInspectionTestBase): @skip_if_not_supported @unittest.skipIf( sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, "Test only runs on Linux with process_vm_readv support", ) def test_remote_stack_trace(self): - # Spawn a process with some realistic Python code port = find_unused_port() script = textwrap.dedent( f"""\ import time, sys, socket, threading - # Connect to the test process + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) @@ -90,66 +531,59 @@ class TestGetStackTrace(unittest.TestCase): foo() def foo(): - sock.sendall(b"ready:thread\\n"); time.sleep(10_000) # same line number + sock.sendall(b"ready:thread\\n"); time.sleep(10_000) t = threading.Thread(target=bar) t.start() - sock.sendall(b"ready:main\\n"); t.join() # same line number + sock.sendall(b"ready:main\\n"); t.join() """ ) - stack_trace = None + with os_helper.temp_dir() as work_dir: script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - # Create a socket server to communicate with the target process - server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) - + server_socket = _create_server_socket(port) script_name = _make_test_script(script_dir, "script", script) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() - response = b"" - while ( - b"ready:main" not in response - or b"ready:thread" not in response - ): - response += client_socket.recv(1024) - stack_trace = get_stack_trace(p.pid) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + _wait_for_signal( + client_socket, [b"ready:main", b"ready:thread"] + ) + + stack_trace = get_stack_trace(p.pid) + + # Find expected thread stack by funcname + found_thread = self._find_thread_with_frame( + stack_trace, + lambda f: f.funcname == "foo" and f.location.lineno == 15, + ) + self.assertIsNotNone( + found_thread, "Expected thread stack trace not found" + ) + # Check the funcnames in order + funcnames = [f.funcname for f in found_thread.frame_info] + self.assertEqual( + funcnames[:6], + ["foo", "baz", "bar", "Thread.run", "Thread._bootstrap_inner", "Thread._bootstrap"] + ) + + # Check main thread + found_main = self._find_frame_in_trace( + stack_trace, + lambda f: f.funcname == "" and f.location.lineno == 19, + ) + self.assertIsNotNone( + found_main, "Main thread stack trace not found" + ) finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) - - thread_expected_stack_trace = [ - FrameInfo([script_name, 15, "foo"]), - FrameInfo([script_name, 12, "baz"]), - FrameInfo([script_name, 9, "bar"]), - FrameInfo([threading.__file__, ANY, "Thread.run"]), - ] - # Is possible that there are more threads, so we check that the - # expected stack traces are in the result (looking at you Windows!) - self.assertIn((ANY, thread_expected_stack_trace), stack_trace) - - # Check that the main thread stack trace is in the result - frame = FrameInfo([script_name, 19, ""]) - for _, stack in stack_trace: - if frame in stack: - break - else: - self.fail("Main thread stack trace not found in result") + _cleanup_sockets(client_socket, server_socket) @skip_if_not_supported @unittest.skipIf( @@ -193,13 +627,81 @@ class TestGetStackTrace(unittest.TestCase): f"stdout: {result.stdout}\nstderr: {result.stderr}" ) + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_remote_stack_trace_non_ascii_names(self): + # Exercise each PyUnicode kind (1-byte non-ASCII, 2-byte BMP, + # 4-byte non-BMP) for both the filename and the function name + # reported in the stack trace. + latin1 = "zażółć" # 1-byte non-ASCII (forces non-ASCII path) + bmp = "λάμβδα" # 2-byte BMP + astral = "𐌀𐌁𐌂𐌃" # 4-byte non-BMP (Old Italic; XID, no NFKC fold) + func_name = f"{latin1}_{bmp}_{astral}" + script_basename = f"mod_{latin1}_{bmp}_{astral}" + + port = find_unused_port() + script = textwrap.dedent( + f"""\ + import socket + import time + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def {func_name}(): + sock.sendall(b"ready") + time.sleep(10_000) + + {func_name}() + """ + ) + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + script_name = _make_test_script(script_dir, script_basename, script) + + server_socket = _create_server_socket(port) + client_socket = None + try: + p = subprocess.Popen([sys.executable, script_name]) + client_socket, _ = server_socket.accept() + server_socket.close() + _wait_for_signal(client_socket, b"ready") + + stack_trace = get_stack_trace(p.pid) + except PermissionError: + self.skipTest("Insufficient permissions to read the stack trace") + finally: + if client_socket is not None: + client_socket.close() + p.kill() + p.wait(timeout=SHORT_TIMEOUT) + + frames = [ + frame + for interp in stack_trace + for thread in interp.threads + for frame in thread.frame_info + ] + target = next( + (f for f in frames if f.funcname == func_name), None + ) + self.assertIsNotNone( + target, + f"Frame for {func_name!r} missing; got " + f"{[(f.filename, f.funcname) for f in frames]}", + ) + self.assertEqual(target.filename, script_name) + @skip_if_not_supported @unittest.skipIf( sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, "Test only runs on Linux with process_vm_readv support", ) def test_async_remote_stack_trace(self): - # Spawn a process with some realistic Python code port = find_unused_port() script = textwrap.dedent( f"""\ @@ -207,12 +709,12 @@ class TestGetStackTrace(unittest.TestCase): import time import sys import socket - # Connect to the test process + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) def c5(): - sock.sendall(b"ready"); time.sleep(10_000) # same line number + sock.sendall(b"ready"); time.sleep(10_000) async def c4(): await asyncio.sleep(0) @@ -243,7 +745,7 @@ class TestGetStackTrace(unittest.TestCase): asyncio.run(main(), loop_factory={{TASK_FACTORY}}) """ ) - stack_trace = None + for task_factory_variant in "asyncio.new_event_loop", "new_eager_loop": with ( self.subTest(task_factory_variant=task_factory_variant), @@ -251,195 +753,152 @@ class TestGetStackTrace(unittest.TestCase): ): script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - server_socket = socket.socket( - socket.AF_INET, socket.SOCK_STREAM - ) - server_socket.setsockopt( - socket.SOL_SOCKET, socket.SO_REUSEADDR, 1 - ) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) + + server_socket = _create_server_socket(port) script_name = _make_test_script( script_dir, "script", script.format(TASK_FACTORY=task_factory_variant), ) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() - response = client_socket.recv(1024) - self.assertEqual(response, b"ready") - stack_trace = get_async_stack_trace(p.pid) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) - finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) + with _managed_subprocess( + [sys.executable, script_name] + ) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None - # First check all the tasks are present - tasks_names = [ - task.task_name for task in stack_trace[0].awaited_by - ] - for task_name in ["c2_root", "sub_main_1", "sub_main_2"]: - self.assertIn(task_name, tasks_names) + response = _wait_for_signal(client_socket, b"ready") + self.assertIn(b"ready", response) - # Now ensure that the awaited_by_relationships are correct - id_to_task = { - task.task_id: task for task in stack_trace[0].awaited_by - } - task_name_to_awaited_by = { - task.task_name: set( - id_to_task[awaited.task_name].task_name - for awaited in task.awaited_by - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - task_name_to_awaited_by, - { - "c2_root": {"Task-1", "sub_main_1", "sub_main_2"}, - "Task-1": set(), - "sub_main_1": {"Task-1"}, - "sub_main_2": {"Task-1"}, - }, - ) + stack_trace = get_async_stack_trace(p.pid) - # Now ensure that the coroutine stacks are correct - coroutine_stacks = { - task.task_name: sorted( - tuple(tuple(frame) for frame in coro.call_stack) - for coro in task.coroutine_stack - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - coroutine_stacks, - { - "Task-1": [ - ( - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] - ), - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] - ), - tuple([script_name, 26, "main"]), - ) - ], - "c2_root": [ - ( - tuple([script_name, 10, "c5"]), - tuple([script_name, 14, "c4"]), - tuple([script_name, 17, "c3"]), - tuple([script_name, 20, "c2"]), - ) - ], - "sub_main_1": [(tuple([script_name, 23, "c1"]),)], - "sub_main_2": [(tuple([script_name, 23, "c1"]),)], - }, - ) + # Check all tasks are present + tasks_names = [ + task.task_name + for task in stack_trace[0].awaited_by + ] + for task_name in [ + "c2_root", + "sub_main_1", + "sub_main_2", + ]: + self.assertIn(task_name, tasks_names) - # Now ensure the coroutine stacks for the awaited_by relationships are correct. - awaited_by_coroutine_stacks = { - task.task_name: sorted( - ( - id_to_task[coro.task_name].task_name, - tuple(tuple(frame) for frame in coro.call_stack), + # Check awaited_by relationships + relationships = self._get_awaited_by_relationships( + stack_trace ) - for coro in task.awaited_by - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - awaited_by_coroutine_stacks, - { - "Task-1": [], - "c2_root": [ - ( - "Task-1", + self.assertEqual( + relationships, + { + "c2_root": { + "Task-1", + "sub_main_1", + "sub_main_2", + }, + "Task-1": set(), + "sub_main_1": {"Task-1"}, + "sub_main_2": {"Task-1"}, + }, + ) + + # Check coroutine stacks (using line numbers only to avoid + # flakiness from column offset variations when sampling + # catches different statements on the same line) + coroutine_stacks = self._extract_coroutine_stacks_lineno_only( + stack_trace + ) + self.assertEqual( + coroutine_stacks, + { + "Task-1": [ + ( + (taskgroups.__file__, ANY, "TaskGroup._aexit", None), + (taskgroups.__file__, ANY, "TaskGroup.__aexit__", None), + (script_name, 26, "main", None), + ) + ], + "c2_root": [ + ( + (script_name, 10, "c5", None), + (script_name, 14, "c4", None), + (script_name, 17, "c3", None), + (script_name, 20, "c2", None), + ) + ], + "sub_main_1": [ + ((script_name, 23, "c1", None),) + ], + "sub_main_2": [ + ((script_name, 23, "c1", None),) + ], + }, + ) + + # Check awaited_by coroutine stacks (line numbers only) + id_to_task = self._get_task_id_map(stack_trace) + awaited_by_coroutine_stacks = { + task.task_name: sorted( ( + id_to_task[coro.task_name].task_name, tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] + self._frame_to_lineno_tuple(frame) + for frame in coro.call_stack ), - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] - ), - tuple([script_name, 26, "main"]), - ), - ), - ("sub_main_1", (tuple([script_name, 23, "c1"]),)), - ("sub_main_2", (tuple([script_name, 23, "c1"]),)), - ], - "sub_main_1": [ - ( - "Task-1", - ( - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] - ), - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] - ), - tuple([script_name, 26, "main"]), - ), + ) + for coro in task.awaited_by ) - ], - "sub_main_2": [ - ( - "Task-1", - ( - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] + for task in stack_trace[0].awaited_by + } + self.assertEqual( + awaited_by_coroutine_stacks, + { + "Task-1": [], + "c2_root": [ + ( + "Task-1", + ( + (taskgroups.__file__, ANY, "TaskGroup._aexit", None), + (taskgroups.__file__, ANY, "TaskGroup.__aexit__", None), + (script_name, 26, "main", None), + ), ), - tuple( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] + ( + "sub_main_1", + ((script_name, 23, "c1", None),), ), - tuple([script_name, 26, "main"]), - ), - ) - ], - }, - ) + ( + "sub_main_2", + ((script_name, 23, "c1", None),), + ), + ], + "sub_main_1": [ + ( + "Task-1", + ( + (taskgroups.__file__, ANY, "TaskGroup._aexit", None), + (taskgroups.__file__, ANY, "TaskGroup.__aexit__", None), + (script_name, 26, "main", None), + ), + ) + ], + "sub_main_2": [ + ( + "Task-1", + ( + (taskgroups.__file__, ANY, "TaskGroup._aexit", None), + (taskgroups.__file__, ANY, "TaskGroup.__aexit__", None), + (script_name, 26, "main", None), + ), + ) + ], + }, + ) + finally: + _cleanup_sockets(client_socket, server_socket) @skip_if_not_supported @unittest.skipIf( @@ -447,7 +906,6 @@ class TestGetStackTrace(unittest.TestCase): "Test only runs on Linux with process_vm_readv support", ) def test_asyncgen_remote_stack_trace(self): - # Spawn a process with some realistic Python code port = find_unused_port() script = textwrap.dedent( f"""\ @@ -455,12 +913,12 @@ class TestGetStackTrace(unittest.TestCase): import time import sys import socket - # Connect to the test process + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) async def gen_nested_call(): - sock.sendall(b"ready"); time.sleep(10_000) # same line number + sock.sendall(b"ready"); time.sleep(10_000) async def gen(): for num in range(2): @@ -475,59 +933,53 @@ class TestGetStackTrace(unittest.TestCase): asyncio.run(main()) """ ) - stack_trace = None + with os_helper.temp_dir() as work_dir: script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - # Create a socket server to communicate with the target process - server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) + + server_socket = _create_server_socket(port) script_name = _make_test_script(script_dir, "script", script) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() - response = client_socket.recv(1024) - self.assertEqual(response, b"ready") - stack_trace = get_async_stack_trace(p.pid) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) - finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None - # For this simple asyncgen test, we only expect one task with the full coroutine stack - self.assertEqual(len(stack_trace[0].awaited_by), 1) - task = stack_trace[0].awaited_by[0] - self.assertEqual(task.task_name, "Task-1") + response = _wait_for_signal(client_socket, b"ready") + self.assertIn(b"ready", response) - # Check the coroutine stack - based on actual output, only shows main - coroutine_stack = sorted( - tuple(tuple(frame) for frame in coro.call_stack) - for coro in task.coroutine_stack - ) - self.assertEqual( - coroutine_stack, - [ - ( - tuple([script_name, 10, "gen_nested_call"]), - tuple([script_name, 16, "gen"]), - tuple([script_name, 19, "main"]), + stack_trace = get_async_stack_trace(p.pid) + + # For this simple asyncgen test, we only expect one task + self.assertEqual(len(stack_trace[0].awaited_by), 1) + task = stack_trace[0].awaited_by[0] + self.assertEqual(task.task_name, "Task-1") + + # Check the coroutine stack (using line numbers only to avoid + # flakiness from column offset variations when sampling + # catches different statements on the same line) + coroutine_stack = sorted( + tuple(self._frame_to_lineno_tuple(frame) for frame in coro.call_stack) + for coro in task.coroutine_stack + ) + self.assertEqual( + coroutine_stack, + [ + ( + (script_name, 10, "gen_nested_call", None), + (script_name, 16, "gen", None), + (script_name, 19, "main", None), + ) + ], ) - ], - ) - # No awaited_by relationships expected for this simple case - self.assertEqual(task.awaited_by, []) + # No awaited_by relationships expected + self.assertEqual(task.awaited_by, []) + finally: + _cleanup_sockets(client_socket, server_socket) @skip_if_not_supported @unittest.skipIf( @@ -535,7 +987,6 @@ class TestGetStackTrace(unittest.TestCase): "Test only runs on Linux with process_vm_readv support", ) def test_async_gather_remote_stack_trace(self): - # Spawn a process with some realistic Python code port = find_unused_port() script = textwrap.dedent( f"""\ @@ -543,13 +994,13 @@ class TestGetStackTrace(unittest.TestCase): import time import sys import socket - # Connect to the test process + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) async def deep(): await asyncio.sleep(0) - sock.sendall(b"ready"); time.sleep(10_000) # same line number + sock.sendall(b"ready"); time.sleep(10_000) async def c1(): await asyncio.sleep(0) @@ -564,103 +1015,89 @@ class TestGetStackTrace(unittest.TestCase): asyncio.run(main()) """ ) - stack_trace = None + with os_helper.temp_dir() as work_dir: script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - # Create a socket server to communicate with the target process - server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) + + server_socket = _create_server_socket(port) script_name = _make_test_script(script_dir, "script", script) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() - response = client_socket.recv(1024) - self.assertEqual(response, b"ready") - stack_trace = get_async_stack_trace(p.pid) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) - finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None - # First check all the tasks are present - tasks_names = [ - task.task_name for task in stack_trace[0].awaited_by - ] - for task_name in ["Task-1", "Task-2"]: - self.assertIn(task_name, tasks_names) + response = _wait_for_signal(client_socket, b"ready") + self.assertIn(b"ready", response) - # Now ensure that the awaited_by_relationships are correct - id_to_task = { - task.task_id: task for task in stack_trace[0].awaited_by - } - task_name_to_awaited_by = { - task.task_name: set( - id_to_task[awaited.task_name].task_name - for awaited in task.awaited_by - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - task_name_to_awaited_by, - { - "Task-1": set(), - "Task-2": {"Task-1"}, - }, - ) + stack_trace = get_async_stack_trace(p.pid) - # Now ensure that the coroutine stacks are correct - coroutine_stacks = { - task.task_name: sorted( - tuple(tuple(frame) for frame in coro.call_stack) - for coro in task.coroutine_stack - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - coroutine_stacks, - { - "Task-1": [(tuple([script_name, 21, "main"]),)], - "Task-2": [ - ( - tuple([script_name, 11, "deep"]), - tuple([script_name, 15, "c1"]), - ) - ], - }, - ) + # Check all tasks are present + tasks_names = [ + task.task_name for task in stack_trace[0].awaited_by + ] + for task_name in ["Task-1", "Task-2"]: + self.assertIn(task_name, tasks_names) - # Now ensure the coroutine stacks for the awaited_by relationships are correct. - awaited_by_coroutine_stacks = { - task.task_name: sorted( - ( - id_to_task[coro.task_name].task_name, - tuple(tuple(frame) for frame in coro.call_stack), + # Check awaited_by relationships + relationships = self._get_awaited_by_relationships( + stack_trace ) - for coro in task.awaited_by - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - awaited_by_coroutine_stacks, - { - "Task-1": [], - "Task-2": [ - ("Task-1", (tuple([script_name, 21, "main"]),)) - ], - }, - ) + self.assertEqual( + relationships, + { + "Task-1": set(), + "Task-2": {"Task-1"}, + }, + ) + + # Check coroutine stacks (using line numbers only to avoid + # flakiness from column offset variations when sampling + # catches different statements on the same line) + coroutine_stacks = self._extract_coroutine_stacks_lineno_only( + stack_trace + ) + self.assertEqual( + coroutine_stacks, + { + "Task-1": [((script_name, 21, "main", None),)], + "Task-2": [ + ( + (script_name, 11, "deep", None), + (script_name, 15, "c1", None), + ) + ], + }, + ) + + # Check awaited_by coroutine stacks (line numbers only) + id_to_task = self._get_task_id_map(stack_trace) + awaited_by_coroutine_stacks = { + task.task_name: sorted( + ( + id_to_task[coro.task_name].task_name, + tuple( + self._frame_to_lineno_tuple(frame) for frame in coro.call_stack + ), + ) + for coro in task.awaited_by + ) + for task in stack_trace[0].awaited_by + } + self.assertEqual( + awaited_by_coroutine_stacks, + { + "Task-1": [], + "Task-2": [ + ("Task-1", ((script_name, 21, "main", None),)) + ], + }, + ) + finally: + _cleanup_sockets(client_socket, server_socket) @skip_if_not_supported @unittest.skipIf( @@ -668,7 +1105,6 @@ class TestGetStackTrace(unittest.TestCase): "Test only runs on Linux with process_vm_readv support", ) def test_async_staggered_race_remote_stack_trace(self): - # Spawn a process with some realistic Python code port = find_unused_port() script = textwrap.dedent( f"""\ @@ -676,13 +1112,13 @@ class TestGetStackTrace(unittest.TestCase): import time import sys import socket - # Connect to the test process + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) async def deep(): await asyncio.sleep(0) - sock.sendall(b"ready"); time.sleep(10_000) # same line number + sock.sendall(b"ready"); time.sleep(10_000) async def c1(): await asyncio.sleep(0) @@ -700,123 +1136,101 @@ class TestGetStackTrace(unittest.TestCase): asyncio.run(main()) """ ) - stack_trace = None + with os_helper.temp_dir() as work_dir: script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - # Create a socket server to communicate with the target process - server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) + + server_socket = _create_server_socket(port) script_name = _make_test_script(script_dir, "script", script) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() - response = client_socket.recv(1024) - self.assertEqual(response, b"ready") - stack_trace = get_async_stack_trace(p.pid) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) - finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None - # First check all the tasks are present - tasks_names = [ - task.task_name for task in stack_trace[0].awaited_by - ] - for task_name in ["Task-1", "Task-2"]: - self.assertIn(task_name, tasks_names) + response = _wait_for_signal(client_socket, b"ready") + self.assertIn(b"ready", response) - # Now ensure that the awaited_by_relationships are correct - id_to_task = { - task.task_id: task for task in stack_trace[0].awaited_by - } - task_name_to_awaited_by = { - task.task_name: set( - id_to_task[awaited.task_name].task_name - for awaited in task.awaited_by - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - task_name_to_awaited_by, - { - "Task-1": set(), - "Task-2": {"Task-1"}, - }, - ) + stack_trace = get_async_stack_trace(p.pid) - # Now ensure that the coroutine stacks are correct - coroutine_stacks = { - task.task_name: sorted( - tuple(tuple(frame) for frame in coro.call_stack) - for coro in task.coroutine_stack - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - coroutine_stacks, - { - "Task-1": [ - ( - tuple([staggered.__file__, ANY, "staggered_race"]), - tuple([script_name, 21, "main"]), - ) - ], - "Task-2": [ - ( - tuple([script_name, 11, "deep"]), - tuple([script_name, 15, "c1"]), - tuple( - [ - staggered.__file__, - ANY, - "staggered_race..run_one_coro", - ] - ), - ) - ], - }, - ) + # Check all tasks are present + tasks_names = [ + task.task_name for task in stack_trace[0].awaited_by + ] + for task_name in ["Task-1", "Task-2"]: + self.assertIn(task_name, tasks_names) - # Now ensure the coroutine stacks for the awaited_by relationships are correct. - awaited_by_coroutine_stacks = { - task.task_name: sorted( - ( - id_to_task[coro.task_name].task_name, - tuple(tuple(frame) for frame in coro.call_stack), + # Check awaited_by relationships + relationships = self._get_awaited_by_relationships( + stack_trace ) - for coro in task.awaited_by - ) - for task in stack_trace[0].awaited_by - } - self.assertEqual( - awaited_by_coroutine_stacks, - { - "Task-1": [], - "Task-2": [ - ( - "Task-1", + self.assertEqual( + relationships, + { + "Task-1": set(), + "Task-2": {"Task-1"}, + }, + ) + + # Check coroutine stacks (using line numbers only to avoid + # flakiness from column offset variations when sampling + # catches different statements on the same line) + coroutine_stacks = self._extract_coroutine_stacks_lineno_only( + stack_trace + ) + self.assertEqual( + coroutine_stacks, + { + "Task-1": [ + ( + (staggered.__file__, ANY, "staggered_race", None), + (script_name, 21, "main", None), + ) + ], + "Task-2": [ + ( + (script_name, 11, "deep", None), + (script_name, 15, "c1", None), + (staggered.__file__, ANY, "staggered_race..run_one_coro", None), + ) + ], + }, + ) + + # Check awaited_by coroutine stacks (line numbers only) + id_to_task = self._get_task_id_map(stack_trace) + awaited_by_coroutine_stacks = { + task.task_name: sorted( ( + id_to_task[coro.task_name].task_name, tuple( - [staggered.__file__, ANY, "staggered_race"] + self._frame_to_lineno_tuple(frame) for frame in coro.call_stack ), - tuple([script_name, 21, "main"]), - ), + ) + for coro in task.awaited_by ) - ], - }, - ) + for task in stack_trace[0].awaited_by + } + self.assertEqual( + awaited_by_coroutine_stacks, + { + "Task-1": [], + "Task-2": [ + ( + "Task-1", + ( + (staggered.__file__, ANY, "staggered_race", None), + (script_name, 21, "main", None), + ), + ) + ], + }, + ) + finally: + _cleanup_sockets(client_socket, server_socket) @skip_if_not_supported @unittest.skipIf( @@ -824,6 +1238,10 @@ class TestGetStackTrace(unittest.TestCase): "Test only runs on Linux with process_vm_readv support", ) def test_async_global_awaited_by(self): + # Reduced from 1000 to 100 to avoid file descriptor exhaustion + # when running tests in parallel (e.g., -j 20) + NUM_TASKS = 100 + port = find_unused_port() script = textwrap.dedent( f"""\ @@ -839,7 +1257,6 @@ class TestGetStackTrace(unittest.TestCase): PORT = socket_helper.find_unused_port() connections = 0 - # Connect to the test process sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) @@ -862,23 +1279,16 @@ class TestGetStackTrace(unittest.TestCase): assert message == data.decode() writer.close() await writer.wait_closed() - # Signal we are ready to sleep sock.sendall(b"ready") await asyncio.sleep(SHORT_TIMEOUT) async def echo_client_spam(server): async with asyncio.TaskGroup() as tg: - while connections < 1000: + while connections < {NUM_TASKS}: msg = list(ascii_lowercase + digits) random.shuffle(msg) tg.create_task(echo_client("".join(msg))) await asyncio.sleep(0) - # at least a 1000 tasks created. Each task will signal - # when is ready to avoid the race caused by the fact that - # tasks are waited on tg.__exit__ and we cannot signal when - # that happens otherwise - # at this point all client tasks completed without assertion errors - # let's wrap up the test server.close() await server.wait_closed() @@ -893,231 +1303,361 @@ class TestGetStackTrace(unittest.TestCase): asyncio.run(main()) """ ) - stack_trace = None + with os_helper.temp_dir() as work_dir: script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - # Create a socket server to communicate with the target process - server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) + + server_socket = _create_server_socket(port) script_name = _make_test_script(script_dir, "script", script) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() - for _ in range(1000): - expected_response = b"ready" - response = client_socket.recv(len(expected_response)) - self.assertEqual(response, expected_response) - for _ in busy_retry(SHORT_TIMEOUT): + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + # Wait for NUM_TASKS "ready" signals try: - all_awaited_by = get_all_awaited_by(p.pid) - except RuntimeError as re: - # This call reads a linked list in another process with - # no synchronization. That occasionally leads to invalid - # reads. Here we avoid making the test flaky. - msg = str(re) - if msg.startswith("Task list appears corrupted"): - continue - elif msg.startswith( - "Invalid linked list structure reading remote memory" - ): - continue - elif msg.startswith("Unknown error reading memory"): - continue - elif msg.startswith("Unhandled frame owner"): - continue - raise # Unrecognized exception, safest not to ignore it - else: - break - # expected: a list of two elements: 1 thread, 1 interp - self.assertEqual(len(all_awaited_by), 2) - # expected: a tuple with the thread ID and the awaited_by list - self.assertEqual(len(all_awaited_by[0]), 2) - # expected: no tasks in the fallback per-interp task list - self.assertEqual(all_awaited_by[1], (0, [])) - entries = all_awaited_by[0][1] - # expected: at least 1000 pending tasks - self.assertGreaterEqual(len(entries), 1000) - # the first three tasks stem from the code structure - main_stack = [ - FrameInfo([taskgroups.__file__, ANY, "TaskGroup._aexit"]), - FrameInfo( - [taskgroups.__file__, ANY, "TaskGroup.__aexit__"] - ), - FrameInfo([script_name, 60, "main"]), - ] - self.assertIn( - TaskInfo( - [ANY, "Task-1", [CoroInfo([main_stack, ANY])], []] - ), - entries, - ) - self.assertIn( - TaskInfo( - [ - ANY, - "server task", - [ - CoroInfo( - [ - [ - FrameInfo( - [ - base_events.__file__, - ANY, - "Server.serve_forever", - ] - ) - ], - ANY, - ] - ) - ], - [ - CoroInfo( - [ - [ - FrameInfo( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] - ), - FrameInfo( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] - ), - FrameInfo( - [script_name, ANY, "main"] - ), - ], - ANY, - ] - ) - ], - ] - ), - entries, - ) - self.assertIn( - TaskInfo( - [ - ANY, - "Task-4", - [ - CoroInfo( - [ - [ - FrameInfo( - [tasks.__file__, ANY, "sleep"] - ), - FrameInfo( - [ - script_name, - 38, - "echo_client", - ] - ), - ], - ANY, - ] - ) - ], - [ - CoroInfo( - [ - [ - FrameInfo( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] - ), - FrameInfo( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] - ), - FrameInfo( - [ - script_name, - 41, - "echo_client_spam", - ] - ), - ], - ANY, - ] - ) - ], - ] - ), - entries, - ) + _wait_for_n_signals(client_socket, b"ready", NUM_TASKS) + except RuntimeError as e: + self.fail(str(e)) - expected_awaited_by = [ - CoroInfo( - [ - [ - FrameInfo( - [ - taskgroups.__file__, - ANY, - "TaskGroup._aexit", - ] - ), - FrameInfo( - [ - taskgroups.__file__, - ANY, - "TaskGroup.__aexit__", - ] - ), - FrameInfo( - [script_name, 41, "echo_client_spam"] - ), - ], - ANY, - ] - ) - ] - tasks_with_awaited = [ - task - for task in entries - if task.awaited_by == expected_awaited_by - ] - self.assertGreaterEqual(len(tasks_with_awaited), 1000) + all_awaited_by = get_all_awaited_by(p.pid) - # the final task will have some random number, but it should for - # sure be one of the echo client spam horde (In windows this is not true - # for some reason) - if sys.platform != "win32": - self.assertEqual( - tasks_with_awaited[-1].awaited_by, - entries[-1].awaited_by, + # Expected: a list of two elements: 1 thread, 1 interp + self.assertEqual(len(all_awaited_by), 2) + # Expected: a tuple with the thread ID and the awaited_by list + self.assertEqual(len(all_awaited_by[0]), 2) + # Expected: no tasks in the fallback per-interp task list + self.assertEqual(all_awaited_by[1], (0, [])) + + entries = all_awaited_by[0][1] + # Expected: at least NUM_TASKS pending tasks + self.assertGreaterEqual(len(entries), NUM_TASKS) + + # Check the main task structure + main_stack = [ + FrameInfo( + [taskgroups.__file__, ANY, "TaskGroup._aexit", ANY] + ), + FrameInfo( + [taskgroups.__file__, ANY, "TaskGroup.__aexit__", ANY] + ), + FrameInfo([script_name, ANY, "main", ANY]), + ] + self.assertIn( + TaskInfo( + [ANY, "Task-1", [CoroInfo([main_stack, ANY])], []] + ), + entries, ) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) + self.assertIn( + TaskInfo( + [ + ANY, + "server task", + [ + CoroInfo( + [ + [ + FrameInfo( + [ + base_events.__file__, + ANY, + "Server.serve_forever", + ANY, + ] + ) + ], + ANY, + ] + ) + ], + [ + CoroInfo( + [ + [ + FrameInfo( + [ + taskgroups.__file__, + ANY, + "TaskGroup._aexit", + ANY, + ] + ), + FrameInfo( + [ + taskgroups.__file__, + ANY, + "TaskGroup.__aexit__", + ANY, + ] + ), + FrameInfo( + [script_name, ANY, "main", ANY] + ), + ], + ANY, + ] + ) + ], + ] + ), + entries, + ) + self.assertIn( + TaskInfo( + [ + ANY, + "Task-4", + [ + CoroInfo( + [ + [ + FrameInfo( + [ + tasks.__file__, + ANY, + "sleep", + ANY, + ] + ), + FrameInfo( + [ + script_name, + ANY, + "echo_client", + ANY, + ] + ), + ], + ANY, + ] + ) + ], + [ + CoroInfo( + [ + [ + FrameInfo( + [ + taskgroups.__file__, + ANY, + "TaskGroup._aexit", + ANY, + ] + ), + FrameInfo( + [ + taskgroups.__file__, + ANY, + "TaskGroup.__aexit__", + ANY, + ] + ), + FrameInfo( + [ + script_name, + ANY, + "echo_client_spam", + ANY, + ] + ), + ], + ANY, + ] + ) + ], + ] + ), + entries, + ) + + # Find tasks awaited by echo_client_spam via TaskGroup + def matches_awaited_by_pattern(task): + if len(task.awaited_by) != 1: + return False + coro = task.awaited_by[0] + if len(coro.call_stack) != 3: + return False + funcnames = [f.funcname for f in coro.call_stack] + return funcnames == [ + "TaskGroup._aexit", + "TaskGroup.__aexit__", + "echo_client_spam", + ] + + tasks_with_awaited = [ + task + for task in entries + if matches_awaited_by_pattern(task) + ] + self.assertGreaterEqual(len(tasks_with_awaited), NUM_TASKS) + + # Final task should be from echo client spam (not on Windows) + if sys.platform != "win32": + self.assertEqual( + tasks_with_awaited[-1].awaited_by, + entries[-1].awaited_by, + ) finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) + _cleanup_sockets(client_socket, server_socket) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_async_global_awaited_by_from_non_main_thread(self): + port = find_unused_port() + script = textwrap.dedent( + f"""\ + import asyncio + import socket + import threading + import time + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + async def worker_main(): + task = asyncio.create_task( + asyncio.sleep(10_000), + name="worker task", + ) + await asyncio.sleep(0) + sock.sendall(f"ready:{{threading.get_native_id()}}\\n".encode()) + await task + + def run_worker_loop(): + asyncio.run(worker_main()) + + threading.Thread( + target=run_worker_loop, + name="async-worker", + daemon=True, + ).start() + time.sleep(10_000) + """ + ) + + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port) + script_name = _make_test_script(script_dir, "script", script) + client_socket = None + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + response = _wait_for_signal(client_socket, b"ready:") + worker_thread_id = int( + response.split(b"ready:", 1)[1].splitlines()[0] + ) + + for _ in busy_retry(SHORT_TIMEOUT): + all_awaited_by = get_all_awaited_by(p.pid) + if any( + task.task_name == "worker task" + for info in all_awaited_by + if info.thread_id == worker_thread_id + for task in info.awaited_by + ): + break + else: + self.fail( + "get_all_awaited_by() did not report " + "the asyncio task from the non-main thread" + ) + finally: + _cleanup_sockets(client_socket, server_socket) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_async_remote_stack_trace_from_non_main_thread(self): + port = find_unused_port() + script = textwrap.dedent( + f"""\ + import asyncio + import socket + import threading + import time + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def blocking_call(): + sock.sendall(f"ready:{{threading.get_native_id()}}\\n".encode()) + time.sleep(10_000) + + async def worker_task(): + await asyncio.sleep(0) + blocking_call() + + async def worker_main(): + task = asyncio.create_task( + worker_task(), + name="worker task", + ) + await task + + def run_worker_loop(): + asyncio.run(worker_main()) + + threading.Thread( + target=run_worker_loop, + name="async-worker", + daemon=True, + ).start() + time.sleep(10_000) + """ + ) + + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port) + script_name = _make_test_script(script_dir, "script", script) + client_socket = None + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + response = _wait_for_signal(client_socket, b"ready:") + worker_thread_id = int( + response.split(b"ready:", 1)[1].splitlines()[0] + ) + + for _ in busy_retry(SHORT_TIMEOUT): + stack_trace = get_async_stack_trace(p.pid) + if any( + task.task_name == "worker task" + for info in stack_trace + if info.thread_id == worker_thread_id + for task in info.awaited_by + ): + break + else: + self.fail( + "get_async_stack_trace() did not report " + "the running asyncio task from the non-main thread" + ) + finally: + _cleanup_sockets(client_socket, server_socket) @skip_if_not_supported @unittest.skipIf( @@ -1126,34 +1666,342 @@ class TestGetStackTrace(unittest.TestCase): ) def test_self_trace(self): stack_trace = get_stack_trace(os.getpid()) - # Is possible that there are more threads, so we check that the - # expected stack traces are in the result (looking at you Windows!) - this_tread_stack = None - for thread_id, stack in stack_trace: - if thread_id == threading.get_native_id(): - this_tread_stack = stack + + this_thread_stack = None + for interpreter_info in stack_trace: + for thread_info in interpreter_info.threads: + if thread_info.thread_id == threading.get_native_id(): + this_thread_stack = thread_info.frame_info + break + if this_thread_stack: break - self.assertIsNotNone(this_tread_stack) - self.assertEqual( - stack[:2], - [ - FrameInfo( - [ - __file__, - get_stack_trace.__code__.co_firstlineno + 2, - "get_stack_trace", - ] - ), - FrameInfo( - [ - __file__, - self.test_self_trace.__code__.co_firstlineno + 6, - "TestGetStackTrace.test_self_trace", - ] - ), - ], + + self.assertIsNotNone(this_thread_stack) + # Check the top two frames + self.assertGreaterEqual(len(this_thread_stack), 2) + self.assertEqual(this_thread_stack[0].funcname, "get_stack_trace") + self.assertTrue(this_thread_stack[0].filename.endswith("test_external_inspection.py")) + self.assertEqual(this_thread_stack[1].funcname, "TestGetStackTrace.test_self_trace") + self.assertTrue(this_thread_stack[1].filename.endswith("test_external_inspection.py")) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + @requires_subinterpreters + def test_subinterpreter_stack_trace(self): + port = find_unused_port() + + import pickle + + subinterp_code = textwrap.dedent(f""" + import socket + import time + + def sub_worker(): + def nested_func(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + sock.sendall(b"ready:sub\\n") + time.sleep(10_000) + nested_func() + + sub_worker() + """).strip() + + pickled_code = pickle.dumps(subinterp_code) + + script = textwrap.dedent( + f""" + from concurrent import interpreters + import time + import sys + import socket + import threading + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def main_worker(): + sock.sendall(b"ready:main\\n") + time.sleep(10_000) + + def run_subinterp(): + subinterp = interpreters.create() + import pickle + pickled_code = {pickled_code!r} + subinterp_code = pickle.loads(pickled_code) + subinterp.exec(subinterp_code) + + sub_thread = threading.Thread(target=run_subinterp) + sub_thread.start() + + main_thread = threading.Thread(target=main_worker) + main_thread.start() + + main_thread.join() + sub_thread.join() + """ ) + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port) + script_name = _make_test_script(script_dir, "script", script) + client_sockets = [] + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + # Accept connections from both main and subinterpreter + responses = set() + while len(responses) < 2: + try: + client_socket, _ = server_socket.accept() + client_sockets.append(client_socket) + response = client_socket.recv(1024) + if b"ready:main" in response: + responses.add("main") + if b"ready:sub" in response: + responses.add("sub") + except socket.timeout: + break + + server_socket.close() + server_socket = None + + stack_trace = get_stack_trace(p.pid) + + # Verify we have at least one interpreter + self.assertGreaterEqual(len(stack_trace), 1) + + # Look for main interpreter (ID 0) and subinterpreter (ID > 0) + main_interp = None + sub_interp = None + for interpreter_info in stack_trace: + if interpreter_info.interpreter_id == 0: + main_interp = interpreter_info + elif interpreter_info.interpreter_id > 0: + sub_interp = interpreter_info + + self.assertIsNotNone( + main_interp, "Main interpreter should be present" + ) + + # Check main interpreter has expected stack trace + main_found = self._find_frame_in_trace( + [main_interp], lambda f: f.funcname == "main_worker" + ) + self.assertIsNotNone( + main_found, + "Main interpreter should have main_worker in stack", + ) + + # If subinterpreter is present, check its stack trace + if sub_interp: + sub_found = self._find_frame_in_trace( + [sub_interp], + lambda f: f.funcname + in ("sub_worker", "nested_func"), + ) + self.assertIsNotNone( + sub_found, + "Subinterpreter should have sub_worker or nested_func in stack", + ) + finally: + _cleanup_sockets(*client_sockets, server_socket) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + @requires_subinterpreters + def test_multiple_subinterpreters_with_threads(self): + port = find_unused_port() + + import pickle + + subinterp1_code = textwrap.dedent(f""" + import socket + import time + import threading + + def worker1(): + def nested_func(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + sock.sendall(b"ready:sub1-t1\\n") + time.sleep(10_000) + nested_func() + + def worker2(): + def nested_func(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + sock.sendall(b"ready:sub1-t2\\n") + time.sleep(10_000) + nested_func() + + t1 = threading.Thread(target=worker1) + t2 = threading.Thread(target=worker2) + t1.start() + t2.start() + t1.join() + t2.join() + """).strip() + + subinterp2_code = textwrap.dedent(f""" + import socket + import time + import threading + + def worker1(): + def nested_func(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + sock.sendall(b"ready:sub2-t1\\n") + time.sleep(10_000) + nested_func() + + def worker2(): + def nested_func(): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + sock.sendall(b"ready:sub2-t2\\n") + time.sleep(10_000) + nested_func() + + t1 = threading.Thread(target=worker1) + t2 = threading.Thread(target=worker2) + t1.start() + t2.start() + t1.join() + t2.join() + """).strip() + + pickled_code1 = pickle.dumps(subinterp1_code) + pickled_code2 = pickle.dumps(subinterp2_code) + + script = textwrap.dedent( + f""" + from concurrent import interpreters + import time + import sys + import socket + import threading + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def main_worker(): + sock.sendall(b"ready:main\\n") + time.sleep(10_000) + + def run_subinterp1(): + subinterp = interpreters.create() + import pickle + pickled_code = {pickled_code1!r} + subinterp_code = pickle.loads(pickled_code) + subinterp.exec(subinterp_code) + + def run_subinterp2(): + subinterp = interpreters.create() + import pickle + pickled_code = {pickled_code2!r} + subinterp_code = pickle.loads(pickled_code) + subinterp.exec(subinterp_code) + + sub1_thread = threading.Thread(target=run_subinterp1) + sub2_thread = threading.Thread(target=run_subinterp2) + sub1_thread.start() + sub2_thread.start() + + main_thread = threading.Thread(target=main_worker) + main_thread.start() + + main_thread.join() + sub1_thread.join() + sub2_thread.join() + """ + ) + + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port, backlog=5) + script_name = _make_test_script(script_dir, "script", script) + client_sockets = [] + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + # Accept connections from main and all subinterpreter threads + expected_responses = { + "ready:main", + "ready:sub1-t1", + "ready:sub1-t2", + "ready:sub2-t1", + "ready:sub2-t2", + } + responses = set() + + while len(responses) < 5: + try: + client_socket, _ = server_socket.accept() + client_sockets.append(client_socket) + response = client_socket.recv(1024) + response_str = response.decode().strip() + if response_str in expected_responses: + responses.add(response_str) + except socket.timeout: + break + + server_socket.close() + server_socket = None + + stack_trace = get_stack_trace(p.pid) + + # Verify we have multiple interpreters + self.assertGreaterEqual(len(stack_trace), 2) + + # Count interpreters by ID + interpreter_ids = { + interp.interpreter_id for interp in stack_trace + } + self.assertIn( + 0, + interpreter_ids, + "Main interpreter should be present", + ) + self.assertGreaterEqual(len(interpreter_ids), 3) + + # Count total threads + total_threads = sum( + len(interp.threads) for interp in stack_trace + ) + self.assertGreaterEqual(total_threads, 5) + + # Look for expected function names + all_funcnames = set() + for interpreter_info in stack_trace: + for thread_info in interpreter_info.threads: + for frame in thread_info.frame_info: + all_funcnames.add(frame.funcname) + + expected_funcs = { + "main_worker", + "worker1", + "worker2", + "nested_func", + } + found_funcs = expected_funcs.intersection(all_funcnames) + self.assertGreater(len(found_funcs), 0) + finally: + _cleanup_sockets(*client_sockets, server_socket) + @skip_if_not_supported @unittest.skipIf( sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, @@ -1161,54 +2009,41 @@ class TestGetStackTrace(unittest.TestCase): ) @requires_gil_enabled("Free threaded builds don't have an 'active thread'") def test_only_active_thread(self): - # Test that only_active_thread parameter works correctly port = find_unused_port() script = textwrap.dedent( f"""\ import time, sys, socket, threading - # Connect to the test process sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.connect(('localhost', {port})) def worker_thread(name, barrier, ready_event): - barrier.wait() # Synchronize thread start - ready_event.wait() # Wait for main thread signal - # Sleep to keep thread alive + barrier.wait() + ready_event.wait() time.sleep(10_000) def main_work(): - # Do busy work to hold the GIL sock.sendall(b"working\\n") count = 0 while count < 100000000: count += 1 if count % 10000000 == 0: - pass # Keep main thread busy + pass sock.sendall(b"done\\n") - # Create synchronization primitives num_threads = 3 - barrier = threading.Barrier(num_threads + 1) # +1 for main thread + barrier = threading.Barrier(num_threads + 1) ready_event = threading.Event() - # Start worker threads threads = [] for i in range(num_threads): t = threading.Thread(target=worker_thread, args=(f"Worker-{{i}}", barrier, ready_event)) t.start() threads.append(t) - # Wait for all threads to be ready barrier.wait() - - # Signal ready to parent process sock.sendall(b"ready\\n") - - # Signal threads to start waiting ready_event.set() - - # Now do busy work to hold the GIL main_work() """ ) @@ -1217,87 +2052,1957 @@ class TestGetStackTrace(unittest.TestCase): script_dir = os.path.join(work_dir, "script_pkg") os.mkdir(script_dir) - # Create a socket server to communicate with the target process - server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - server_socket.bind(("localhost", port)) - server_socket.settimeout(SHORT_TIMEOUT) - server_socket.listen(1) - + server_socket = _create_server_socket(port) script_name = _make_test_script(script_dir, "script", script) client_socket = None + try: - p = subprocess.Popen([sys.executable, script_name]) - client_socket, _ = server_socket.accept() - server_socket.close() + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None - # Wait for ready signal - response = b"" - while b"ready" not in response: - response += client_socket.recv(1024) + # Wait for ready and working signals + _wait_for_signal(client_socket, [b"ready", b"working"]) - # Wait for the main thread to start its busy work - while b"working" not in response: - response += client_socket.recv(1024) + # Get stack trace with all threads + unwinder_all = RemoteUnwinder(p.pid, all_threads=True) + for _ in busy_retry(SHORT_TIMEOUT): + with contextlib.suppress(*TRANSIENT_ERRORS): + all_traces = unwinder_all.get_stack_trace() + found = self._find_frame_in_trace( + all_traces, + lambda f: f.funcname == "main_work" + and f.location.lineno > 12, + ) + if found: + break + else: + self.fail( + "Main thread did not start its busy work on time" + ) - # Get stack trace with all threads - unwinder_all = RemoteUnwinder(p.pid, all_threads=True) - for _ in range(10): - # Wait for the main thread to start its busy work - all_traces = unwinder_all.get_stack_trace() - found = False - for thread_id, stack in all_traces: - if not stack: - continue - current_frame = stack[0] - if ( - current_frame.funcname == "main_work" - and current_frame.lineno > 15 - ): - found = True - - if found: - break - # Give a bit of time to take the next sample - time.sleep(0.1) - else: - self.fail( - "Main thread did not start its busy work on time" + # Get stack trace with only GIL holder + unwinder_gil = RemoteUnwinder( + p.pid, only_active_thread=True + ) + # Use condition to retry until we capture a thread holding the GIL + # (sampling may catch moments with no GIL holder on slow CI) + gil_traces = _get_stack_trace_with_retry( + unwinder_gil, + condition=lambda t: sum(len(i.threads) for i in t) >= 1, ) - # Get stack trace with only GIL holder - unwinder_gil = RemoteUnwinder(p.pid, only_active_thread=True) - gil_traces = unwinder_gil.get_stack_trace() + # Count threads + total_threads = sum( + len(interp.threads) for interp in all_traces + ) + self.assertGreater(total_threads, 1) - except PermissionError: - self.skipTest( - "Insufficient permissions to read the stack trace" - ) + total_gil_threads = sum( + len(interp.threads) for interp in gil_traces + ) + self.assertEqual(total_gil_threads, 1) + + # Get the GIL holder thread ID + gil_thread_id = None + for interpreter_info in gil_traces: + if interpreter_info.threads: + gil_thread_id = interpreter_info.threads[ + 0 + ].thread_id + break + + # Get all thread IDs + all_thread_ids = [] + for interpreter_info in all_traces: + for thread_info in interpreter_info.threads: + all_thread_ids.append(thread_info.thread_id) + + self.assertIn(gil_thread_id, all_thread_ids) finally: - if client_socket is not None: - client_socket.close() - p.kill() - p.terminate() - p.wait(timeout=SHORT_TIMEOUT) + _cleanup_sockets(client_socket, server_socket) - # Verify we got multiple threads in all_traces - self.assertGreater( - len(all_traces), 1, "Should have multiple threads" + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_opcodes_collection(self): + """Test that opcodes are collected when the opcodes flag is set.""" + script = textwrap.dedent( + """\ + import time, sys, socket + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def foo(): + sock.sendall(b"ready") + time.sleep(10_000) + + foo() + """ + ) + + def get_trace_with_opcodes(pid): + return RemoteUnwinder(pid, opcodes=True).get_stack_trace() + + stack_trace, _ = self._run_script_and_get_trace( + script, get_trace_with_opcodes, wait_for_signals=b"ready" + ) + + # Find our foo frame and verify it has an opcode + foo_frame = self._find_frame_in_trace( + stack_trace, lambda f: f.funcname == "foo" + ) + self.assertIsNotNone(foo_frame, "Could not find foo frame") + self.assertIsInstance(foo_frame.opcode, int) + self.assertGreaterEqual(foo_frame.opcode, 0) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_location_tuple_format(self): + """Test that location is a 4-tuple (lineno, end_lineno, col_offset, end_col_offset).""" + script = textwrap.dedent( + """\ + import time, sys, socket + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def foo(): + sock.sendall(b"ready") + time.sleep(10_000) + + foo() + """ + ) + + def get_trace_with_opcodes(pid): + return RemoteUnwinder(pid, opcodes=True).get_stack_trace() + + stack_trace, _ = self._run_script_and_get_trace( + script, get_trace_with_opcodes, wait_for_signals=b"ready" + ) + + # Find our foo frame + foo_frame = self._find_frame_in_trace( + stack_trace, lambda f: f.funcname == "foo" + ) + self.assertIsNotNone(foo_frame, "Could not find foo frame") + + # Check location is a 4-tuple with valid values + location = foo_frame.location + self.assertIsInstance(location, tuple) + self.assertEqual(len(location), 4) + lineno, end_lineno, col_offset, end_col_offset = location + self.assertIsInstance(lineno, int) + self.assertGreater(lineno, 0) + self.assertIsInstance(end_lineno, int) + self.assertGreaterEqual(end_lineno, lineno) + self.assertIsInstance(col_offset, int) + self.assertGreaterEqual(col_offset, 0) + self.assertIsInstance(end_col_offset, int) + self.assertGreaterEqual(end_col_offset, col_offset) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_location_tuple_exact_values(self): + """Test exact values of location tuple including column offsets.""" + script = textwrap.dedent( + """\ + import time, sys, socket + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def foo(): + sock.sendall(b"ready") + time.sleep(10_000) + + foo() + """ + ) + + def get_trace_with_opcodes(pid): + return RemoteUnwinder(pid, opcodes=True).get_stack_trace() + + stack_trace, _ = self._run_script_and_get_trace( + script, get_trace_with_opcodes, wait_for_signals=b"ready" + ) + + foo_frame = self._find_frame_in_trace( + stack_trace, lambda f: f.funcname == "foo" + ) + self.assertIsNotNone(foo_frame, "Could not find foo frame") + + # Can catch either sock.sendall (line 7) or time.sleep (line 8) + location = foo_frame.location + valid_locations = [ + (7, 7, 4, 26), # sock.sendall(b"ready") + (8, 8, 4, 22), # time.sleep(10_000) + ] + actual = (location.lineno, location.end_lineno, + location.col_offset, location.end_col_offset) + self.assertIn(actual, valid_locations) + + @skip_if_not_supported + @unittest.skipIf(sys._is_gil_enabled(), "Requires free-threading") + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Requires process_vm_readv", + ) + def test_tlbc_cache_refresh_after_growth(self): + # Reproducer from gh-157660. + script = textwrap.dedent("""\ + import os, threading + from _remote_debugging import RemoteUnwinder + from _queue import SimpleQueue + from test import support + + go = threading.Lock() + stop = threading.Lock() + go.acquire() + stop.acquire() + ready = SimpleQueue() + + def leaf(): + ready.put(None) + stop.acquire() + + def start_leaf(): + ready.put(None) + go.acquire() + leaf() + + def park(): + ready.put(None) + stop.acquire() + + def leaf_count(u): + return sum( + f.funcname == "leaf" + for i in u.get_stack_trace() + for t in i.threads for f in t.frame_info + ) + + # SimpleQueue.put() and Lock.acquire() do not push Python frames. + # Once notified, the worker's stack stays stable until go is released. + threading.Thread(target=leaf, daemon=True).start() + ready.get(timeout=support.SHORT_TIMEOUT) + for _ in range(16): + threading.Thread(target=park, daemon=True).start() + ready.get(timeout=support.SHORT_TIMEOUT) + threading.Thread(target=start_leaf, daemon=True).start() + ready.get(timeout=support.SHORT_TIMEOUT) + + u = RemoteUnwinder(os.getpid(), all_threads=True, cache_frames=False) + assert leaf_count(u) == 1 + go.release() + ready.get(timeout=support.SHORT_TIMEOUT) + assert leaf_count(u) == 2 + """) + result = subprocess.run( + [sys.executable, "-X", "gil=0", "-X", "tlbc=1", "-c", script], + capture_output=True, + text=True, + timeout=SHORT_TIMEOUT, + ) + self.assertEqual( + result.returncode, 0, + f"stdout: {result.stdout}\nstderr: {result.stderr}", + ) + + @skip_if_not_supported + @unittest.skipIf(sys._is_gil_enabled(), "Requires free-threading") + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Requires process_vm_readv", + ) + def test_tlbc_cache_refresh_after_slot_fill(self): + # Reproducer from gh-157660. + script = textwrap.dedent("""\ + import os, threading + from _remote_debugging import RemoteUnwinder + from _queue import SimpleQueue + + go = threading.Lock() + stop = threading.Lock() + go.acquire() + stop.acquire() + ready = SimpleQueue() + + def leaf(): + ready.put(None) + stop.acquire() + + def start_leaf(): + ready.put(None) + go.acquire() + leaf() + + from test import support + + def lines(u): + return sorted( + f.location.lineno + for i in u.get_stack_trace() + for t in i.threads for f in t.frame_info + if f.funcname == "leaf" + ) + + # SimpleQueue.put() and Lock.acquire() do not push Python frames. + # Once notified, the worker's stack stays stable until go is released. + threading.Thread(target=leaf, daemon=True).start() + ready.get(timeout=support.SHORT_TIMEOUT) + threading.Thread(target=start_leaf, daemon=True).start() + ready.get(timeout=support.SHORT_TIMEOUT) + u = RemoteUnwinder(os.getpid(), all_threads=True, cache_frames=False) + before = lines(u) + # The notification can be observed before put() returns, so either + # line in leaf() is a valid sample. + assert before in ([12], [13]), before + go.release() + ready.get(timeout=support.SHORT_TIMEOUT) + cached = lines(u) + assert len(cached) == 2, cached + assert all(line in (12, 13) for line in cached), cached + """) + result = subprocess.run( + [sys.executable, "-X", "gil=0", "-X", "tlbc=1", "-c", script], + capture_output=True, + text=True, + timeout=SHORT_TIMEOUT, + ) + self.assertEqual( + result.returncode, 0, + f"stdout: {result.stdout}\nstderr: {result.stderr}", + ) + + +class TestUnsupportedPlatformHandling(unittest.TestCase): + @unittest.skipIf( + sys.platform in ("linux", "darwin", "win32"), + "Test only runs on unsupported platforms (not Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_unsupported_platform_error(self): + with self.assertRaises(RuntimeError) as cm: + RemoteUnwinder(os.getpid()) + + self.assertIn( + "Reading the PyRuntime section is not supported on this platform", + str(cm.exception), + ) + + +@requires_remote_subprocess_debugging() +class TestDetectionOfThreadStatus(RemoteInspectionTestBase): + def _run_thread_status_test(self, mode, check_condition): + """ + Common pattern for thread status detection tests. + + Args: + mode: Profiling mode (PROFILING_MODE_CPU, PROFILING_MODE_GIL, etc.) + check_condition: Function(statuses, sleeper_tid, busy_tid) -> bool + """ + port = find_unused_port() + script = textwrap.dedent( + f"""\ + import time, sys, socket, threading + import os + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(('localhost', {port})) + + def sleeper(): + tid = threading.get_native_id() + sock.sendall(f'ready:sleeper:{{tid}}\\n'.encode()) + time.sleep(10000) + + def busy(): + tid = threading.get_native_id() + sock.sendall(f'ready:busy:{{tid}}\\n'.encode()) + x = 0 + while True: + x = x + 1 + time.sleep(0.5) + + t1 = threading.Thread(target=sleeper) + t2 = threading.Thread(target=busy) + t1.start() + t2.start() + sock.sendall(b'ready:main\\n') + t1.join() + t2.join() + sock.close() + """ + ) + + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port) + script_name = _make_test_script( + script_dir, "thread_status_script", script + ) + client_socket = None + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + # Wait for all ready signals and parse TIDs + response = _wait_for_signal( + client_socket, + [b"ready:main", b"ready:sleeper", b"ready:busy"], + ) + + sleeper_tid = None + busy_tid = None + for line in response.split(b"\n"): + if line.startswith(b"ready:sleeper:"): + try: + sleeper_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + elif line.startswith(b"ready:busy:"): + try: + busy_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + + self.assertIsNotNone( + sleeper_tid, "Sleeper thread id not received" + ) + self.assertIsNotNone( + busy_tid, "Busy thread id not received" + ) + + # Sample until we see expected thread states + statuses = {} + unwinder = RemoteUnwinder( + p.pid, + all_threads=True, + mode=mode, + skip_non_matching_threads=False, + ) + for _ in busy_retry(SHORT_TIMEOUT): + with contextlib.suppress(*TRANSIENT_ERRORS): + traces = unwinder.get_stack_trace() + statuses = self._get_thread_statuses(traces) + + if check_condition( + statuses, sleeper_tid, busy_tid + ): + break + + return statuses, sleeper_tid, busy_tid + finally: + _cleanup_sockets(client_socket, server_socket) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_thread_status_detection(self): + def check_cpu_status(statuses, sleeper_tid, busy_tid): + return ( + sleeper_tid in statuses + and busy_tid in statuses + and not (statuses[sleeper_tid] & THREAD_STATUS_ON_CPU) + and (statuses[busy_tid] & THREAD_STATUS_ON_CPU) ) - # Verify we got exactly one thread in gil_traces - self.assertEqual( - len(gil_traces), 1, "Should have exactly one GIL holder" + statuses, sleeper_tid, busy_tid = self._run_thread_status_test( + PROFILING_MODE_CPU, check_cpu_status + ) + + self.assertIn(sleeper_tid, statuses) + self.assertIn(busy_tid, statuses) + self.assertFalse( + statuses[sleeper_tid] & THREAD_STATUS_ON_CPU, + "Sleeper thread should be off CPU", + ) + self.assertTrue( + statuses[busy_tid] & THREAD_STATUS_ON_CPU, + "Busy thread should be on CPU", + ) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_thread_status_gil_detection(self): + def check_gil_status(statuses, sleeper_tid, busy_tid): + return ( + sleeper_tid in statuses + and busy_tid in statuses + and not (statuses[sleeper_tid] & THREAD_STATUS_HAS_GIL) + and (statuses[busy_tid] & THREAD_STATUS_HAS_GIL) ) - # The GIL holder should be in the all_traces list - gil_thread_id = gil_traces[0][0] - all_thread_ids = [trace[0] for trace in all_traces] - self.assertIn( - gil_thread_id, - all_thread_ids, - "GIL holder should be among all threads", + statuses, sleeper_tid, busy_tid = self._run_thread_status_test( + PROFILING_MODE_GIL, check_gil_status + ) + + self.assertIn(sleeper_tid, statuses) + self.assertIn(busy_tid, statuses) + self.assertFalse( + statuses[sleeper_tid] & THREAD_STATUS_HAS_GIL, + "Sleeper thread should not have GIL", + ) + self.assertTrue( + statuses[busy_tid] & THREAD_STATUS_HAS_GIL, + "Busy thread should have GIL", + ) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_thread_status_all_mode_detection(self): + port = find_unused_port() + script = textwrap.dedent( + f"""\ + import socket + import threading + import time + import sys + + def sleeper_thread(): + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"sleeper:" + str(threading.get_native_id()).encode()) + while True: + time.sleep(1) + + def busy_thread(): + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"busy:" + str(threading.get_native_id()).encode()) + while True: + sum(range(100000)) + + t1 = threading.Thread(target=sleeper_thread) + t2 = threading.Thread(target=busy_thread) + t1.start() + t2.start() + t1.join() + t2.join() + """ + ) + + with os_helper.temp_dir() as tmp_dir: + script_file = make_script(tmp_dir, "script", script) + server_socket = _create_server_socket(port, backlog=2) + client_sockets = [] + + try: + with _managed_subprocess( + [sys.executable, script_file], + ) as p: + sleeper_tid = None + busy_tid = None + + # Receive thread IDs from the child process + for _ in range(2): + client_socket, _ = server_socket.accept() + client_sockets.append(client_socket) + line = client_socket.recv(1024) + if line: + if line.startswith(b"sleeper:"): + try: + sleeper_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + elif line.startswith(b"busy:"): + try: + busy_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + + server_socket.close() + server_socket = None + + statuses = {} + unwinder = RemoteUnwinder( + p.pid, + all_threads=True, + mode=PROFILING_MODE_ALL, + skip_non_matching_threads=False, + ) + for _ in busy_retry(SHORT_TIMEOUT): + with contextlib.suppress(*TRANSIENT_ERRORS): + traces = unwinder.get_stack_trace() + statuses = self._get_thread_statuses(traces) + + # Check ALL mode provides both GIL and CPU info + if ( + sleeper_tid in statuses + and busy_tid in statuses + and not ( + statuses[sleeper_tid] + & THREAD_STATUS_ON_CPU + ) + and not ( + statuses[sleeper_tid] + & THREAD_STATUS_HAS_GIL + ) + and (statuses[busy_tid] & THREAD_STATUS_ON_CPU) + and ( + statuses[busy_tid] & THREAD_STATUS_HAS_GIL + ) + ): + break + + self.assertIsNotNone( + sleeper_tid, "Sleeper thread id not received" + ) + self.assertIsNotNone( + busy_tid, "Busy thread id not received" + ) + self.assertIn(sleeper_tid, statuses) + self.assertIn(busy_tid, statuses) + + # Sleeper: off CPU, no GIL + self.assertFalse( + statuses[sleeper_tid] & THREAD_STATUS_ON_CPU, + "Sleeper should be off CPU", + ) + self.assertFalse( + statuses[sleeper_tid] & THREAD_STATUS_HAS_GIL, + "Sleeper should not have GIL", + ) + + # Busy: on CPU, has GIL + self.assertTrue( + statuses[busy_tid] & THREAD_STATUS_ON_CPU, + "Busy should be on CPU", + ) + self.assertTrue( + statuses[busy_tid] & THREAD_STATUS_HAS_GIL, + "Busy should have GIL", + ) + finally: + _cleanup_sockets(*client_sockets, server_socket) + + def _make_exception_test_script(self, port): + """Create script with exception and normal threads for testing.""" + return textwrap.dedent( + f"""\ + import socket + import threading + import time + + def exception_thread(): + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"exception:" + str(threading.get_native_id()).encode()) + try: + raise ValueError("test exception") + except ValueError: + while True: + time.sleep(0.01) + + def normal_thread(): + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"normal:" + str(threading.get_native_id()).encode()) + while True: + sum(range(1000)) + + t1 = threading.Thread(target=exception_thread) + t2 = threading.Thread(target=normal_thread) + t1.start() + t2.start() + t1.join() + t2.join() + """ + ) + + @contextmanager + def _run_exception_test_process(self): + """Context manager to run exception test script and yield thread IDs and process.""" + port = find_unused_port() + script = self._make_exception_test_script(port) + + with os_helper.temp_dir() as tmp_dir: + script_file = make_script(tmp_dir, "script", script) + server_socket = _create_server_socket(port, backlog=2) + client_sockets = [] + + try: + with _managed_subprocess([sys.executable, script_file]) as p: + exception_tid = None + normal_tid = None + + for _ in range(2): + client_socket, _ = server_socket.accept() + client_sockets.append(client_socket) + line = client_socket.recv(1024) + if line: + if line.startswith(b"exception:"): + try: + exception_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + elif line.startswith(b"normal:"): + try: + normal_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + + server_socket.close() + server_socket = None + + yield p, exception_tid, normal_tid + finally: + _cleanup_sockets(*client_sockets, server_socket) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_thread_status_exception_detection(self): + """Test that THREAD_STATUS_HAS_EXCEPTION is set when thread has an active exception.""" + with self._run_exception_test_process() as (p, exception_tid, normal_tid): + self.assertIsNotNone(exception_tid, "Exception thread id not received") + self.assertIsNotNone(normal_tid, "Normal thread id not received") + + statuses = {} + unwinder = RemoteUnwinder( + p.pid, + all_threads=True, + mode=PROFILING_MODE_ALL, + skip_non_matching_threads=False, ) + for _ in busy_retry(SHORT_TIMEOUT): + with contextlib.suppress(*TRANSIENT_ERRORS): + traces = unwinder.get_stack_trace() + statuses = self._get_thread_statuses(traces) + + if ( + exception_tid in statuses + and normal_tid in statuses + and (statuses[exception_tid] & THREAD_STATUS_HAS_EXCEPTION) + and not (statuses[normal_tid] & THREAD_STATUS_HAS_EXCEPTION) + ): + break + + self.assertIn(exception_tid, statuses) + self.assertIn(normal_tid, statuses) + self.assertTrue( + statuses[exception_tid] & THREAD_STATUS_HAS_EXCEPTION, + "Exception thread should have HAS_EXCEPTION flag", + ) + self.assertFalse( + statuses[normal_tid] & THREAD_STATUS_HAS_EXCEPTION, + "Normal thread should not have HAS_EXCEPTION flag", + ) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_thread_status_exception_mode_filtering(self): + """Test that PROFILING_MODE_EXCEPTION correctly filters threads.""" + with self._run_exception_test_process() as (p, exception_tid, normal_tid): + self.assertIsNotNone(exception_tid, "Exception thread id not received") + self.assertIsNotNone(normal_tid, "Normal thread id not received") + + unwinder = RemoteUnwinder( + p.pid, + all_threads=True, + mode=PROFILING_MODE_EXCEPTION, + skip_non_matching_threads=True, + ) + for _ in busy_retry(SHORT_TIMEOUT): + with contextlib.suppress(*TRANSIENT_ERRORS): + traces = unwinder.get_stack_trace() + statuses = self._get_thread_statuses(traces) + + if exception_tid in statuses: + self.assertNotIn( + normal_tid, + statuses, + "Normal thread should be filtered out in exception mode", + ) + return + + self.fail("Never found exception thread in exception mode") + +@requires_remote_subprocess_debugging() +class TestExceptionDetectionScenarios(RemoteInspectionTestBase): + """Test exception detection across all scenarios. + + This class verifies the exact conditions under which THREAD_STATUS_HAS_EXCEPTION + is set. Each test covers a specific scenario: + + 1. except_block: Thread inside except block + -> SHOULD have HAS_EXCEPTION (exc_info->exc_value is set) + + 2. finally_propagating: Exception propagating through finally block + -> SHOULD have HAS_EXCEPTION (current_exception is set) + + 3. finally_after_except: Finally block after except handled exception + -> Should NOT have HAS_EXCEPTION (exc_info cleared after except) + + 4. finally_no_exception: Finally block with no exception raised + -> Should NOT have HAS_EXCEPTION (no exception state) + """ + + def _make_single_scenario_script(self, port, scenario): + """Create script for a single exception scenario.""" + scenarios = { + "except_block": f"""\ +import socket +import threading +import time + +def target_thread(): + '''Inside except block - exception info is present''' + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"ready:" + str(threading.get_native_id()).encode()) + try: + raise ValueError("test") + except ValueError: + while True: + time.sleep(0.01) + +t = threading.Thread(target=target_thread) +t.start() +t.join() +""", + "finally_propagating": f"""\ +import socket +import threading +import time + +def target_thread(): + '''Exception propagating through finally - current_exception is set''' + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"ready:" + str(threading.get_native_id()).encode()) + try: + try: + raise ValueError("propagating") + finally: + # Exception is propagating through here + while True: + time.sleep(0.01) + except: + pass # Never reached due to infinite loop + +t = threading.Thread(target=target_thread) +t.start() +t.join() +""", + "finally_after_except": f"""\ +import socket +import threading +import time + +def target_thread(): + '''Finally runs after except handled - exc_info is cleared''' + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"ready:" + str(threading.get_native_id()).encode()) + try: + raise ValueError("test") + except ValueError: + pass # Exception caught and handled + finally: + while True: + time.sleep(0.01) + +t = threading.Thread(target=target_thread) +t.start() +t.join() +""", + "finally_no_exception": f"""\ +import socket +import threading +import time + +def target_thread(): + '''Finally with no exception at all''' + conn = socket.create_connection(("localhost", {port})) + conn.sendall(b"ready:" + str(threading.get_native_id()).encode()) + try: + pass # No exception + finally: + while True: + time.sleep(0.01) + +t = threading.Thread(target=target_thread) +t.start() +t.join() +""", + } + + return scenarios[scenario] + + @contextmanager + def _run_scenario_process(self, scenario): + """Context manager to run a single scenario and yield thread ID and process.""" + port = find_unused_port() + script = self._make_single_scenario_script(port, scenario) + + with os_helper.temp_dir() as tmp_dir: + script_file = make_script(tmp_dir, "script", script) + server_socket = _create_server_socket(port, backlog=1) + client_socket = None + + try: + with _managed_subprocess([sys.executable, script_file]) as p: + thread_tid = None + + client_socket, _ = server_socket.accept() + line = client_socket.recv(1024) + if line and line.startswith(b"ready:"): + try: + thread_tid = int(line.split(b":")[-1]) + except (ValueError, IndexError): + pass + + server_socket.close() + server_socket = None + + yield p, thread_tid + finally: + _cleanup_sockets(client_socket, server_socket) + + def _check_thread_status( + self, p, thread_tid, condition, condition_name="condition" + ): + """Helper to check thread status with a custom condition. + + This waits until we see 3 consecutive samples where the condition + returns True, which confirms the thread has reached and is stable + in the expected state. Samples that don't match are ignored (the + thread may not have reached the expected state yet). + + Args: + p: Process object with pid attribute + thread_tid: Thread ID to check + condition: Callable(statuses, thread_tid) -> bool that returns + True when the thread is in the expected state + condition_name: Description of condition for error messages + """ + unwinder = RemoteUnwinder( + p.pid, + all_threads=True, + mode=PROFILING_MODE_ALL, + skip_non_matching_threads=False, + ) + + # Wait for 3 consecutive samples matching expected state + matching_samples = 0 + for _ in busy_retry(SHORT_TIMEOUT): + with contextlib.suppress(*TRANSIENT_ERRORS): + traces = unwinder.get_stack_trace() + statuses = self._get_thread_statuses(traces) + + if thread_tid in statuses: + if condition(statuses, thread_tid): + matching_samples += 1 + if matching_samples >= 3: + return # Success - confirmed stable in expected state + else: + # Thread not yet in expected state, reset counter + matching_samples = 0 + + self.fail( + f"Thread did not stabilize in expected state " + f"({condition_name}) within timeout" + ) + + def _check_exception_status(self, p, thread_tid, expect_exception): + """Helper to check if thread has expected exception status.""" + def condition(statuses, tid): + has_exc = bool(statuses[tid] & THREAD_STATUS_HAS_EXCEPTION) + return has_exc == expect_exception + + self._check_thread_status( + p, thread_tid, condition, + condition_name=f"expect_exception={expect_exception}" + ) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_except_block_has_exception(self): + """Test that thread inside except block has HAS_EXCEPTION flag. + + When a thread is executing inside an except block, exc_info->exc_value + is set, so THREAD_STATUS_HAS_EXCEPTION should be True. + """ + with self._run_scenario_process("except_block") as (p, thread_tid): + self.assertIsNotNone(thread_tid, "Thread ID not received") + self._check_exception_status(p, thread_tid, expect_exception=True) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_finally_propagating_has_exception(self): + """Test that finally block with propagating exception has HAS_EXCEPTION flag. + + When an exception is propagating through a finally block (not yet caught), + current_exception is set, so THREAD_STATUS_HAS_EXCEPTION should be True. + """ + with self._run_scenario_process("finally_propagating") as (p, thread_tid): + self.assertIsNotNone(thread_tid, "Thread ID not received") + self._check_exception_status(p, thread_tid, expect_exception=True) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_finally_after_except_no_exception(self): + """Test that finally block after except has NO HAS_EXCEPTION flag. + + When a finally block runs after an except block has handled the exception, + Python clears exc_info before entering finally, so THREAD_STATUS_HAS_EXCEPTION + should be False. + """ + with self._run_scenario_process("finally_after_except") as (p, thread_tid): + self.assertIsNotNone(thread_tid, "Thread ID not received") + self._check_exception_status(p, thread_tid, expect_exception=False) + + @unittest.skipIf( + sys.platform not in ("linux", "darwin", "win32"), + "Test only runs on supported platforms (Linux, macOS, or Windows)", + ) + @unittest.skipIf( + sys.platform == "android", "Android raises Linux-specific exception" + ) + def test_finally_no_exception_no_flag(self): + """Test that finally block with no exception has NO HAS_EXCEPTION flag. + + When a finally block runs during normal execution (no exception raised), + there is no exception state, so THREAD_STATUS_HAS_EXCEPTION should be False. + """ + with self._run_scenario_process("finally_no_exception") as (p, thread_tid): + self.assertIsNotNone(thread_tid, "Thread ID not received") + self._check_exception_status(p, thread_tid, expect_exception=False) + + +@requires_remote_subprocess_debugging() +class TestFrameCaching(RemoteInspectionTestBase): + """Test that frame caching produces correct results. + + Uses socket-based synchronization for deterministic testing. + All tests verify cache reuse via object identity checks (assertIs). + """ + + @contextmanager + def _target_process(self, script_body): + """Context manager for running a target process with socket sync.""" + port = find_unused_port() + script = f"""\ +import socket +sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) +sock.connect(('localhost', {port})) +{textwrap.dedent(script_body)} +""" + + with os_helper.temp_dir() as work_dir: + script_dir = os.path.join(work_dir, "script_pkg") + os.mkdir(script_dir) + + server_socket = _create_server_socket(port) + script_name = _make_test_script(script_dir, "script", script) + client_socket = None + + try: + with _managed_subprocess([sys.executable, script_name]) as p: + client_socket, _ = server_socket.accept() + server_socket.close() + server_socket = None + + def make_unwinder(cache_frames=True): + return RemoteUnwinder( + p.pid, all_threads=True, cache_frames=cache_frames + ) + + yield p, client_socket, make_unwinder + finally: + _cleanup_sockets(client_socket, server_socket) + + def _get_frames_with_retry(self, unwinder, required_funcs): + """Get frames containing required_funcs, with retry for transient errors.""" + for _ in range(MAX_TRIES): + with contextlib.suppress(*TRANSIENT_ERRORS): + traces = unwinder.get_stack_trace() + for interp in traces: + for thread in interp.threads: + funcs = {f.funcname for f in thread.frame_info} + if required_funcs.issubset(funcs): + return thread.frame_info + time.sleep(RETRY_DELAY) + return None + + def _sample_frames( + self, + client_socket, + unwinder, + wait_signal, + send_ack, + required_funcs, + expected_frames=1, + ): + """Wait for signal, sample frames with retry until required funcs present, send ack.""" + _wait_for_signal(client_socket, wait_signal) + frames = None + for _ in range(MAX_TRIES): + frames = self._get_frames_with_retry(unwinder, required_funcs) + if frames and len(frames) >= expected_frames: + break + time.sleep(RETRY_DELAY) + client_socket.sendall(send_ack) + return frames + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_cache_hit_same_stack(self): + """Test that consecutive samples reuse cached parent frame objects. + + The current frame (index 0) is always re-read from memory to get + updated line numbers, so it may be a different object. Parent frames + (index 1+) should be identical objects from cache. + """ + script_body = """\ + def level3(): + sock.sendall(b"sync1") + sock.recv(16) + sock.sendall(b"sync2") + sock.recv(16) + sock.sendall(b"sync3") + sock.recv(16) + + def level2(): + level3() + + def level1(): + level2() + + level1() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + expected = {"level1", "level2", "level3"} + + frames1 = self._sample_frames( + client_socket, unwinder, b"sync1", b"ack", expected + ) + frames2 = self._sample_frames( + client_socket, unwinder, b"sync2", b"ack", expected + ) + frames3 = self._sample_frames( + client_socket, unwinder, b"sync3", b"done", expected + ) + + self.assertIsNotNone(frames1) + self.assertIsNotNone(frames2) + self.assertIsNotNone(frames3) + self.assertEqual(len(frames1), len(frames2)) + self.assertEqual(len(frames2), len(frames3)) + + # Current frame (index 0) is always re-read, so check value equality + self.assertEqual(frames1[0].funcname, frames2[0].funcname) + self.assertEqual(frames2[0].funcname, frames3[0].funcname) + + # Parent frames (index 1+) must be identical objects (cache reuse) + for i in range(1, len(frames1)): + f1, f2, f3 = frames1[i], frames2[i], frames3[i] + self.assertIs( + f1, f2, f"Frame {i}: samples 1-2 must be same object" + ) + self.assertIs( + f2, f3, f"Frame {i}: samples 2-3 must be same object" + ) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_line_number_updates_in_same_frame(self): + """Test that line numbers are correctly updated when execution moves within a function. + + When the profiler samples at different points within the same function, + it must report the correct line number for each sample, not stale cached values. + """ + script_body = """\ + def outer(): + inner() + + def inner(): + sock.sendall(b"line_a"); sock.recv(16) + sock.sendall(b"line_b"); sock.recv(16) + sock.sendall(b"line_c"); sock.recv(16) + sock.sendall(b"line_d"); sock.recv(16) + + outer() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + + frames_a = self._sample_frames( + client_socket, unwinder, b"line_a", b"ack", {"inner"} + ) + frames_b = self._sample_frames( + client_socket, unwinder, b"line_b", b"ack", {"inner"} + ) + frames_c = self._sample_frames( + client_socket, unwinder, b"line_c", b"ack", {"inner"} + ) + frames_d = self._sample_frames( + client_socket, unwinder, b"line_d", b"done", {"inner"} + ) + + self.assertIsNotNone(frames_a) + self.assertIsNotNone(frames_b) + self.assertIsNotNone(frames_c) + self.assertIsNotNone(frames_d) + + # Get the 'inner' frame from each sample (should be index 0) + inner_a = frames_a[0] + inner_b = frames_b[0] + inner_c = frames_c[0] + inner_d = frames_d[0] + + self.assertEqual(inner_a.funcname, "inner") + self.assertEqual(inner_b.funcname, "inner") + self.assertEqual(inner_c.funcname, "inner") + self.assertEqual(inner_d.funcname, "inner") + + # Line numbers must be different and increasing (execution moves forward) + self.assertLess( + inner_a.location.lineno, inner_b.location.lineno, "Line B should be after line A" + ) + self.assertLess( + inner_b.location.lineno, inner_c.location.lineno, "Line C should be after line B" + ) + self.assertLess( + inner_c.location.lineno, inner_d.location.lineno, "Line D should be after line C" + ) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_cache_invalidation_on_return(self): + """Test cache invalidation when stack shrinks (function returns).""" + script_body = """\ + def inner(): + sock.sendall(b"at_inner") + sock.recv(16) + + def outer(): + inner() + sock.sendall(b"at_outer") + sock.recv(16) + + outer() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + + frames_deep = self._sample_frames( + client_socket, + unwinder, + b"at_inner", + b"ack", + {"inner", "outer"}, + ) + frames_shallow = self._sample_frames( + client_socket, unwinder, b"at_outer", b"done", {"outer"} + ) + + self.assertIsNotNone(frames_deep) + self.assertIsNotNone(frames_shallow) + + funcs_deep = [f.funcname for f in frames_deep] + funcs_shallow = [f.funcname for f in frames_shallow] + + self.assertIn("inner", funcs_deep) + self.assertIn("outer", funcs_deep) + self.assertNotIn("inner", funcs_shallow) + self.assertIn("outer", funcs_shallow) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_cache_invalidation_on_call(self): + """Test cache invalidation when stack grows (new function called).""" + script_body = """\ + def deeper(): + sock.sendall(b"at_deeper") + sock.recv(16) + + def middle(): + sock.sendall(b"at_middle") + sock.recv(16) + deeper() + + def top(): + middle() + + top() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + + frames_before = self._sample_frames( + client_socket, + unwinder, + b"at_middle", + b"ack", + {"middle", "top"}, + ) + frames_after = self._sample_frames( + client_socket, + unwinder, + b"at_deeper", + b"done", + {"deeper", "middle", "top"}, + ) + + self.assertIsNotNone(frames_before) + self.assertIsNotNone(frames_after) + + funcs_before = [f.funcname for f in frames_before] + funcs_after = [f.funcname for f in frames_after] + + self.assertIn("middle", funcs_before) + self.assertIn("top", funcs_before) + self.assertNotIn("deeper", funcs_before) + + self.assertIn("deeper", funcs_after) + self.assertIn("middle", funcs_after) + self.assertIn("top", funcs_after) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_partial_stack_reuse(self): + """Test that unchanged parent frames are reused from cache when top frame moves.""" + script_body = """\ + def level4(): + sock.sendall(b"sync1") + sock.recv(16) + sock.sendall(b"sync2") + sock.recv(16) + + def level3(): + level4() + + def level2(): + level3() + + def level1(): + level2() + + level1() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + + # Sample 1: level4 at first sendall + frames1 = self._sample_frames( + client_socket, + unwinder, + b"sync1", + b"ack", + {"level1", "level2", "level3", "level4"}, + ) + # Sample 2: level4 at second sendall (same stack, different line) + frames2 = self._sample_frames( + client_socket, + unwinder, + b"sync2", + b"done", + {"level1", "level2", "level3", "level4"}, + ) + + self.assertIsNotNone(frames1) + self.assertIsNotNone(frames2) + + def find_frame(frames, funcname): + for f in frames: + if f.funcname == funcname: + return f + return None + + # level4 should have different line numbers (it moved) + l4_1 = find_frame(frames1, "level4") + l4_2 = find_frame(frames2, "level4") + self.assertIsNotNone(l4_1) + self.assertIsNotNone(l4_2) + self.assertNotEqual( + l4_1.location.lineno, + l4_2.location.lineno, + "level4 should be at different lines", + ) + + # Parent frames (level1, level2, level3) should be reused from cache + for name in ["level1", "level2", "level3"]: + f1 = find_frame(frames1, name) + f2 = find_frame(frames2, name) + self.assertIsNotNone(f1, f"{name} missing from sample 1") + self.assertIsNotNone(f2, f"{name} missing from sample 2") + self.assertIs(f1, f2, f"{name} should be reused from cache") + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_recursive_frames(self): + """Test caching with same function appearing multiple times (recursion).""" + script_body = """\ + def recurse(n): + if n <= 0: + sock.sendall(b"sync1") + sock.recv(16) + sock.sendall(b"sync2") + sock.recv(16) + else: + recurse(n - 1) + + recurse(5) + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + + frames1 = self._sample_frames( + client_socket, unwinder, b"sync1", b"ack", {"recurse"} + ) + frames2 = self._sample_frames( + client_socket, unwinder, b"sync2", b"done", {"recurse"} + ) + + self.assertIsNotNone(frames1) + self.assertIsNotNone(frames2) + + # Should have multiple "recurse" frames (6 total: recurse(5) down to recurse(0)) + recurse_count = sum(1 for f in frames1 if f.funcname == "recurse") + self.assertEqual(recurse_count, 6, "Should have 6 recursive frames") + + self.assertEqual(len(frames1), len(frames2)) + + # Current frame (index 0) is re-read, check value equality + self.assertEqual(frames1[0].funcname, frames2[0].funcname) + + # Parent frames (index 1+) should be identical objects (cache reuse) + for i in range(1, len(frames1)): + self.assertIs( + frames1[i], + frames2[i], + f"Frame {i}: recursive frames must be same object", + ) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_cache_vs_no_cache_equivalence(self): + """Test that cache_frames=True and cache_frames=False produce equivalent results.""" + script_body = """\ + def level3(): + sock.sendall(b"ready"); sock.recv(16) + + def level2(): + level3() + + def level1(): + level2() + + level1() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + _wait_for_signal(client_socket, b"ready") + + # Sample with cache + unwinder_cache = make_unwinder(cache_frames=True) + frames_cached = self._get_frames_with_retry( + unwinder_cache, {"level1", "level2", "level3"} + ) + + # Sample without cache + unwinder_no_cache = make_unwinder(cache_frames=False) + frames_no_cache = self._get_frames_with_retry( + unwinder_no_cache, {"level1", "level2", "level3"} + ) + + client_socket.sendall(b"done") + + self.assertIsNotNone(frames_cached) + self.assertIsNotNone(frames_no_cache) + + # Same number of frames + self.assertEqual(len(frames_cached), len(frames_no_cache)) + + # Same function names in same order + funcs_cached = [f.funcname for f in frames_cached] + funcs_no_cache = [f.funcname for f in frames_no_cache] + self.assertEqual(funcs_cached, funcs_no_cache) + + # For level3 (leaf frame), due to timing races we can be at either + # sock.sendall() or sock.recv() - both are valid. For parent frames, + # the locations should match exactly. + # Valid locations for level3: line 3 has two statements + # sock.sendall(b"ready") -> col 4-26 + # sock.recv(16) -> col 28-41 + level3_valid_cols = {(4, 26), (28, 41)} + + for i in range(len(frames_cached)): + loc_cached = frames_cached[i].location + loc_no_cache = frames_no_cache[i].location + + if frames_cached[i].funcname == "level3": + # Leaf frame: can be at either statement + self.assertIn( + (loc_cached.col_offset, loc_cached.end_col_offset), + level3_valid_cols, + ) + self.assertIn( + (loc_no_cache.col_offset, loc_no_cache.end_col_offset), + level3_valid_cols, + ) + else: + # Parent frames: must match exactly + self.assertEqual(loc_cached, loc_no_cache) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_cache_per_thread_isolation(self): + """Test that frame cache is per-thread and cache invalidation works independently.""" + script_body = """\ + import threading + + lock = threading.Lock() + + def sync(msg): + with lock: + sock.sendall(msg + b"\\n") + sock.recv(1) + + # Thread 1 functions + def baz1(): + sync(b"t1:baz1") + + def bar1(): + baz1() + + def blech1(): + sync(b"t1:blech1") + + def foo1(): + bar1() # Goes down to baz1, syncs + blech1() # Returns up, goes down to blech1, syncs + + # Thread 2 functions + def baz2(): + sync(b"t2:baz2") + + def bar2(): + baz2() + + def blech2(): + sync(b"t2:blech2") + + def foo2(): + bar2() # Goes down to baz2, syncs + blech2() # Returns up, goes down to blech2, syncs + + t1 = threading.Thread(target=foo1) + t2 = threading.Thread(target=foo2) + t1.start() + t2.start() + t1.join() + t2.join() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder = make_unwinder(cache_frames=True) + + # Message dispatch table: signal -> required functions for that thread + dispatch = { + b"t1:baz1": {"baz1", "bar1", "foo1"}, + b"t2:baz2": {"baz2", "bar2", "foo2"}, + b"t1:blech1": {"blech1", "foo1"}, + b"t2:blech2": {"blech2", "foo2"}, + } + + # Track results for each sync point + results = {} + + # Process 4 sync points (order depends on thread scheduling) + buffer = _wait_for_signal(client_socket, b"\n") + for i in range(4): + # Extract first message from buffer + msg, sep, buffer = buffer.partition(b"\n") + self.assertIn(msg, dispatch, f"Unexpected message: {msg!r}") + + # Sample frames for the thread at this sync point + required_funcs = dispatch[msg] + frames = self._get_frames_with_retry(unwinder, required_funcs) + self.assertIsNotNone(frames, f"Thread not found for {msg!r}") + results[msg] = [f.funcname for f in frames] + + # Release thread and wait for next message (if not last) + client_socket.sendall(b"k") + if i < 3: + buffer += _wait_for_signal(client_socket, b"\n") + + # Validate Phase 1: baz snapshots + t1_baz = results.get(b"t1:baz1") + t2_baz = results.get(b"t2:baz2") + self.assertIsNotNone(t1_baz, "Missing t1:baz1 snapshot") + self.assertIsNotNone(t2_baz, "Missing t2:baz2 snapshot") + + # Thread 1 at baz1: should have foo1->bar1->baz1 + self.assertIn("baz1", t1_baz) + self.assertIn("bar1", t1_baz) + self.assertIn("foo1", t1_baz) + self.assertNotIn("blech1", t1_baz) + # No cross-contamination + self.assertNotIn("baz2", t1_baz) + self.assertNotIn("bar2", t1_baz) + self.assertNotIn("foo2", t1_baz) + + # Thread 2 at baz2: should have foo2->bar2->baz2 + self.assertIn("baz2", t2_baz) + self.assertIn("bar2", t2_baz) + self.assertIn("foo2", t2_baz) + self.assertNotIn("blech2", t2_baz) + # No cross-contamination + self.assertNotIn("baz1", t2_baz) + self.assertNotIn("bar1", t2_baz) + self.assertNotIn("foo1", t2_baz) + + # Validate Phase 2: blech snapshots (cache invalidation test) + t1_blech = results.get(b"t1:blech1") + t2_blech = results.get(b"t2:blech2") + self.assertIsNotNone(t1_blech, "Missing t1:blech1 snapshot") + self.assertIsNotNone(t2_blech, "Missing t2:blech2 snapshot") + + # Thread 1 at blech1: bar1/baz1 should be GONE (cache invalidated) + self.assertIn("blech1", t1_blech) + self.assertIn("foo1", t1_blech) + self.assertNotIn( + "bar1", t1_blech, "Cache not invalidated: bar1 still present" + ) + self.assertNotIn( + "baz1", t1_blech, "Cache not invalidated: baz1 still present" + ) + # No cross-contamination + self.assertNotIn("blech2", t1_blech) + + # Thread 2 at blech2: bar2/baz2 should be GONE (cache invalidated) + self.assertIn("blech2", t2_blech) + self.assertIn("foo2", t2_blech) + self.assertNotIn( + "bar2", t2_blech, "Cache not invalidated: bar2 still present" + ) + self.assertNotIn( + "baz2", t2_blech, "Cache not invalidated: baz2 still present" + ) + # No cross-contamination + self.assertNotIn("blech1", t2_blech) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_new_unwinder_with_stale_last_profiled_frame(self): + """Test that a new unwinder returns complete stack when cache lookup misses.""" + script_body = """\ + def level4(): + sock.sendall(b"sync1") + sock.recv(16) + sock.sendall(b"sync2") + sock.recv(16) + + def level3(): + level4() + + def level2(): + level3() + + def level1(): + level2() + + level1() + """ + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + expected = {"level1", "level2", "level3", "level4"} + + # First unwinder samples - this sets last_profiled_frame in target + unwinder1 = make_unwinder(cache_frames=True) + frames1 = self._sample_frames( + client_socket, unwinder1, b"sync1", b"ack", expected + ) + + # Create NEW unwinder (empty cache) and sample + # The target still has last_profiled_frame set from unwinder1 + unwinder2 = make_unwinder(cache_frames=True) + frames2 = self._sample_frames( + client_socket, unwinder2, b"sync2", b"done", expected + ) + + self.assertIsNotNone(frames1) + self.assertIsNotNone(frames2) + + funcs1 = [f.funcname for f in frames1] + funcs2 = [f.funcname for f in frames2] + + # Both should have all levels + for level in ["level1", "level2", "level3", "level4"]: + self.assertIn(level, funcs1, f"{level} missing from first sample") + self.assertIn(level, funcs2, f"{level} missing from second sample") + + # Should have same stack depth + self.assertEqual( + len(frames1), + len(frames2), + "New unwinder should return complete stack despite stale last_profiled_frame", + ) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_cache_exhaustion(self): + """Test cache works when frame limit (1024) is exceeded. + + FRAME_CACHE_MAX_FRAMES=1024. With 1100 recursive frames, + the cache can't store all of them but should still work. + """ + # Use 1100 to exceed FRAME_CACHE_MAX_FRAMES=1024 + depth = 1100 + script_body = f"""\ +import sys +sys.setrecursionlimit(2000) + +def recurse(n): + if n <= 0: + sock.sendall(b"ready") + sock.recv(16) # wait for ack + sock.sendall(b"ready2") + sock.recv(16) # wait for done + return + recurse(n - 1) + +recurse({depth}) +""" + + with self._target_process(script_body) as ( + p, + client_socket, + make_unwinder, + ): + unwinder_cache = make_unwinder(cache_frames=True) + unwinder_no_cache = make_unwinder(cache_frames=False) + + frames_cached = self._sample_frames( + client_socket, + unwinder_cache, + b"ready", + b"ack", + {"recurse"}, + expected_frames=1102, + ) + # Sample again with no cache for comparison + frames_no_cache = self._sample_frames( + client_socket, + unwinder_no_cache, + b"ready2", + b"done", + {"recurse"}, + expected_frames=1102, + ) + + self.assertIsNotNone(frames_cached) + self.assertIsNotNone(frames_no_cache) + + # Both should have many recurse frames (> 1024 limit) + cached_count = [f.funcname for f in frames_cached].count("recurse") + no_cache_count = [f.funcname for f in frames_no_cache].count("recurse") + + self.assertGreater( + cached_count, 1000, "Should have >1000 recurse frames" + ) + self.assertGreater( + no_cache_count, 1000, "Should have >1000 recurse frames" + ) + + # Both modes should produce same frame count + self.assertEqual( + len(frames_cached), + len(frames_no_cache), + "Cache exhaustion should not affect stack completeness", + ) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_get_stats(self): + """Test that get_stats() returns statistics when stats=True.""" + script_body = """\ + sock.sendall(b"ready") + sock.recv(16) + """ + + with self._target_process(script_body) as (p, client_socket, _): + unwinder = RemoteUnwinder(p.pid, all_threads=True, stats=True) + _wait_for_signal(client_socket, b"ready") + + # Take a sample + _get_stack_trace_with_retry(unwinder) + + stats = unwinder.get_stats() + client_socket.sendall(b"done") + + # Verify expected keys exist + expected_keys = [ + "total_samples", + "frame_cache_hits", + "frame_cache_misses", + "frame_cache_partial_hits", + "frames_read_from_cache", + "frames_read_from_memory", + "frame_cache_hit_rate", + "batched_read_attempts", + "batched_read_successes", + "batched_read_misses", + "batched_read_segments_requested", + "batched_read_segments_completed", + "batched_read_success_rate", + "batched_read_segment_completion_rate", + ] + for key in expected_keys: + self.assertIn(key, stats) + + self.assertEqual(stats["total_samples"], 1) + + @skip_if_not_supported + @unittest.skipIf( + sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, + "Test only runs on Linux with process_vm_readv support", + ) + def test_get_stats_disabled_raises(self): + """Test that get_stats() raises RuntimeError when stats=False.""" + script_body = """\ + sock.sendall(b"ready") + sock.recv(16) + """ + + with self._target_process(script_body) as (p, client_socket, _): + unwinder = RemoteUnwinder( + p.pid, all_threads=True + ) # stats=False by default + _wait_for_signal(client_socket, b"ready") + + with self.assertRaises(RuntimeError): + unwinder.get_stats() + + client_socket.sendall(b"done") if __name__ == "__main__": diff --git a/crates/host_env/Cargo.toml b/crates/host_env/Cargo.toml index 7f9634213..b7ed30e23 100644 --- a/crates/host_env/Cargo.toml +++ b/crates/host_env/Cargo.toml @@ -77,8 +77,12 @@ libloading = { workspace = true, optional = true } libffi = { workspace = true, features = ["system"], optional = true } [target.'cfg(target_os = "macos")'.dependencies] +mach2 = { workspace = true } system-configuration = { workspace = true } +[target.'cfg(any(all(target_os = "linux", not(target_env = "uclibc")), target_os = "macos", windows))'.dependencies] +object = { workspace = true } + [target.'cfg(windows)'.dependencies] memchr.workspace = true junction = { workspace = true } @@ -97,6 +101,7 @@ windows-sys = { workspace = true, features = [ "Win32_System_Com", "Win32_System_Console", "Win32_System_Diagnostics_Debug", + "Win32_System_Diagnostics_ToolHelp", "Win32_System_Environment", "Win32_System_EventLog", "Win32_System_Ole", diff --git a/crates/host_env/src/lib.rs b/crates/host_env/src/lib.rs index 33b5ad60b..912fd6274 100644 --- a/crates/host_env/src/lib.rs +++ b/crates/host_env/src/lib.rs @@ -36,6 +36,7 @@ pub mod locale; #[cfg(feature = "native-certs")] pub mod native_certs; pub mod readline; +pub mod remote_debugging; #[cfg(feature = "ssl")] pub mod ssl; diff --git a/crates/host_env/src/remote_debugging/common.rs b/crates/host_env/src/remote_debugging/common.rs new file mode 100644 index 000000000..a3f9be698 --- /dev/null +++ b/crates/host_env/src/remote_debugging/common.rs @@ -0,0 +1,62 @@ +use std::io; + +pub(super) const COOKIE: &[u8; 8] = b"xdebugpy"; + +pub(super) trait RemoteRead { + fn read_exact(&mut self, address: u64, destination: &mut [u8]) -> io::Result<()>; +} + +pub(super) fn has_cookie(reader: &mut impl RemoteRead, address: u64) -> io::Result { + if address == 0 { + return Ok(false); + } + let mut bytes = [0; COOKIE.len()]; + match reader.read_exact(address, &mut bytes) { + Ok(()) => Ok(&bytes == COOKIE), + Err(error) if stops_probe(&error) => Err(error), + Err(_) => Ok(false), + } +} + +pub(super) fn stops_probe(error: &io::Error) -> bool { + matches!( + error.kind(), + io::ErrorKind::PermissionDenied | io::ErrorKind::OutOfMemory + ) +} + +pub(super) fn invalid_data(message: &'static str) -> io::Error { + io::Error::new(io::ErrorKind::InvalidData, message) +} + +// These are resource limits, not claims about an unwinder's capabilities. +// Exhaustion is a failed probe, just like inaccessible or missing metadata. +pub(super) struct ReadBudget { + bytes: usize, + calls: usize, +} + +impl ReadBudget { + pub(super) fn new() -> Self { + Self { + bytes: 256 * 1024 * 1024, + calls: 65_536, + } + } + + pub(super) fn consume(&mut self, bytes: usize) -> io::Result<()> { + self.bytes = self.bytes.checked_sub(bytes).ok_or_else(|| { + io::Error::new( + io::ErrorKind::OutOfMemory, + "runtime probe byte budget exhausted", + ) + })?; + self.calls = self.calls.checked_sub(1).ok_or_else(|| { + io::Error::new( + io::ErrorKind::OutOfMemory, + "runtime probe read budget exhausted", + ) + })?; + Ok(()) + } +} diff --git a/crates/host_env/src/remote_debugging/linux.rs b/crates/host_env/src/remote_debugging/linux.rs new file mode 100644 index 000000000..a237ac47c --- /dev/null +++ b/crates/host_env/src/remote_debugging/linux.rs @@ -0,0 +1,190 @@ +use super::{COOKIE, ReadBudget, RemoteRead, has_cookie, invalid_data, object_file, stops_probe}; +use std::{ + ffi::OsStr, + fs::{File, OpenOptions}, + io::{self, Read}, + os::unix::{ffi::OsStrExt, fs::FileExt}, + path::Path, +}; + +const MAX_MAPS_BYTES: u64 = 8 * 1024 * 1024; +const SCAN_CHUNK: usize = 1024 * 1024; + +struct Process { + pid: i32, + memory: Option, + budget: ReadBudget, +} + +impl RemoteRead for Process { + fn read_exact(&mut self, address: u64, destination: &mut [u8]) -> io::Result<()> { + let address = usize::try_from(address) + .map_err(|_| invalid_data("remote address exceeds native pointer width"))?; + address + .checked_add(destination.len()) + .ok_or_else(|| invalid_data("remote read address overflow"))?; + let mut completed = 0; + while completed < destination.len() { + let remaining = &mut destination[completed..]; + self.budget.consume(remaining.len())?; + let current = address + completed; + let count = if let Some(memory) = &self.memory { + memory.read_at(remaining, current as u64)? + } else { + let local = libc::iovec { + iov_base: remaining.as_mut_ptr().cast(), + iov_len: remaining.len(), + }; + let remote = libc::iovec { + iov_base: current as *mut libc::c_void, + iov_len: remaining.len(), + }; + // SAFETY: local names the writable destination slice. The + // remote pointer is only an address for the kernel, never a + // pointer dereferenced by this process. + let result = unsafe { libc::process_vm_readv(self.pid, &local, 1, &remote, 1, 0) }; + if result < 0 { + let error = io::Error::last_os_error(); + if error.raw_os_error() == Some(libc::ENOSYS) { + // Match CPython's fallback only for a missing syscall. + // Permission denials must not select another route. + self.memory = Some( + OpenOptions::new() + .read(true) + .write(true) + .open(format!("/proc/{}/mem", self.pid))?, + ); + continue; + } + return Err(error); + } + result as usize + }; + if count == 0 || count > remaining.len() { + return Err(io::ErrorKind::UnexpectedEof.into()); + } + completed += count; + } + Ok(()) + } +} + +struct Mapping<'a> { + start: u64, + end: u64, + writable: bool, + path: &'a [u8], +} + +fn token<'a>(remaining: &mut &'a [u8]) -> Option<&'a [u8]> { + let source = (*remaining).trim_ascii_start(); + if source.is_empty() { + return None; + } + let end = source + .iter() + .position(u8::is_ascii_whitespace) + .unwrap_or(source.len()); + let (value, rest) = source.split_at(end); + *remaining = rest; + Some(value) +} + +fn mapping(line: &[u8]) -> Option> { + let mut fields = line; + let range = token(&mut fields)?; + let permissions = token(&mut fields)?; + for _ in 0..3 { + token(&mut fields)?; + } + let path = fields.trim_ascii_start(); + if path.is_empty() || (path.starts_with(b"[") && path.ends_with(b"]")) { + return None; + } + let dash = range.iter().position(|byte| *byte == b'-')?; + let parse = |bytes| u64::from_str_radix(core::str::from_utf8(bytes).ok()?, 16).ok(); + let start = parse(&range[..dash])?; + let end = parse(&range[dash + 1..])?; + if start == 0 || end <= start { + return None; + } + Some(Mapping { + start, + end, + writable: permissions.starts_with(b"rw"), + path, + }) +} + +fn scan_deleted(process: &mut Process, mapping: &Mapping<'_>) -> io::Result> { + let mut buffer = vec![0; SCAN_CHUNK]; + let mut address = mapping.start; + while address < mapping.end { + let wanted = usize::try_from((mapping.end - address).min(SCAN_CHUNK as u64)) + .map_err(|_| invalid_data("mapping scan size overflow"))?; + match process.read_exact(address, &mut buffer[..wanted]) { + Ok(()) => { + if let Some(offset) = buffer[..wanted] + .windows(COOKIE.len()) + .position(|bytes| bytes == COOKIE) + { + return Ok(Some(address + offset as u64)); + } + } + Err(error) if stops_probe(&error) => return Err(error), + Err(_) => {} + } + let overlap = COOKIE.len() - 1; + if wanted <= overlap { + break; + } + address = address + .checked_add((wanted - overlap) as u64) + .ok_or_else(|| invalid_data("mapping scan address overflow"))?; + } + Ok(None) +} + +pub(super) fn probe(pid: i32) -> io::Result { + let mut maps = Vec::new(); + File::open(format!("/proc/{pid}/maps"))? + .take(MAX_MAPS_BYTES + 1) + .read_to_end(&mut maps)?; + if maps.len() as u64 > MAX_MAPS_BYTES { + return Err(invalid_data("process map exceeds its size bound")); + } + let mut process = Process { + pid, + memory: None, + budget: ReadBudget::new(), + }; + for line in maps.split(|byte| *byte == b'\n') { + process.budget.consume(0)?; + let Some(mapping) = mapping(line) else { + continue; + }; + let path = Path::new(OsStr::from_bytes(mapping.path)); + let Some(name) = path.file_name() else { + continue; + }; + if !name.as_bytes().windows(6).any(|bytes| bytes == b"python") { + continue; + } + let candidate = if mapping.path.ends_with(b" (deleted)") { + if !mapping.writable { + continue; + } + scan_deleted(&mut process, &mapping) + } else { + object_file::read(path, &mut process.budget).and_then(|bytes| { + object_file::section_address(&bytes, mapping.start, object::BinaryFormat::Elf) + }) + }; + match candidate { + Ok(Some(address)) if has_cookie(&mut process, address)? => return Ok(true), + Err(error) if stops_probe(&error) => return Err(error), + _ => {} + } + } + Ok(false) +} diff --git a/crates/host_env/src/remote_debugging/macos.rs b/crates/host_env/src/remote_debugging/macos.rs new file mode 100644 index 000000000..c43916628 --- /dev/null +++ b/crates/host_env/src/remote_debugging/macos.rs @@ -0,0 +1,281 @@ +use super::{ReadBudget, RemoteRead, has_cookie, invalid_data, object_file, stops_probe}; +use core::mem::size_of; +use mach2::{ + kern_return::{ + KERN_INVALID_ADDRESS, KERN_NO_ACCESS, KERN_PROTECTION_FAILURE, KERN_RESOURCE_SHORTAGE, + KERN_SUCCESS, kern_return_t, + }, + mach_port::mach_port_deallocate, + port::{MACH_PORT_DEAD, MACH_PORT_NULL, mach_port_t}, + traps::{mach_task_self, task_for_pid}, + vm::{mach_vm_read_overwrite, mach_vm_region}, + vm_prot::{VM_PROT_EXECUTE, VM_PROT_READ}, + vm_region::{VM_REGION_BASIC_INFO_64, vm_region_basic_info_64}, +}; +use object::BinaryFormat; +use std::{ffi::OsStr, io, os::unix::ffi::OsStrExt, path::Path}; + +const MAX_FAT_ARCHES: usize = 64; + +struct OwnedPort(mach_port_t); + +impl Drop for OwnedPort { + fn drop(&mut self) { + if self.0 != MACH_PORT_NULL && self.0 != MACH_PORT_DEAD { + // task_for_pid and mach_vm_region each return an owned send right. + // mach_task_self is borrowed and must not itself be deallocated. + unsafe { mach_port_deallocate(mach_task_self(), self.0) }; + } + } +} + +struct Process { + task: OwnedPort, + budget: ReadBudget, +} + +fn mach_error(code: kern_return_t) -> io::Error { + // Mach error codes carry the detailed code in their low fourteen bits. + let kind = match code & 0x3fff { + KERN_PROTECTION_FAILURE | KERN_NO_ACCESS => io::ErrorKind::PermissionDenied, + KERN_RESOURCE_SHORTAGE => io::ErrorKind::OutOfMemory, + KERN_INVALID_ADDRESS => io::ErrorKind::UnexpectedEof, + _ => io::ErrorKind::Other, + }; + io::Error::new(kind, format!("Mach runtime probe failed: {code}")) +} + +impl RemoteRead for Process { + fn read_exact(&mut self, address: u64, destination: &mut [u8]) -> io::Result<()> { + if destination.is_empty() { + return Ok(()); + } + address + .checked_add(destination.len() as u64 - 1) + .ok_or_else(|| invalid_data("remote read address overflow"))?; + self.budget.consume(destination.len())?; + let mut count = 0; + // Only the local destination is a dereferenceable pointer. The remote + // address remains an integer interpreted by the kernel. + let result = unsafe { + mach_vm_read_overwrite( + self.task.0, + address, + destination.len() as u64, + destination.as_mut_ptr() as usize as u64, + &mut count, + ) + }; + if result != KERN_SUCCESS { + return Err(mach_error(result)); + } + if count != destination.len() as u64 { + return Err(io::ErrorKind::UnexpectedEof.into()); + } + Ok(()) + } +} + +fn word(bytes: &[u8], offset: usize) -> io::Result<[u8; N]> { + let end = offset + .checked_add(N) + .ok_or_else(|| invalid_data("Mach-O field offset overflow"))?; + bytes + .get(offset..end) + .and_then(|value| value.try_into().ok()) + .ok_or_else(|| invalid_data("truncated Mach-O field")) +} + +fn read_u32(bytes: &[u8], offset: usize, little: bool) -> io::Result { + let value = word(bytes, offset)?; + Ok(if little { + u32::from_le_bytes(value) + } else { + u32::from_be_bytes(value) + }) +} + +fn sysctl_int(name: &core::ffi::CStr) -> io::Result { + let mut value = 0i32; + let mut length = size_of::(); + let result = unsafe { + libc::sysctlbyname( + name.as_ptr(), + (&mut value as *mut i32).cast(), + &mut length, + core::ptr::null_mut(), + 0, + ) + }; + if result != 0 { + return Err(io::Error::last_os_error()); + } + if length != size_of::() { + return Err(invalid_data("invalid CPU sysctl size")); + } + Ok(value) +} + +fn mach_slice(bytes: &[u8]) -> io::Result> { + let magic = u32::from_be_bytes(word(bytes, 0)?); + let little = match magic { + object::macho::FAT_MAGIC => false, + object::macho::FAT_CIGAM => true, + _ => return Ok(Some(bytes)), + }; + let count = usize::try_from(read_u32(bytes, 4, little)?) + .map_err(|_| invalid_data("Mach-O architecture count overflow"))?; + if count > MAX_FAT_ARCHES { + return Err(invalid_data("Mach-O architecture count exceeds its bound")); + } + let entry_size = 20; + let table_end = count + .checked_mul(entry_size) + .and_then(|size| size.checked_add(8)) + .ok_or_else(|| invalid_data("Mach-O architecture table overflow"))?; + if table_end > bytes.len() { + return Err(invalid_data("truncated Mach-O architecture table")); + } + // Match the pinned CPython locator's host-CPU selection, rather than guessing + // from the inspector's build target (which may itself run under translation). + let cpu = sysctl_int(c"hw.cputype")? as u32; + let capable = sysctl_int(c"hw.cpu64bit_capable")?; + let cpu = cpu + | if capable != 0 { + object::macho::CPU_ARCH_ABI64 + } else { + 0 + }; + for index in 0..count { + // The bounded, validated architecture table makes these offsets safe. + let entry = 8 + index * entry_size; + if read_u32(bytes, entry, little)? != cpu { + continue; + } + let offset = u64::from(read_u32(bytes, entry + 8, little)?); + let length = u64::from(read_u32(bytes, entry + 12, little)?); + let end = offset + .checked_add(length) + .ok_or_else(|| invalid_data("Mach-O architecture range overflow"))?; + let offset = usize::try_from(offset) + .map_err(|_| invalid_data("Mach-O architecture offset overflow"))?; + let end = + usize::try_from(end).map_err(|_| invalid_data("Mach-O architecture end overflow"))?; + if offset < table_end || length == 0 { + return Err(invalid_data("invalid Mach-O architecture range")); + } + return bytes + .get(offset..end) + .map(Some) + .ok_or_else(|| invalid_data("truncated Mach-O architecture")); + } + Ok(None) +} + +fn candidate(process: &mut Process, path: &Path, base: u64) -> io::Result { + let bytes = object_file::read(path, &mut process.budget)?; + let Some(bytes) = mach_slice(&bytes)? else { + return Ok(false); + }; + let Some(address) = object_file::section_address(bytes, base, BinaryFormat::MachO)? else { + return Ok(false); + }; + has_cookie(process, address) +} + +fn search(process: &mut Process, pid: i32, prefix: &[u8]) -> io::Result { + let mut address = 0; + loop { + process.budget.consume(0)?; + let requested = address; + let mut size = 0; + let mut info = vm_region_basic_info_64::default(); + let mut count = vm_region_basic_info_64::count(); + let mut object = MACH_PORT_NULL; + let result = unsafe { + mach_vm_region( + process.task.0, + &mut address, + &mut size, + VM_REGION_BASIC_INFO_64, + (&mut info as *mut vm_region_basic_info_64).cast(), + &mut count, + &mut object, + ) + }; + // Even irrelevant regions may return an owned memory-object port. + let _object = OwnedPort(object); + if result == KERN_INVALID_ADDRESS { + return Ok(false); + } + if result != KERN_SUCCESS { + return Err(mach_error(result)); + } + if count < vm_region_basic_info_64::count() || address < requested || size == 0 { + return Err(invalid_data("invalid Mach memory region")); + } + let next = address + .checked_add(size) + .ok_or_else(|| invalid_data("Mach memory region address overflow"))?; + if info.protection & (VM_PROT_READ | VM_PROT_EXECUTE) == (VM_PROT_READ | VM_PROT_EXECUTE) { + let mut path = [0u8; libc::MAXPATHLEN as usize + 1]; + process.budget.consume(0)?; + let length = unsafe { + // A zero return can describe an anonymous mapping; clear errno + // first so an older failure cannot be mistaken for a denial. + *libc::__error() = 0; + libc::proc_regionfilename( + pid, + address, + path.as_mut_ptr().cast(), + libc::MAXPATHLEN as u32, + ) + }; + if length <= 0 { + let error = io::Error::last_os_error(); + if stops_probe(&error) { + return Err(error); + } + } else if (length as usize) < path.len() + && let Some(end) = path.iter().position(|&byte| byte == 0) + && end < libc::MAXPATHLEN as usize + { + let path = &path[..end]; + let name = path.rsplit(|&byte| byte == b'/').next().unwrap_or(path); + if name.starts_with(prefix) { + match candidate(process, Path::new(OsStr::from_bytes(path)), address) { + Ok(true) => return Ok(true), + Err(error) if stops_probe(&error) => return Err(error), + Ok(false) | Err(_) => {} + } + } + } + } + address = next; + } +} + +pub(super) fn probe(pid: i32) -> io::Result { + let mut task = MACH_PORT_NULL; + let result = unsafe { task_for_pid(mach_task_self(), pid, &mut task) }; + let task = OwnedPort(task); + if result != KERN_SUCCESS || task.0 == MACH_PORT_NULL || task.0 == MACH_PORT_DEAD { + // CPython classifies every task_for_pid failure as a permission failure. + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!("cannot acquire task port for PID {pid}: {result}"), + )); + } + let mut process = Process { + task, + budget: ReadBudget::new(), + }; + for prefix in [b"libpython".as_slice(), b"python", b"Python"] { + match search(&mut process, pid, prefix) { + Ok(true) => return Ok(true), + Err(error) if stops_probe(&error) => return Err(error), + Ok(false) | Err(_) => {} + } + } + Ok(false) +} diff --git a/crates/host_env/src/remote_debugging/mod.rs b/crates/host_env/src/remote_debugging/mod.rs new file mode 100644 index 000000000..c65101a04 --- /dev/null +++ b/crates/host_env/src/remote_debugging/mod.rs @@ -0,0 +1,56 @@ +//! Read-only discovery of CPython's exported runtime metadata. +//! +//! A successful cookie probe identifies a runtime; it does not validate its +//! debug offsets or promise that stack unwinding is supported. + +#[cfg(all(target_os = "linux", not(target_env = "uclibc")))] +mod linux; +#[cfg(target_os = "macos")] +mod macos; +#[cfg(windows)] +mod windows; + +#[cfg(any( + all(target_os = "linux", not(target_env = "uclibc")), + target_os = "macos", + windows +))] +mod object_file; + +#[cfg(any( + all(target_os = "linux", not(target_env = "uclibc")), + target_os = "macos", + windows +))] +mod common; + +#[cfg(all(target_os = "linux", not(target_env = "uclibc")))] +use common::COOKIE; +#[cfg(any( + all(target_os = "linux", not(target_env = "uclibc")), + target_os = "macos", + windows +))] +use common::{ReadBudget, RemoteRead, has_cookie, invalid_data, stops_probe}; + +pub const PROCESS_VM_READV_SUPPORTED: u8 = + cfg!(all(target_os = "linux", not(target_env = "uclibc"))) as u8; + +pub fn is_python_process(pid: i32) -> bool { + #[cfg(all(target_os = "linux", not(target_env = "uclibc")))] + let result = linux::probe(pid); + #[cfg(target_os = "macos")] + let result = macos::probe(pid); + #[cfg(windows)] + let result = windows::probe(pid); + #[cfg(not(any( + all(target_os = "linux", not(target_env = "uclibc")), + target_os = "macos", + windows + )))] + let result: Result = { + let _ = pid; + Err(std::io::ErrorKind::Unsupported.into()) + }; + result.unwrap_or(false) +} diff --git a/crates/host_env/src/remote_debugging/object_file.rs b/crates/host_env/src/remote_debugging/object_file.rs new file mode 100644 index 000000000..37678b93a --- /dev/null +++ b/crates/host_env/src/remote_debugging/object_file.rs @@ -0,0 +1,94 @@ +use super::{ReadBudget, invalid_data}; +use object::{BinaryFormat, Object, ObjectSection, ObjectSegment}; +use std::{fs::OpenOptions, io, io::Read, path::Path}; + +const MAX_OBJECT_BYTES: u64 = 256 * 1024 * 1024; + +pub(super) fn read(path: &Path, budget: &mut ReadBudget) -> io::Result> { + budget.consume(0)?; + let mut options = OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + // A mapped pathname can be replaced with a FIFO while being inspected. + options.custom_flags(libc::O_NONBLOCK); + } + let file = options.open(path)?; + let metadata = file.metadata()?; + if !metadata.is_file() || metadata.len() > MAX_OBJECT_BYTES { + return Err(invalid_data( + "runtime candidate is not a bounded regular file", + )); + } + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(metadata.len() as usize) + .map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?; + let result = file.take(MAX_OBJECT_BYTES + 1).read_to_end(&mut bytes); + // A failed read may already have appended bytes. Account for them before + // propagating the error, as well as before rejecting a grown file. + budget.consume(bytes.len())?; + result?; + if bytes.len() as u64 > MAX_OBJECT_BYTES { + return Err(invalid_data("runtime candidate grew beyond its size bound")); + } + Ok(bytes) +} + +pub(super) fn section_address( + bytes: &[u8], + mapping_base: u64, + format: BinaryFormat, +) -> io::Result> { + let file = object::File::parse(bytes).map_err(|_| invalid_data("invalid object file"))?; + if file.format() != format { + return Ok(None); + } + let name = match format { + BinaryFormat::Elf => b".PyRuntime".as_slice(), + BinaryFormat::MachO => b"PyRuntime".as_slice(), + BinaryFormat::Pe => b"PyRuntim".as_slice(), + _ => return Ok(None), + }; + let section = if format == BinaryFormat::MachO { + file.sections().find(|section| { + section.name_bytes().ok() == Some(name) + && section.segment_name_bytes().ok() == Some(Some(b"__DATA".as_slice())) + }) + } else { + file.section_by_name_bytes(name) + }; + let Some(section) = section else { + return Ok(None); + }; + let image_base = match format { + BinaryFormat::Elf => { + // The ELF reader's segment iterator contains only PT_LOAD entries. + let first = file + .segments() + .next() + .ok_or_else(|| invalid_data("ELF image has no load segment"))?; + let alignment = first.align(); + if alignment == 0 || !alignment.is_power_of_two() { + return Err(invalid_data("invalid ELF load alignment")); + } + first.address() - first.address() % alignment + } + BinaryFormat::MachO => file + .segments() + .find(|segment| segment.name_bytes().ok() == Some(Some(b"__TEXT".as_slice()))) + .ok_or_else(|| invalid_data("Mach-O image has no text segment"))? + .address(), + BinaryFormat::Pe => file.relative_address_base(), + _ => unreachable!(), + }; + let relative = section + .address() + .checked_sub(image_base) + .ok_or_else(|| invalid_data("runtime section precedes its image base"))?; + let address = mapping_base + .checked_add(relative) + .ok_or_else(|| invalid_data("runtime section address overflow"))?; + Ok((address != 0).then_some(address)) +} diff --git a/crates/host_env/src/remote_debugging/windows.rs b/crates/host_env/src/remote_debugging/windows.rs new file mode 100644 index 000000000..77a55fc0a --- /dev/null +++ b/crates/host_env/src/remote_debugging/windows.rs @@ -0,0 +1,178 @@ +use super::{ReadBudget, RemoteRead, has_cookie, invalid_data, object_file, stops_probe}; +use core::mem::size_of; +use object::BinaryFormat; +use std::{ffi::OsString, io, os::windows::ffi::OsStringExt, path::PathBuf}; +use windows_sys::Win32::{ + Foundation::{ + CloseHandle, ERROR_BAD_LENGTH, ERROR_NO_MORE_FILES, GetLastError, HANDLE, + INVALID_HANDLE_VALUE, + }, + System::{ + Diagnostics::{ + Debug::ReadProcessMemory, + ToolHelp::{ + CreateToolhelp32Snapshot, MODULEENTRY32W, Module32FirstW, Module32NextW, + TH32CS_SNAPMODULE, + }, + }, + Threading::{ + OpenProcess, PROCESS_QUERY_INFORMATION, PROCESS_SUSPEND_RESUME, PROCESS_VM_OPERATION, + PROCESS_VM_READ, PROCESS_VM_WRITE, + }, + }, +}; + +const MAX_SNAPSHOT_ATTEMPTS: usize = 16; +const MAX_MODULES: usize = 16_384; + +struct OwnedHandle(HANDLE); + +impl Drop for OwnedHandle { + fn drop(&mut self) { + // Only successfully acquired process and snapshot handles are owned here. + unsafe { CloseHandle(self.0) }; + } +} + +struct Process { + handle: OwnedHandle, + budget: ReadBudget, +} + +impl RemoteRead for Process { + fn read_exact(&mut self, address: u64, mut destination: &mut [u8]) -> io::Result<()> { + if destination.is_empty() { + return Ok(()); + } + let end = address + .checked_add(destination.len() as u64 - 1) + .ok_or_else(|| invalid_data("remote read address overflow"))?; + usize::try_from(end).map_err(|_| invalid_data("remote address exceeds pointer width"))?; + let mut address = usize::try_from(address) + .map_err(|_| invalid_data("remote address exceeds pointer width"))?; + while !destination.is_empty() { + self.budget.consume(destination.len())?; + let mut count = 0; + // The address is passed to the OS, never dereferenced in this process. + let success = unsafe { + ReadProcessMemory( + self.handle.0, + address as *const _, + destination.as_mut_ptr().cast(), + destination.len(), + &mut count, + ) + }; + if success == 0 { + // In particular, an ERROR_PARTIAL_COPY failure is not a successful + // short read. CPython only continues after successful reads. + return Err(io::Error::last_os_error()); + } + if count == 0 { + return Err(io::ErrorKind::UnexpectedEof.into()); + } + if count > destination.len() { + return Err(invalid_data("remote read exceeded its destination")); + } + destination = &mut destination[count..]; + if !destination.is_empty() { + address = address + .checked_add(count) + .ok_or_else(|| invalid_data("remote read address overflow"))?; + } + } + Ok(()) + } +} + +fn snapshot(pid: u32, budget: &mut ReadBudget) -> io::Result { + for _ in 0..MAX_SNAPSHOT_ATTEMPTS { + budget.consume(0)?; + // Keep CPython's module selection flags, including its cross-bitness limits. + let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid) }; + if handle != INVALID_HANDLE_VALUE { + return Ok(OwnedHandle(handle)); + } + let error = unsafe { GetLastError() }; + if error != ERROR_BAD_LENGTH { + return Err(io::Error::from_raw_os_error(error as i32)); + } + } + Err(io::Error::from_raw_os_error(ERROR_BAD_LENGTH as i32)) +} + +fn terminated(value: &[u16]) -> io::Result<&[u16]> { + let end = value + .iter() + .position(|&unit| unit == 0) + .ok_or_else(|| invalid_data("unterminated module name or path"))?; + Ok(&value[..end]) +} + +fn candidate(process: &mut Process, module: &MODULEENTRY32W) -> io::Result { + const PYTHON: &[u16] = &[112, 121, 116, 104, 111, 110]; + let name = terminated(&module.szModule)?; + if !name.windows(PYTHON.len()).any(|part| part == PYTHON) { + return Ok(false); + } + let path = PathBuf::from(OsString::from_wide(terminated(&module.szExePath)?)); + let bytes = object_file::read(&path, &mut process.budget)?; + let base = module.modBaseAddr as usize as u64; + let Some(address) = object_file::section_address(&bytes, base, BinaryFormat::Pe)? else { + return Ok(false); + }; + has_cookie(process, address) +} + +pub(super) fn probe(pid: i32) -> io::Result { + // CPython requests these rights even for this read-only operation. Reducing + // them could make detection succeed where CPython's handle initialization fails. + // The DWORD conversion also preserves the Windows API's signed-PID behavior. + let handle = unsafe { + OpenProcess( + PROCESS_VM_READ + | PROCESS_VM_WRITE + | PROCESS_VM_OPERATION + | PROCESS_QUERY_INFORMATION + | PROCESS_SUSPEND_RESUME, + 0, + pid as u32, + ) + }; + if handle.is_null() { + return Err(io::Error::last_os_error()); + } + let mut process = Process { + handle: OwnedHandle(handle), + budget: ReadBudget::new(), + }; + let snapshot = snapshot(pid as u32, &mut process.budget)?; + let mut module = MODULEENTRY32W { + dwSize: size_of::() as u32, + ..Default::default() + }; + process.budget.consume(0)?; + let mut success = unsafe { Module32FirstW(snapshot.0, &mut module) }; + for _ in 0..MAX_MODULES { + if success == 0 { + let error = unsafe { GetLastError() }; + return if error == ERROR_NO_MORE_FILES { + Ok(false) + } else { + Err(io::Error::from_raw_os_error(error as i32)) + }; + } + match candidate(&mut process, &module) { + Ok(true) => return Ok(true), + Err(error) if stops_probe(&error) => return Err(error), + Ok(false) | Err(_) => {} + } + process.budget.consume(0)?; + module.dwSize = size_of::() as u32; + success = unsafe { Module32NextW(snapshot.0, &mut module) }; + } + Err(io::Error::new( + io::ErrorKind::OutOfMemory, + "runtime probe module budget exhausted", + )) +} diff --git a/crates/stdlib/src/_remote_debugging.rs b/crates/stdlib/src/_remote_debugging.rs index 898bb9efc..ba98b3b3b 100644 --- a/crates/stdlib/src/_remote_debugging.rs +++ b/crates/stdlib/src/_remote_debugging.rs @@ -7,7 +7,7 @@ mod _remote_debugging { use super::binary; use crate::vm::{ AsObject, Py, PyObject, PyObjectRef, PyRef, PyResult, VirtualMachine, - builtins::{PyBaseExceptionRef, PyDictRef, PyList, PyTuple, PyType}, + builtins::{PyBaseExceptionRef, PyDictRef, PyFloat, PyList, PyTuple, PyType}, common::{lock::PyMutex, rc::PyRc}, convert::{IntoPyException, ToPyObject}, function::{FsPath, FuncArgs}, @@ -117,6 +117,10 @@ mod _remote_debugging { #[pyclass(with(PyStructSequence))] impl InterpreterInfo {} + #[pyattr] + const PROCESS_VM_READV_SUPPORTED: u8 = + rustpython_host_env::remote_debugging::PROCESS_VM_READV_SUPPORTED; + #[pyattr] const THREAD_STATUS_HAS_GIL: u8 = 1; #[pyattr] @@ -130,6 +134,22 @@ mod _remote_debugging { #[pyattr] const THREAD_STATUS_MAIN_THREAD: u8 = 32; + #[derive(FromArgs)] + struct IsPythonProcessArgs { + #[pyarg(any)] + pid: PyObjectRef, + } + + /// Check if a process is a Python process. + #[pyfunction] + fn is_python_process(args: IsPythonProcessArgs, vm: &VirtualMachine) -> PyResult { + // Finish argument binding before invoking a user-defined __index__. + let pid: i32 = args.pid.try_into_value(vm)?; + Ok(rustpython_host_env::remote_debugging::is_python_process( + pid, + )) + } + #[pyfunction] fn zstd_available() -> bool { false @@ -293,20 +313,36 @@ mod _remote_debugging { #[derive(FromArgs)] struct BinaryExitArgs { - #[pyarg(positional, optional)] + #[pyarg(any, optional)] exc_type: Option, - #[pyarg(positional, optional)] + #[pyarg(any, optional, name = "exc_val")] _exc_value: Option, - #[pyarg(positional, optional)] + #[pyarg(any, optional, name = "exc_tb")] _exc_tb: Option, } + #[derive(FromArgs)] + struct SampleArgs { + #[pyarg(any)] + stack_frames: PyObjectRef, + #[pyarg(any)] + timestamp_us: PyObjectRef, + } + + fn profile_float(obj: &PyObject, vm: &VirtualMachine) -> PyResult { + // PyFloat_AsDouble reads float subclasses without calling __float__. + Ok(match obj.downcast_ref::() { + Some(value) => value.to_f64(), + None => obj.try_float(vm)?.to_f64(), + }) + } + #[derive(FromArgs)] struct SetStatsArgs { #[pyarg(any)] - duration_sec: f64, + duration_sec: PyObjectRef, #[pyarg(any)] - sample_rate: f64, + sample_rate: PyObjectRef, #[pyarg(any, optional)] error_rate: Option, #[pyarg(any, optional)] @@ -316,19 +352,19 @@ mod _remote_debugging { fn optional_float(value: Option, vm: &VirtualMachine) -> PyResult> { value .filter(|value| !vm.is_none(value)) - .map(|value| value.try_float(vm).map(|value| value.to_f64())) + .map(|value| profile_float(&value, vm)) .transpose() } #[pyclass(with(Constructor, Initializer))] impl BinaryWriter { #[pymethod] - fn write_sample( - zelf: &Py, - stack_frames: PyObjectRef, - timestamp_us: u64, - vm: &VirtualMachine, - ) -> PyResult<()> { + fn write_sample(zelf: &Py, args: SampleArgs, vm: &VirtualMachine) -> PyResult<()> { + let SampleArgs { + stack_frames, + timestamp_us, + } = args; + let timestamp_us = timestamp_us.try_into_value::(vm)?; if zelf.state.lock().writer.is_none() { return Err(vm.new_value_error("Writer is closed")); } @@ -378,13 +414,15 @@ mod _remote_debugging { #[pymethod] fn set_stats(zelf: &Py, args: SetStatsArgs, vm: &VirtualMachine) -> PyResult<()> { + let duration_sec = profile_float(&args.duration_sec, vm)?; + let sample_rate = profile_float(&args.sample_rate, vm)?; let is_open = zelf.state.lock().writer.is_some(); if !is_open { return Err(vm.new_value_error("Writer is closed")); } let profile = binary::ProfileStats { - duration_sec: args.duration_sec, - sample_rate: args.sample_rate, + duration_sec, + sample_rate, error_rate: optional_float(args.error_rate, vm)?, missed_samples: optional_float(args.missed_samples, vm)?, }; diff --git a/crates/stdlib/src/_remote_debugging/binary.rs b/crates/stdlib/src/_remote_debugging/binary.rs index d77b66226..1d6158ee2 100644 --- a/crates/stdlib/src/_remote_debugging/binary.rs +++ b/crates/stdlib/src/_remote_debugging/binary.rs @@ -792,7 +792,7 @@ impl Reader { let mut footer = Cursor::new(&data[footer_start..]); let string_count = endian.u32(footer.fixed()?); let frame_count = endian.u32(footer.fixed()?); - let (profile_stats, tables_end) = parse_profile_stats(&data, endian, footer_start)?; + let (profile_stats, _) = parse_profile_stats(&data, endian, footer_start)?; for (name, offset) in [("string", string_offset), ("frame", frame_offset)] { if offset > data.len() as u64 { return Err(invalid(format!( @@ -815,15 +815,14 @@ impl Reader { .map_err(|_| Error::Overflow("File is too large to read"))?; let frame_offset = usize::try_from(frame_offset) .map_err(|_| Error::Overflow("File is too large to read"))?; - if frame_offset > tables_end { - return Err(invalid("Frame table overlaps file trailer")); - } if compression_type == 1 { return Err(Error::Runtime( "File uses zstd compression but zstd support not compiled in", )); } - let mut string_data = Cursor::new(&data[string_offset..frame_offset]); + // CPython bounds table entries by the remaining file bytes, including + // subsequent sections. Keep the count and cursor checks anchored to EOF. + let mut string_data = Cursor::new(&data[string_offset..]); validate_count("string", string_count, string_data.remaining(), 1)?; let mut strings = Vec::new(); reserve(&mut strings, string_count as usize)?; @@ -836,7 +835,7 @@ impl Reader { .map_err(|_| invalid("String table overflow"))?; strings.push(decode_utf8_replacing(bytes)?); } - let mut frame_data = Cursor::new(&data[frame_offset..tables_end]); + let mut frame_data = Cursor::new(&data[frame_offset..]); validate_count("frame", frame_count, frame_data.remaining(), 7)?; let mut frames = Vec::new(); reserve(&mut frames, frame_count as usize)?; diff --git a/scripts/update_lib/cmd_migrate.py b/scripts/update_lib/cmd_migrate.py index 949fe8315..efbce2d03 100644 --- a/scripts/update_lib/cmd_migrate.py +++ b/scripts/update_lib/cmd_migrate.py @@ -19,6 +19,12 @@ Non-marker adaptations require reviewed blocks on each migration, for example Use a host Python that parses the source version (Python 3.15 for 3.15 syntax), as required by the existing marker extraction and application steps. +For additive updates, --assignments or --functions selects contiguous canonical +top-level definitions before an unchanged anchor. Both input SHA256 pins are +required. --exports copies only complete selected __all__ lines. Function +dependencies are checked without importing or executing either module; unchanged +imports and simple assignments are supported, not arbitrary existing helpers. + Canonical removals are a separate, explicit operation. JSON source metadata pins path (Lib/test/test_*), commit and tree; removals maps explicit relative filenames to sha256 and upstream_commit. The caller must first verify and @@ -43,6 +49,7 @@ import pathlib import re import shutil import subprocess +import symtable import sys sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) @@ -57,6 +64,68 @@ def _preserve_block(old, new, canonical, local): return new.replace(canonical, local, 1) +def _single_assignment(tree, name): + matches = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Name) + and isinstance(node.ctx, ast.Store) + and node.id == name + ] + nodes = [ + node + for node in tree.body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id == name + ] + rebound = any( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) + and node.name == name + or isinstance(node, ast.alias) + and (node.asname or node.name.split(".")[0]) == name + for node in ast.walk(tree) + ) + if len(nodes) != 1 or matches != [nodes[0].targets[0]] or rebound: + raise ValueError(f"Expected one top-level assignment to {name}") + return nodes[0] + + +def _selected_export_addition(src, dst, source_lines, lines, exports): + src_all, dst_all = ( + _single_assignment(src, "__all__"), + _single_assignment(dst, "__all__"), + ) + if not isinstance(src_all.value, ast.List) or not isinstance( + dst_all.value, ast.List + ): + raise ValueError("__all__ must be a literal list") + if any( + not isinstance(node, ast.Constant) or not isinstance(node.value, str) + for node in src_all.value.elts + dst_all.value.elts + ): + raise ValueError("__all__ entries must be literal strings") + src_values, dst_values = ( + ast.literal_eval(src_all.value), + ast.literal_eval(dst_all.value), + ) + if any(src_values.count(name) != 1 or name in dst_values for name in exports): + raise ValueError( + "Exports must occur once in canonical __all__ and be absent locally" + ) + export_nodes = [node for node in src_all.value.elts if node.value in exports] + export_lines = sorted({node.lineno - 1 for node in export_nodes}) + fragment = "".join(source_lines[line] for line in export_lines) + if ast.literal_eval("[\n" + fragment + "]") != exports: + raise ValueError( + "Canonical export lines include unselected entries or differ in order" + ) + if lines[dst_all.lineno - 1].strip() != "__all__ = [": + raise ValueError("Expected a multiline destination __all__ list") + return dst_all.lineno, [fragment] + + def copy_selected_assignments(source, destination, names, exports): """Copy contiguous canonical assignments and export lines, failing closed. @@ -74,34 +143,7 @@ def copy_selected_assignments(source, destination, names, exports): raise ValueError("Exports must be unique selected assignment names") src, dst = ast.parse(source), ast.parse(destination) - def assignment(tree, name): - matches = [ - node - for node in ast.walk(tree) - if isinstance(node, ast.Name) - and isinstance(node.ctx, ast.Store) - and node.id == name - ] - nodes = [ - node - for node in tree.body - if isinstance(node, ast.Assign) - and len(node.targets) == 1 - and isinstance(node.targets[0], ast.Name) - and node.targets[0].id == name - ] - rebound = any( - isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) - and node.name == name - or isinstance(node, ast.alias) - and (node.asname or node.name.split(".")[0]) == name - for node in ast.walk(tree) - ) - if len(nodes) != 1 or matches != [nodes[0].targets[0]] or rebound: - raise ValueError(f"Expected one top-level assignment to {name}") - return nodes[0] - - selected = [assignment(src, name) for name in names] + selected = [_single_assignment(src, name) for name in names] start = src.body.index(selected[0]) if src.body[start : start + len(names)] != selected or start + len(names) == len( src.body @@ -223,34 +265,9 @@ def copy_selected_assignments(source, destination, names, exports): ) ] if exports: - src_all, dst_all = assignment(src, "__all__"), assignment(dst, "__all__") - if not isinstance(src_all.value, ast.List) or not isinstance( - dst_all.value, ast.List - ): - raise ValueError("__all__ must be a literal list") - if any( - not isinstance(node, ast.Constant) or not isinstance(node.value, str) - for node in src_all.value.elts + dst_all.value.elts - ): - raise ValueError("__all__ entries must be literal strings") - src_values, dst_values = ( - ast.literal_eval(src_all.value), - ast.literal_eval(dst_all.value), + additions.append( + _selected_export_addition(src, dst, source_lines, lines, exports) ) - if any(src_values.count(name) != 1 or name in dst_values for name in exports): - raise ValueError( - "Exports must occur once in canonical __all__ and be absent locally" - ) - export_nodes = [node for node in src_all.value.elts if node.value in exports] - export_lines = sorted({node.lineno - 1 for node in export_nodes}) - fragment = "".join(source_lines[line] for line in export_lines) - if ast.literal_eval("[\n" + fragment + "]") != exports: - raise ValueError( - "Canonical export lines include unselected entries or differ in order" - ) - if lines[dst_all.lineno - 1].strip() != "__all__ = [": - raise ValueError("Expected a multiline destination __all__ list") - additions.append((dst_all.lineno, [fragment])) for offset, fragment in sorted(additions, reverse=True): lines[offset:offset] = fragment result = "".join(lines) @@ -258,10 +275,291 @@ def copy_selected_assignments(source, destination, names, exports): return result -def migrate_selected_assignments(path, names, exports, source_hash, destination_hash): +def _module_bindings(tree): + """Collect possible module bindings without mistaking function locals for them.""" + bindings = collections.defaultdict(list) + + def visit(node): + children = ast.iter_child_nodes(node) + names = [] + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + names = [node.name] + # A declaration in a nested scope can rebind a module dependency. + for child in ast.walk(node): + if isinstance(child, ast.Global): + for name in child.names: + bindings[name].append(child) + children = list(node.decorator_list) + if isinstance(node, ast.ClassDef): + children += node.bases + node.keywords + else: + children += node.args.defaults + [ + value for value in node.args.kw_defaults if value is not None + ] + elif isinstance(node, ast.Lambda): + children = node.args.defaults + [ + value for value in node.args.kw_defaults if value is not None + ] + elif isinstance(node, ast.Name) and isinstance(node.ctx, (ast.Store, ast.Del)): + names = [node.id] + elif isinstance(node, ast.alias): + names = [node.asname or node.name.split(".")[0]] + elif isinstance(node, (ast.ExceptHandler, ast.MatchAs, ast.MatchStar)): + names = [node.name] if node.name is not None else [] + elif isinstance(node, ast.MatchMapping) and node.rest is not None: + names = [node.rest] + for name in names: + bindings[name].append(node) + for child in children: + visit(child) + + visit(tree) + return bindings + + +def _function_dependencies(fragment): + """Find definition-time and deferred global reads without executing source.""" + table = symtable.symtable(fragment, "", "exec") + if any( + symbol.is_assigned() and not symbol.is_namespace() + for symbol in table.get_symbols() + ): + raise ValueError("Function headers cannot bind additional module names") + immediate = { + symbol.get_name() for symbol in table.get_symbols() if symbol.is_referenced() + } + deferred = set() + + def visit(scope): + for symbol in scope.get_symbols(): + if symbol.is_global(): + if symbol.is_assigned() or symbol.is_imported(): + raise ValueError("Selected functions cannot rebind module globals") + if symbol.is_referenced(): + deferred.add(symbol.get_name()) + for child in scope.get_children(): + visit(child) + + for child in table.get_children(): + visit(child) + if (immediate | deferred) & {"eval", "exec", "globals", "locals", "vars"}: + raise ValueError("Dynamic namespace dependencies require a full migration") + return immediate, deferred + + +def copy_selected_functions(source, destination, names, exports): + """Add contiguous complete functions, preserving unrelated destination bytes. + + Global dependencies must be builtins, selected functions, or unambiguously + bound identical imports/simple assignments (checked transitively). Existing + helper functions and dynamic namespace dependencies require a full migration. + Scoped expressions in definition-time headers and classes anywhere inside + selected functions also require a full migration. Symbol-table analysis + distinguishes body locals, closures and comprehensions from module + dependencies; it never imports or calls the selected functions. + """ + if ( + not names + or len(set(names)) != len(names) + or not all(name.isidentifier() for name in names) + ): + raise ValueError("Function names must be unique identifiers") + if len(set(exports)) != len(exports) or not set(exports) <= set(names): + raise ValueError("Exports must be unique selected function names") + src, dst = ast.parse(source), ast.parse(destination) + src_bindings, dst_bindings = _module_bindings(src), _module_bindings(dst) + selected = [] + for name in names: + matches = src_bindings.get(name, []) + if ( + len(matches) != 1 + or not isinstance(matches[0], (ast.FunctionDef, ast.AsyncFunctionDef)) + or matches[0] not in src.body + or getattr(matches[0], "type_params", ()) + ): + raise ValueError(f"Expected one top-level non-generic function: {name}") + if dst_bindings.get(name): + raise ValueError( + f"A selected function already exists in the destination: {name}" + ) + node = matches[0] + headers = ( + node.decorator_list + + node.args.defaults + + [value for value in node.args.kw_defaults if value is not None] + ) + arguments = node.args.posonlyargs + node.args.args + node.args.kwonlyargs + arguments += [arg for arg in (node.args.vararg, node.args.kwarg) if arg] + headers += [arg.annotation for arg in arguments if arg.annotation is not None] + if node.returns is not None: + headers.append(node.returns) + if any( + isinstance( + child, + ( + ast.Lambda, + ast.ListComp, + ast.SetComp, + ast.DictComp, + ast.GeneratorExp, + ), + ) + for header in headers + for child in ast.walk(header) + ): + raise ValueError( + "Scoped expressions in function headers require a full migration" + ) + if any(isinstance(child, ast.ClassDef) for child in ast.walk(node)): + raise ValueError("Classes in selected functions require a full migration") + selected.append(node) + start = src.body.index(selected[0]) + if src.body[start : start + len(names)] != selected or start + len(names) == len( + src.body + ): + raise ValueError("Functions must be contiguous and followed by an anchor") + first_line = min( + [selected[0].lineno, *(node.lineno for node in selected[0].decorator_list)] + ) + anchor = src.body[start + len(names)] + matches = [node for node in dst.body if ast.dump(node) == ast.dump(anchor)] + if len(matches) != 1: + raise ValueError("Expected one unchanged destination anchor") + for tree, node in ((src, anchor), (dst, matches[0])): + index = tree.body.index(node) + if ( + getattr(node, "decorator_list", ()) + or node.col_offset + or index + and tree.body[index - 1].end_lineno >= node.lineno + or index + 1 < len(tree.body) + and tree.body[index + 1].lineno <= node.end_lineno + ): + raise ValueError( + "Anchor must be undecorated and occupy complete source lines" + ) + if start and src.body[start - 1].end_lineno >= first_line: + raise ValueError("Selected functions must occupy complete source lines") + + source_lines, lines = source.splitlines(True), destination.splitlines(True) + futures = [ + [ + node + for node in tree.body + if isinstance(node, ast.ImportFrom) and node.module == "__future__" + ] + for tree in (src, dst) + ] + if [ast.dump(node) for node in futures[0]] != [ + ast.dump(node) for node in futures[1] + ]: + raise ValueError("Source and destination future imports differ") + prefix = "".join( + "".join(source_lines[node.lineno - 1 : node.end_lineno]) for node in futures[0] + ) + if "*" in src_bindings or "*" in dst_bindings: + raise ValueError("Wildcard import makes dependency bindings ambiguous") + + def dependency(name, source_boundary, destination_boundary): + if name in {"eval", "exec", "globals", "locals", "vars"}: + raise ValueError("Dynamic namespace dependencies require a full migration") + declarations = [] + for tree, bindings, boundary in ( + (src, src_bindings, source_boundary), + (dst, dst_bindings, destination_boundary), + ): + bound = bindings.get(name, []) + if not bound and name in vars(builtins): + declarations.append(None) + continue + if len(bound) != 1: + raise ValueError(f"Missing or ambiguous dependency binding: {name}") + candidates = [ + node + for node in tree.body + if isinstance(node, (ast.Import, ast.ImportFrom)) + and bound[0] in node.names + or isinstance(node, ast.Assign) + and len(node.targets) == 1 + and node.targets[0] is bound[0] + ] + if len(candidates) != 1 or candidates[0].lineno >= boundary: + raise ValueError( + f"Expected an earlier import or simple assignment: {name}" + ) + declarations.append(candidates[0]) + left, right = declarations + if left is None and right is None: + return + if left is None or right is None or ast.dump(left) != ast.dump(right): + raise ValueError(f"Changed dependency binding: {name}") + if isinstance(left, ast.Assign): + for child in ast.walk(left.value): + if isinstance( + child, + ( + ast.Lambda, + ast.NamedExpr, + ast.ListComp, + ast.SetComp, + ast.DictComp, + ast.GeneratorExp, + ), + ): + raise ValueError(f"Unsupported dependency expression: {name}") + if isinstance(child, ast.Name) and isinstance(child.ctx, ast.Load): + dependency(child.id, left.lineno, right.lineno) + + earlier = set() + for node in selected: + line = min([node.lineno, *(item.lineno for item in node.decorator_list)]) + immediate, deferred = _function_dependencies( + prefix + "".join(source_lines[line - 1 : node.end_lineno]) + ) + if (immediate & set(names)) - earlier: + raise ValueError("Definition-time dependency is not yet bound") + for name in (immediate | deferred) - set(names): + dependency(name, first_line, matches[0].lineno) + earlier.add(node.name) + + # Keep a following anchor comment attached to its original statement. Only + # canonical function bytes and their trailing whitespace/comments are added. + end = anchor.lineno - 1 + comment = next( + ( + i + for i in range(selected[-1].end_lineno, end) + if source_lines[i].startswith("#") + ), + end, + ) + anchor_comments = source_lines[comment:end] + offset = matches[0].lineno - 1 - len(anchor_comments) + if offset < 0 or lines[offset : matches[0].lineno - 1] != anchor_comments: + raise ValueError("Anchor's leading comments differ") + additions = [(offset, source_lines[first_line - 1 : comment])] + if exports: + for tree, bindings in ((src, src_bindings), (dst, dst_bindings)): + if bindings.get("__all__") != [ + _single_assignment(tree, "__all__").targets[0] + ]: + raise ValueError("Ambiguous __all__ binding") + additions.append( + _selected_export_addition(src, dst, source_lines, lines, exports) + ) + for offset, fragment in sorted(additions, reverse=True): + lines[offset:offset] = fragment + result = "".join(lines) + ast.parse(result) + return result + + +def _migrate_selected( + path, names, exports, source_hash, destination_hash, kind, copier +): target = parse_lib_path(path) if target.parts[:1] != ("Lib",) or ".." in target.parts or target.suffix != ".py": - raise ValueError("Selective assignments require an explicit Lib Python file") + raise ValueError("Selective imports require an explicit Lib Python file") source, destination = path.read_bytes(), target.read_bytes() for label, contents, expected in ( ("source", source, source_hash), @@ -269,16 +567,16 @@ def migrate_selected_assignments(path, names, exports, source_hash, destination_ ): if hashlib.sha256(contents).hexdigest() != expected: raise ValueError(f"{label} SHA256 does not match the pinned input") - result = copy_selected_assignments( + result = copier( source.decode("utf-8"), destination.decode("utf-8"), names, exports ).encode("utf-8") if target.read_bytes() != destination: - raise ValueError("Destination changed during assignment preparation") + raise ValueError("Destination changed during selective import preparation") target.write_bytes(result) print( json.dumps( { - "assignments": names, + kind: names, "exports": exports, "source_sha256": source_hash, "destination_before_sha256": destination_hash, @@ -289,6 +587,30 @@ def migrate_selected_assignments(path, names, exports, source_hash, destination_ ) +def migrate_selected_assignments(path, names, exports, source_hash, destination_hash): + return _migrate_selected( + path, + names, + exports, + source_hash, + destination_hash, + "assignments", + copy_selected_assignments, + ) + + +def migrate_selected_functions(path, names, exports, source_hash, destination_hash): + return _migrate_selected( + path, + names, + exports, + source_hash, + destination_hash, + "functions", + copy_selected_functions, + ) + + def patch_single_content( src_path: pathlib.Path, lib_path: pathlib.Path, @@ -781,6 +1103,11 @@ def main(argv: list[str] | None = None) -> int: nargs="+", help="Copy selected canonical assignments", ) + parser.add_argument( + "--functions", + nargs="+", + help="Copy complete selected canonical functions without executing them", + ) parser.add_argument("--exports", nargs="+", default=[]) parser.add_argument("--source-sha256") parser.add_argument("--destination-sha256") @@ -792,6 +1119,7 @@ def main(argv: list[str] | None = None) -> int: args.legacy_map is not None or args.patches is not None or args.assignments + or args.functions or args.exports or args.source_sha256 or args.destination_sha256 @@ -808,15 +1136,25 @@ def main(argv: list[str] | None = None) -> int: return 0 if args.dry_run: raise ValueError("--dry-run requires --canonical-removals") - if args.assignments: - if args.legacy_map is not None or args.patches is not None: + if args.assignments or args.functions: + if ( + args.legacy_map is not None + or args.patches is not None + or args.assignments + and args.functions + ): raise ValueError( - "Selective imports cannot combine with --legacy-map or --patches" + "Selective imports cannot combine with other migration modes" ) try: - migrate_selected_assignments( + migrate = ( + migrate_selected_functions + if args.functions + else migrate_selected_assignments + ) + migrate( args.path, - args.assignments, + args.functions or args.assignments, args.exports, args.source_sha256, args.destination_sha256, @@ -825,7 +1163,7 @@ def main(argv: list[str] | None = None) -> int: raise ValueError(str(exc)) from exc return 0 if args.exports or args.source_sha256 or args.destination_sha256: - raise ValueError("Selective options require --assignments") + raise ValueError("Selective options require --assignments or --functions") if args.patches is not None and ( args.path.is_dir() or args.legacy_map is not None ): diff --git a/scripts/update_lib/tests/test_selected_functions.py b/scripts/update_lib/tests/test_selected_functions.py new file mode 100644 index 000000000..0de48a63f --- /dev/null +++ b/scripts/update_lib/tests/test_selected_functions.py @@ -0,0 +1,320 @@ +"""Source-only selective function migration, without importing copied helpers.""" + +import contextlib +import hashlib +import io +import json +import pathlib +import tempfile +import unittest +from unittest.mock import patch + +from update_lib.cmd_migrate import copy_selected_functions, main + + +class TestSelectedFunctions(unittest.TestCase): + names = ["available", "requires_available"] + exports = names + imports = "import functools\nimport os\nimport sys\nimport unittest\n\n" + export_line = ' "available", "requires_available",\n' + dependencies = 'IS_MOBILE = sys.platform == "mobile"\nAVAILABLE = not IS_MOBILE\n\n' + block = ( + "@functools.cache\n" + "def available(limit=1):\n" + " import errno\n" + " try:\n" + ' return AVAILABLE and bool(errno.EINVAL) and os.name != "" and limit > 0\n' + " except ValueError as error:\n" + " return bool(str(error))\n\n\n" + "def requires_available():\n" + ' return unittest.skipUnless(available(), "needs availability")\n\n\n' + ) + anchor = "# Keep this comment with STATUS.\nSTATUS = 1\n" + source = ( + imports + + '__all__ = [\n "existing",\n' + + export_line + + "]\n\n" + + dependencies + + block + + anchor + ) + destination = ( + imports + + '__all__ = [\n "existing",\n]\n\n' + + dependencies + + 'def existing(sys):\n return "keep π"\n\n' + + anchor + ) + + def copy(self, source=None, destination=None, names=None, exports=None): + return copy_selected_functions( + self.source if source is None else source, + self.destination if destination is None else destination, + self.names if names is None else names, + self.exports if exports is None else exports, + ) + + def test_exact_functions_exports_and_unrelated_bytes(self): + self.assertEqual( + self.copy(), + self.destination.replace( + "__all__ = [\n", "__all__ = [\n" + self.export_line + ).replace(self.anchor, self.block + self.anchor), + ) + self.assertEqual(self.copy().count(self.anchor), 1) + + def test_destination_line_endings_and_local_adaptations_survive(self): + source = self.source.replace(self.anchor, "STATUS = 1\n") + destination = self.destination.replace(self.anchor, "STATUS = 1\n").replace( + "\n", "\r\n" + ) + self.assertEqual( + self.copy(source, destination), + destination.replace( + "__all__ = [\r\n", "__all__ = [\r\n" + self.export_line + ).replace("STATUS = 1\r\n", self.block + "STATUS = 1\r\n"), + ) + + def test_local_imports_closures_and_comprehensions_are_not_global_dependencies( + self, + ): + block = ( + "@functools.cache\n" + "def available(items=()):\n" + " def inner(value):\n" + " return value in items\n" + " return [inner(value) for value in items]\n\n\n" + "async def requires_available():\n" + " import errno as local\n" + " return local.EINVAL\n\n\n" + ) + self.assertIn(block, self.copy(self.source.replace(self.block, block))) + with self.assertRaises(ValueError): + self.copy( + self.source.replace( + self.block, block.replace("value in items", "MISSING", 1) + ) + ) + + def test_multiple_decorators_defaults_and_body_are_copied_without_execution(self): + block = self.block.replace( + "@functools.cache\ndef available(limit=1):", + "@functools.cache\n@functools.lru_cache(\n maxsize=None\n)\n" + "def available(limit=os.name):", + ).replace( + " import errno\n", + " raise RuntimeError('must not run')\n import errno\n", + ) + self.assertIn(block, self.copy(self.source.replace(self.block, block))) + + def test_selection_and_existing_bindings_fail_closed(self): + for source, destination, names in ( + (self.source, self.destination, []), + (self.source, self.destination, self.names * 2), + (self.source, self.destination, list(reversed(self.names))), + (self.source, self.destination, ["missing"]), + (self.source + "available = None\n", self.destination, self.names), + (self.source, self.destination + "del available\n", self.names), + (self.source, self.destination + "import other as available\n", self.names), + ( + self.source, + self.destination + "if condition:\n available = None\n", + self.names, + ), + ( + self.source.replace( + "def requires_available", "UNSELECTED = 1\ndef requires_available" + ), + self.destination, + self.names, + ), + (self.source.replace(self.anchor, ""), self.destination, self.names), + ): + with self.subTest(names=names), self.assertRaises(ValueError): + self.copy(source, destination, names) + + def test_missing_changed_late_and_rebound_dependencies_fail_closed(self): + for destination in ( + self.destination.replace("import functools\n", ""), + self.destination.replace( + "import functools\n", "import other as functools\n" + ), + self.destination.replace("import os\n", "") + "import os\n", + self.destination.replace( + 'sys.platform == "mobile"', 'sys.platform == "other"' + ), + self.destination + "AVAILABLE = True\n", + self.destination + "del sys\n", + self.destination + "bool = None\n", + self.destination + "def rebind():\n global sys\n sys = None\n", + self.destination + "from unknown import *\n", + ): + with self.subTest(destination=destination), self.assertRaises(ValueError): + self.copy(destination=destination) + for source in ( + self.source + "del sys\n", + self.source + "from unknown import *\n", + self.source.replace("AVAILABLE = not IS_MOBILE", "AVAILABLE = MISSING"), + self.source.replace("limit=1", "limit=requires_available()"), + ): + with self.subTest(source=source), self.assertRaises(ValueError): + self.copy(source=source) + + def test_unsupported_namespace_changes_and_future_mismatch_fail_closed(self): + for statement in ( + "global AVAILABLE\n AVAILABLE = True", + "return globals()", + "return eval('AVAILABLE')", + ): + with self.subTest(statement=statement), self.assertRaises(ValueError): + self.copy( + self.source.replace( + " import errno", " " + statement + "\n import errno" + ) + ) + with self.assertRaises(ValueError): + self.copy("from __future__ import annotations\n" + self.source) + with self.assertRaises(ValueError): + self.copy(self.source.replace("limit=1", "limit=(EXTRA := 1)")) + + def test_anchor_ambiguity_partial_lines_and_changed_comments_fail_closed(self): + for destination in ( + self.destination.replace("STATUS = 1", "STATUS = 2"), + self.destination + "STATUS = 1\n", + self.destination.replace("STATUS = 1", "STATUS = 1; other = 2"), + self.destination.replace("STATUS = 1", "other = 2; STATUS = 1"), + self.destination.replace( + "# Keep this comment with STATUS.", "# Changed comment." + ), + ): + with self.subTest(destination=destination), self.assertRaises(ValueError): + self.copy(destination=destination) + replacement = "@functools.cache\ndef anchor(): pass\n" + with self.assertRaises(ValueError): + self.copy( + self.source.replace("STATUS = 1\n", replacement), + self.destination.replace("STATUS = 1\n", replacement), + ) + + def test_exports_require_complete_canonical_lines(self): + for exports in ( + self.exports[:1], + list(reversed(self.exports)), + self.exports * 2, + ["other"], + ): + with self.subTest(exports=exports), self.assertRaises(ValueError): + self.copy(exports=exports) + with self.assertRaises(ValueError): + self.copy(destination=self.destination + "del __all__\n") + self.assertNotIn(self.export_line, self.copy(exports=[])) + + def test_cli_hash_pins_receipt_and_concurrent_destination_guard(self): + with tempfile.TemporaryDirectory() as temp, contextlib.chdir(temp): + source = pathlib.Path("canonical/Lib/example.py") + destination = pathlib.Path("Lib/example.py") + source.parent.mkdir(parents=True) + destination.parent.mkdir(parents=True) + source.write_text(self.source) + destination.write_text(self.destination) + args = [ + str(source), + "--functions", + *self.names, + "--exports", + *self.exports, + "--source-sha256", + hashlib.sha256(source.read_bytes()).hexdigest(), + "--destination-sha256", + "0" * 64, + ] + with contextlib.redirect_stderr(io.StringIO()): + self.assertEqual(main(args), 1) + self.assertEqual(destination.read_text(), self.destination) + args[-1] = hashlib.sha256(destination.read_bytes()).hexdigest() + with contextlib.redirect_stdout(io.StringIO()) as output: + self.assertEqual(main(args), 0) + receipt = json.loads(output.getvalue()) + self.assertEqual(receipt["functions"], self.names) + self.assertEqual(destination.read_text(), self.copy()) + destination.write_text(self.destination) + + def concurrent_change(*unused): + destination.write_text("# changed meanwhile\n") + return self.copy() + + with patch( + "update_lib.cmd_migrate.copy_selected_functions", concurrent_change + ): + with contextlib.redirect_stderr(io.StringIO()) as errors: + self.assertEqual(main(args), 1) + self.assertIn("Destination changed", errors.getvalue()) + self.assertEqual(destination.read_text(), "# changed meanwhile\n") + + def test_cli_modes_are_exclusive_before_reading_inputs(self): + for options in ( + ["--assignments", "value"], + ["--patches", "missing.json"], + ["--legacy-map", "missing.json"], + ["--canonical-removals", "missing.json"], + ): + with ( + self.subTest(options=options), + contextlib.redirect_stderr(io.StringIO()) as errors, + ): + self.assertEqual( + main( + ["missing/Lib/example.py", "--functions", *self.names, *options] + ), + 1, + ) + self.assertIn("cannot combine", errors.getvalue()) + + def test_scoped_header_expressions_fail_closed(self): + for header in ( + "def available(limit=(lambda: requires_available())()):", + "def available(limit=[requires_available() for _ in (0,)]):", + "def available(limit={requires_available() for _ in (0,)}):", + "def available(limit={0: requires_available() for _ in (0,)}):", + "def available(limit=(requires_available() for _ in (0,))):", + "def available(*, limit=(lambda: requires_available())()):", + "def available(limit: (lambda: requires_available())() = 1):", + ): + with ( + self.subTest(header=header), + self.assertRaisesRegex( + ValueError, "Scoped expressions in function headers" + ), + ): + self.copy(self.source.replace("def available(limit=1):", header)) + for decorator in ( + "@(lambda: requires_available())()", + "@[requires_available for _ in (0,)][0]", + ): + with ( + self.subTest(decorator=decorator), + self.assertRaisesRegex( + ValueError, "Scoped expressions in function headers" + ), + ): + self.copy(self.source.replace("@functools.cache", decorator)) + + def test_nested_classes_fail_closed(self): + for body in ( + " class Inner:\n os = os\n", + " def nested():\n class Inner:\n os = os\n", + ): + with ( + self.subTest(body=body), + self.assertRaisesRegex(ValueError, "Classes in selected functions"), + ): + self.copy( + self.source.replace( + " import errno\n", body + " import errno\n" + ) + ) + + +if __name__ == "__main__": + unittest.main()