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 } { 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 } { 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 ()
0 commit comments