22
33import argparse
44import os
5- import select
5+ import selectors
66import socket
77import subprocess
88import 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+
97145def _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
719732def _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
0 commit comments