Skip to content

Commit a7fa762

Browse files
pablogsalpablogsal
committed
Protect against more errors with no exceptions
1 parent 0338ee1 commit a7fa762

File tree

3 files changed

+312
-25
lines changed

3 files changed

+312
-25
lines changed

‎Lib/profile/_sync_coordinator.py

Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,238 @@
1+
"""
2+
Internal synchronization coordinator for the sample profiler.
3+
4+
This module is used internally by the sample profiler to coordinate
5+
the startup of target processes. It should not be called directly by users.
6+
"""
7+
8+
import os
9+
import sys
10+
import socket
11+
import runpy
12+
import time
13+
from typing import List, NoReturn
14+
15+
16+
class CoordinatorError(Exception):
17+
"""Base exception for coordinator errors."""
18+
pass
19+
20+
21+
class ArgumentError(CoordinatorError):
22+
"""Raised when invalid arguments are provided."""
23+
pass
24+
25+
26+
class SyncError(CoordinatorError):
27+
"""Raised when synchronization with profiler fails."""
28+
pass
29+
30+
31+
class TargetError(CoordinatorError):
32+
"""Raised when target execution fails."""
33+
pass
34+
35+
36+
def _validate_arguments(args: List[str]) -> tuple[int, str, List[str]]:
37+
"""
38+
Validate and parse command line arguments.
39+
40+
Args:
41+
args: Command line arguments including script name
42+
43+
Returns:
44+
Tuple of (sync_port, working_directory, target_args)
45+
46+
Raises:
47+
ArgumentError: If arguments are invalid
48+
"""
49+
if len(args) < 4:
50+
raise ArgumentError(
51+
"Insufficient arguments. Expected: <sync_port> <cwd> <target> [args...]"
52+
)
53+
54+
try:
55+
sync_port = int(args[1])
56+
if not (1 <= sync_port <= 65535):
57+
raise ValueError("Port out of range")
58+
except ValueError as e:
59+
raise ArgumentError(f"Invalid sync port '{args[1]}': {e}") from e
60+
61+
cwd = args[2]
62+
if not os.path.isdir(cwd):
63+
raise ArgumentError(f"Working directory does not exist: {cwd}")
64+
65+
target_args = args[3:]
66+
if not target_args:
67+
raise ArgumentError("No target specified")
68+
69+
return sync_port, cwd, target_args
70+
71+
72+
def _signal_readiness(sync_port: int) -> None:
73+
"""
74+
Signal readiness to the profiler via TCP socket.
75+
76+
Args:
77+
sync_port: Port number where profiler is listening
78+
79+
Raises:
80+
SyncError: If unable to signal readiness
81+
"""
82+
max_retries = 3
83+
retry_delay = 0.1
84+
85+
for attempt in range(max_retries):
86+
try:
87+
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
88+
sock.settimeout(2.0) # 2 second timeout
89+
try:
90+
sock.connect(("127.0.0.1", sync_port))
91+
sock.send(b"ready")
92+
return
93+
finally:
94+
sock.close()
95+
except (socket.error, OSError) as e:
96+
if attempt == max_retries - 1:
97+
# On final attempt, raise the error
98+
raise SyncError(f"Failed to signal readiness after {max_retries} attempts: {e}") from e
99+
# Wait before retry
100+
time.sleep(retry_delay * (2 ** attempt)) # Exponential backoff
101+
102+
103+
def _setup_environment(cwd: str) -> None:
104+
"""
105+
Set up the execution environment.
106+
107+
Args:
108+
cwd: Working directory to change to
109+
110+
Raises:
111+
TargetError: If unable to set up environment
112+
"""
113+
try:
114+
os.chdir(cwd)
115+
except OSError as e:
116+
raise TargetError(f"Failed to change to directory {cwd}: {e}") from e
117+
118+
# Add current directory to sys.path if not present (for module imports)
119+
if cwd not in sys.path:
120+
sys.path.insert(0, cwd)
121+
122+
123+
def _execute_module(module_name: str, module_args: List[str]) -> None:
124+
"""
125+
Execute a Python module.
126+
127+
Args:
128+
module_name: Name of the module to execute
129+
module_args: Arguments to pass to the module
130+
131+
Raises:
132+
TargetError: If module execution fails
133+
"""
134+
# Replace sys.argv to match original module call
135+
sys.argv = ["-m", module_name] + module_args
136+
137+
try:
138+
runpy.run_module(module_name, run_name="__main__", alter_sys=True)
139+
except ImportError as e:
140+
raise TargetError(f"Module '{module_name}' not found: {e}") from e
141+
except SystemExit:
142+
# SystemExit is normal for modules
143+
pass
144+
except Exception as e:
145+
raise TargetError(f"Error executing module '{module_name}': {e}") from e
146+
147+
148+
def _execute_script(script_path: str, script_args: List[str], cwd: str) -> None:
149+
"""
150+
Execute a Python script.
151+
152+
Args:
153+
script_path: Path to the script to execute
154+
script_args: Arguments to pass to the script
155+
cwd: Current working directory for path resolution
156+
157+
Raises:
158+
TargetError: If script execution fails
159+
"""
160+
# Make script path absolute if it isn't already
161+
if not os.path.isabs(script_path):
162+
script_path = os.path.join(cwd, script_path)
163+
164+
if not os.path.isfile(script_path):
165+
raise TargetError(f"Script not found: {script_path}")
166+
167+
# Replace sys.argv to match original script call
168+
sys.argv = [script_path] + script_args
169+
170+
try:
171+
with open(script_path, 'rb') as f:
172+
source_code = f.read()
173+
174+
# Compile and execute the script
175+
code = compile(source_code, script_path, 'exec')
176+
exec(code, {'__name__': '__main__', '__file__': script_path})
177+
except FileNotFoundError as e:
178+
raise TargetError(f"Script file not found: {script_path}") from e
179+
except PermissionError as e:
180+
raise TargetError(f"Permission denied reading script: {script_path}") from e
181+
except SyntaxError as e:
182+
raise TargetError(f"Syntax error in script {script_path}: {e}") from e
183+
except SystemExit:
184+
# SystemExit is normal for scripts
185+
pass
186+
except Exception as e:
187+
raise TargetError(f"Error executing script '{script_path}': {e}") from e
188+
189+
190+
def main() -> NoReturn:
191+
"""
192+
Main coordinator function.
193+
194+
This function coordinates the startup of a target Python process
195+
with the sample profiler by signaling when the process is ready
196+
to be profiled.
197+
"""
198+
try:
199+
# Parse and validate arguments
200+
sync_port, cwd, target_args = _validate_arguments(sys.argv)
201+
202+
# Set up execution environment
203+
_setup_environment(cwd)
204+
205+
# Signal readiness to profiler
206+
_signal_readiness(sync_port)
207+
208+
# Execute the target
209+
if target_args[0] == "-m":
210+
# Module execution
211+
if len(target_args) < 2:
212+
raise ArgumentError("Module name required after -m")
213+
214+
module_name = target_args[1]
215+
module_args = target_args[2:]
216+
_execute_module(module_name, module_args)
217+
else:
218+
# Script execution
219+
script_path = target_args[0]
220+
script_args = target_args[1:]
221+
_execute_script(script_path, script_args, cwd)
222+
223+
except CoordinatorError as e:
224+
print(f"Profiler coordinator error: {e}", file=sys.stderr)
225+
sys.exit(1)
226+
except KeyboardInterrupt:
227+
print("Interrupted", file=sys.stderr)
228+
sys.exit(1)
229+
except Exception as e:
230+
print(f"Unexpected error in profiler coordinator: {e}", file=sys.stderr)
231+
sys.exit(1)
232+
233+
# Normal exit
234+
sys.exit(0)
235+
236+
237+
if __name__ == "__main__":
238+
main()

‎Lib/profile/sample.py

Lines changed: 65 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import _remote_debugging
33
import os
44
import pstats
5+
import socket
56
import subprocess
67
import statistics
78
import sys
@@ -60,6 +61,54 @@
6061
python -m profile.sample --sort-nsamples-cumul -p 1234"""
6162

6263

64+
def _run_with_sync(original_cmd):
65+
"""Run a command with socket-based synchronization and return the process."""
66+
# Create a TCP socket for synchronization
67+
sync_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
68+
sync_sock.bind(("127.0.0.1", 0)) # Let OS choose a free port
69+
sync_port = sync_sock.getsockname()[1]
70+
sync_sock.listen(1)
71+
72+
# Get current working directory to preserve it
73+
cwd = os.getcwd()
74+
75+
try:
76+
# Build command using the sync coordinator
77+
target_args = original_cmd[1:] # Remove python executable
78+
cmd = (sys.executable, "-m", "profile._sync_coordinator", str(sync_port), cwd) + tuple(target_args)
79+
80+
# Start the process with coordinator
81+
process = subprocess.Popen(cmd)
82+
83+
# Wait for ready signal with timeout
84+
sync_sock.settimeout(5.0) # 5 second timeout
85+
try:
86+
conn, addr = sync_sock.accept()
87+
ready_signal = conn.recv(1024)
88+
conn.close()
89+
90+
if ready_signal != b"ready":
91+
raise RuntimeError("Invalid ready signal received")
92+
93+
except socket.timeout:
94+
# If we timeout, kill the process and raise an error
95+
if process.poll() is None:
96+
process.terminate()
97+
try:
98+
process.wait(timeout=2)
99+
except subprocess.TimeoutExpired:
100+
process.kill()
101+
process.wait()
102+
raise RuntimeError("Process failed to signal readiness within timeout")
103+
104+
return process
105+
106+
finally:
107+
sync_sock.close()
108+
109+
110+
111+
63112
class SampleProfiler:
64113
def __init__(self, pid, sample_interval_usec, all_threads):
65114
self.pid = pid
@@ -585,27 +634,19 @@ def _validate_collapsed_format_args(args, parser):
585634

586635

587636
def wait_for_process_and_sample(pid, sort_value, args):
588-
for attempt in range(_MAX_STARTUP_ATTEMPTS):
589-
try:
590-
sample(
591-
pid,
592-
sort=sort_value,
593-
sample_interval_usec=args.interval,
594-
duration_sec=args.duration,
595-
filename=args.outfile,
596-
all_threads=args.all_threads,
597-
limit=args.limit,
598-
show_summary=not args.no_summary,
599-
output_format=args.format,
600-
realtime_stats=args.realtime_stats,
601-
)
602-
break
603-
except RuntimeError:
604-
if attempt < _MAX_STARTUP_ATTEMPTS - 1:
605-
print("Waiting for process to start...")
606-
time.sleep(_STARTUP_RETRY_DELAY_SECONDS)
607-
else:
608-
raise RuntimeError("Process failed to start after maximum retries") from None
637+
"""Sample the process immediately since it has already signaled readiness."""
638+
sample(
639+
pid,
640+
sort=sort_value,
641+
sample_interval_usec=args.interval,
642+
duration_sec=args.duration,
643+
filename=args.outfile,
644+
all_threads=args.all_threads,
645+
limit=args.limit,
646+
show_summary=not args.no_summary,
647+
output_format=args.format,
648+
realtime_stats=args.realtime_stats,
649+
)
609650

610651

611652
def main():
@@ -782,10 +823,10 @@ def main():
782823
else:
783824
cmd = (sys.executable, *args.script)
784825

785-
process = subprocess.Popen(cmd)
826+
# Use synchronized process startup
827+
process = _run_with_sync(cmd)
786828

787-
# If we are the ones starting the process, we need to wait until the
788-
# runtime state is initialized
829+
# Process has already signaled readiness, start sampling immediately
789830
try:
790831
wait_for_process_and_sample(process.pid, sort_value, args)
791832
finally:

‎Modules/_remote_debugging_module.c

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@
7575
#endif
7676
#define INTERP_STATE_BUFFER_SIZE MAX(INTERP_STATE_MIN_SIZE, 256)
7777

78-
78+
#define MAX_TLBC_SIZE 2048
7979

8080
// Copied from Modules/_asynciomodule.c because it's not exported
8181

@@ -1594,10 +1594,16 @@ cache_tlbc_array(RemoteUnwinderObject *unwinder, uintptr_t code_addr, uintptr_t
15941594
return 0; // Invalid size
15951595
}
15961596

1597+
if (tlbc_size > MAX_TLBC_SIZE) {
1598+
PyErr_SetString(PyExc_RuntimeError, "TLBC array size exceeds maximum limit");
1599+
return 0; // Invalid size
1600+
}
1601+
15971602
// Allocate and read the entire TLBC array
15981603
size_t array_data_size = tlbc_size * sizeof(void*);
15991604
tlbc_array = PyMem_RawMalloc(sizeof(Py_ssize_t) + array_data_size);
16001605
if (!tlbc_array) {
1606+
PyErr_NoMemory();
16011607
set_exception_cause(unwinder, PyExc_MemoryError, "Failed to allocate TLBC array");
16021608
return 0; // Memory error
16031609
}
@@ -1611,6 +1617,7 @@ cache_tlbc_array(RemoteUnwinderObject *unwinder, uintptr_t code_addr, uintptr_t
16111617
// Create cache entry
16121618
entry = PyMem_RawMalloc(sizeof(TLBCCacheEntry));
16131619
if (!entry) {
1620+
PyErr_NoMemory();
16141621
PyMem_RawFree(tlbc_array);
16151622
set_exception_cause(unwinder, PyExc_MemoryError, "Failed to allocate TLBC cache entry");
16161623
return 0; // Memory error
@@ -1791,6 +1798,7 @@ parse_code_object(RemoteUnwinderObject *unwinder,
17911798

17921799
meta = PyMem_RawMalloc(sizeof(CachedCodeMetadata));
17931800
if (!meta) {
1801+
PyErr_NoMemory();
17941802
set_exception_cause(unwinder, PyExc_MemoryError, "Failed to allocate cached code metadata");
17951803
goto error;
17961804
}

0 commit comments

Comments
 (0)