Skip to content

Commit 1b08d62

Browse files
committed
fixup! Merge branch 'main' into tachyon-file-does-not-exist
1 parent 7fbb890 commit 1b08d62

File tree

3 files changed

+91
-65
lines changed

3 files changed

+91
-65
lines changed

Lib/profiling/sampling/_sync_coordinator.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,9 @@ def main() -> NoReturn:
229229
else:
230230
script_path = target_args[0]
231231
script_args = target_args[1:]
232-
if not os.path.isfile(os.path.join(cwd, script_path)):
232+
# Match the path resolution logic in _execute_script
233+
check_path = script_path if os.path.isabs(script_path) else os.path.join(cwd, script_path)
234+
if not os.path.isfile(check_path):
233235
raise TargetError(f"Script not found: {script_path}")
234236

235237
# Signal readiness to profiler

Lib/profiling/sampling/cli.py

Lines changed: 63 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import argparse
44
import os
5-
import select
5+
import selectors
66
import socket
77
import subprocess
88
import sys
@@ -94,21 +94,64 @@ def _parse_mode(mode_string):
9494
return mode_map[mode_string]
9595

9696

97+
def _check_process_died(process):
98+
"""Check if process died and raise an error with stderr if available."""
99+
if process.poll() is None:
100+
return # Process still running
101+
102+
# Process died - try to get stderr for error message
103+
stderr_msg = ""
104+
if process.stderr:
105+
try:
106+
stderr_msg = process.stderr.read().decode().strip()
107+
except (OSError, UnicodeDecodeError):
108+
pass
109+
110+
if stderr_msg:
111+
raise RuntimeError(stderr_msg)
112+
raise RuntimeError(f"Process exited with code {process.returncode}")
113+
114+
115+
def _wait_for_ready_signal(sync_sock, process, timeout):
116+
"""Wait for the ready signal from the subprocess, checking for early death."""
117+
deadline = time.monotonic() + timeout
118+
sel = selectors.DefaultSelector()
119+
sel.register(sync_sock, selectors.EVENT_READ)
120+
121+
try:
122+
while True:
123+
_check_process_died(process)
124+
125+
remaining = deadline - time.monotonic()
126+
if remaining <= 0:
127+
raise socket.timeout("timed out")
128+
129+
if not sel.select(timeout=min(0.1, remaining)):
130+
continue
131+
132+
conn, _ = sync_sock.accept()
133+
try:
134+
ready_signal = conn.recv(_RECV_BUFFER_SIZE)
135+
finally:
136+
conn.close()
137+
138+
if ready_signal != _READY_MESSAGE:
139+
raise RuntimeError(f"Invalid ready signal received: {ready_signal!r}")
140+
return
141+
finally:
142+
sel.close()
143+
144+
97145
def _run_with_sync(original_cmd, suppress_output=False):
98146
"""Run a command with socket-based synchronization and return the process."""
99-
# Create a TCP socket for synchronization with better socket options
100147
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sync_sock:
101-
# Set SO_REUSEADDR to avoid "Address already in use" errors
102148
sync_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
103-
sync_sock.bind(("127.0.0.1", 0)) # Let OS choose a free port
149+
sync_sock.bind(("127.0.0.1", 0))
104150
sync_port = sync_sock.getsockname()[1]
105151
sync_sock.listen(1)
106152
sync_sock.settimeout(_SYNC_TIMEOUT)
107153

108-
# Get current working directory to preserve it
109154
cwd = os.getcwd()
110-
111-
# Build command using the sync coordinator
112155
target_args = original_cmd[1:] # Remove python executable
113156
cmd = (
114157
sys.executable,
@@ -118,50 +161,18 @@ def _run_with_sync(original_cmd, suppress_output=False):
118161
cwd,
119162
) + tuple(target_args)
120163

121-
# Start the process with coordinator
122-
# Suppress stdout if requested (for live mode), but capture stderr
123-
# so we can report errors if the process fails to start
124-
popen_kwargs = {}
164+
popen_kwargs = {"stderr": subprocess.PIPE}
125165
if suppress_output:
126166
popen_kwargs["stdin"] = subprocess.DEVNULL
127167
popen_kwargs["stdout"] = subprocess.DEVNULL
128-
popen_kwargs["stderr"] = subprocess.PIPE
129168

130169
process = subprocess.Popen(cmd, **popen_kwargs)
131170

132171
try:
133-
# Wait for ready signal, but also check if process dies early
134-
deadline = time.monotonic() + _SYNC_TIMEOUT
135-
while True:
136-
# Check if process died
137-
if process.poll() is not None:
138-
stderr_msg = ""
139-
if process.stderr:
140-
try:
141-
stderr_msg = process.stderr.read().decode().strip()
142-
except (OSError, UnicodeDecodeError):
143-
pass
144-
if stderr_msg:
145-
raise RuntimeError(stderr_msg)
146-
raise RuntimeError(
147-
f"Process exited with code {process.returncode}"
148-
)
149-
150-
# Check remaining timeout
151-
remaining = deadline - time.monotonic()
152-
if remaining <= 0:
153-
raise socket.timeout("timed out")
154-
155-
# Wait for socket with short timeout to allow process checks
156-
ready, _, _ = select.select([sync_sock], [], [], min(0.1, remaining))
157-
if ready:
158-
with sync_sock.accept()[0] as conn:
159-
ready_signal = conn.recv(_RECV_BUFFER_SIZE)
160-
if ready_signal != _READY_MESSAGE:
161-
raise RuntimeError(
162-
f"Invalid ready signal received: {ready_signal!r}"
163-
)
164-
break
172+
_wait_for_ready_signal(sync_sock, process, _SYNC_TIMEOUT)
173+
174+
if process.stderr:
175+
process.stderr.close()
165176

166177
except socket.timeout:
167178
# If we timeout, kill the process and raise an error
@@ -670,8 +681,10 @@ def _handle_run(args):
670681
else:
671682
cmd = (sys.executable, args.target, *args.args)
672683

673-
# Run with synchronization
674-
process = _run_with_sync(cmd, suppress_output=False)
684+
try:
685+
process = _run_with_sync(cmd, suppress_output=False)
686+
except RuntimeError as e:
687+
sys.exit(f"Error: {e}")
675688

676689
# Use PROFILING_MODE_ALL for gecko format
677690
mode = (
@@ -718,14 +731,6 @@ def _handle_run(args):
718731

719732
def _handle_live_attach(args, pid):
720733
"""Handle live mode for an existing process."""
721-
# Check if process exists
722-
try:
723-
os.kill(pid, 0)
724-
except ProcessLookupError:
725-
sys.exit(f"Error: process not found: {pid}")
726-
except PermissionError:
727-
pass # Process exists, permission error will be handled later
728-
729734
mode = _parse_mode(args.mode)
730735

731736
# Determine skip_idle based on mode
@@ -766,9 +771,10 @@ def _handle_live_run(args):
766771
else:
767772
cmd = (sys.executable, args.target, *args.args)
768773

769-
# Run with synchronization, suppressing output for live mode
770-
# Note: _run_with_sync will raise if the process dies before signaling ready
771-
process = _run_with_sync(cmd, suppress_output=True)
774+
try:
775+
process = _run_with_sync(cmd, suppress_output=True)
776+
except RuntimeError as e:
777+
sys.exit(f"Error: {e}")
772778

773779
mode = _parse_mode(args.mode)
774780

Lib/test/test_profiling/test_sampling_profiler/test_cli.py

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ def test_cli_module_argument_parsing(self):
7070
mock.patch("profiling.sampling.cli.sample") as mock_sample,
7171
mock.patch("subprocess.Popen") as mock_popen,
7272
mock.patch("socket.socket") as mock_socket,
73-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
73+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
7474
):
7575
from profiling.sampling.cli import main
7676
self._setup_sync_mocks(mock_socket, mock_popen)
@@ -97,7 +97,7 @@ def test_cli_module_with_arguments(self):
9797
mock.patch("profiling.sampling.cli.sample") as mock_sample,
9898
mock.patch("subprocess.Popen") as mock_popen,
9999
mock.patch("socket.socket") as mock_socket,
100-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
100+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
101101
):
102102
self._setup_sync_mocks(mock_socket, mock_popen)
103103
from profiling.sampling.cli import main
@@ -117,7 +117,7 @@ def test_cli_script_argument_parsing(self):
117117
mock.patch("profiling.sampling.cli.sample") as mock_sample,
118118
mock.patch("subprocess.Popen") as mock_popen,
119119
mock.patch("socket.socket") as mock_socket,
120-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
120+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
121121
):
122122
self._setup_sync_mocks(mock_socket, mock_popen)
123123
from profiling.sampling.cli import main
@@ -142,7 +142,7 @@ def test_cli_script_with_arguments(self):
142142
mock.patch("profiling.sampling.cli.sample") as mock_sample,
143143
mock.patch("subprocess.Popen") as mock_popen,
144144
mock.patch("socket.socket") as mock_socket,
145-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
145+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
146146
):
147147
# Use the helper to set up mocks consistently
148148
mock_process = self._setup_sync_mocks(mock_socket, mock_popen)
@@ -252,7 +252,7 @@ def test_cli_module_with_profiler_options(self):
252252
mock.patch("profiling.sampling.cli.sample") as mock_sample,
253253
mock.patch("subprocess.Popen") as mock_popen,
254254
mock.patch("socket.socket") as mock_socket,
255-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
255+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
256256
):
257257
self._setup_sync_mocks(mock_socket, mock_popen)
258258
from profiling.sampling.cli import main
@@ -283,7 +283,7 @@ def test_cli_script_with_profiler_options(self):
283283
mock.patch("profiling.sampling.cli.sample") as mock_sample,
284284
mock.patch("subprocess.Popen") as mock_popen,
285285
mock.patch("socket.socket") as mock_socket,
286-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
286+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
287287
):
288288
self._setup_sync_mocks(mock_socket, mock_popen)
289289
from profiling.sampling.cli import main
@@ -325,7 +325,7 @@ def test_cli_long_module_option(self):
325325
mock.patch("profiling.sampling.cli.sample") as mock_sample,
326326
mock.patch("subprocess.Popen") as mock_popen,
327327
mock.patch("socket.socket") as mock_socket,
328-
mock.patch("profiling.sampling.cli.select.select", return_value=([True], [], [])),
328+
mock.patch("profiling.sampling.cli._wait_for_ready_signal"),
329329
):
330330
self._setup_sync_mocks(mock_socket, mock_popen)
331331
from profiling.sampling.cli import main
@@ -716,3 +716,21 @@ def test_async_aware_incompatible_with_all_threads(self):
716716
error_msg = mock_stderr.getvalue()
717717
self.assertIn("--all-threads", error_msg)
718718
self.assertIn("incompatible with --async-aware", error_msg)
719+
720+
@unittest.skipIf(is_emscripten, "subprocess not available")
721+
def test_run_nonexistent_script_exits_cleanly(self):
722+
"""Test that running a non-existent script exits with a clean error."""
723+
from profiling.sampling.cli import main
724+
with mock.patch("sys.argv", ["profiling.sampling.cli", "run", "/nonexistent/script.py"]):
725+
with self.assertRaises(SystemExit) as cm:
726+
main()
727+
self.assertIn("Script not found", str(cm.exception.code))
728+
729+
@unittest.skipIf(is_emscripten, "subprocess not available")
730+
def test_run_nonexistent_module_exits_cleanly(self):
731+
"""Test that running a non-existent module exits with a clean error."""
732+
from profiling.sampling.cli import main
733+
with mock.patch("sys.argv", ["profiling.sampling.cli", "run", "-m", "nonexistent_module_xyz"]):
734+
with self.assertRaises(SystemExit) as cm:
735+
main()
736+
self.assertIn("Module not found", str(cm.exception.code))

0 commit comments

Comments
 (0)