@@ -2,7 +2,7 @@ use std::panic::AssertUnwindSafe;
22use std:: path:: { Path , PathBuf } ;
33use std:: process:: Stdio ;
44use std:: time:: Duration ;
5- use std:: { io, panic} ;
5+ use std:: { env , io, panic} ;
66
77use async_channel:: { Receiver , SendError } ;
88use tempfile:: tempdir_in;
@@ -20,7 +20,7 @@ use uv_warnings::warn_user;
2020
2121const COMPILEALL_SCRIPT : & str = include_str ! ( "pip_compileall.py" ) ;
2222/// This is longer than any compilation should ever take.
23- const COMPILE_TIMEOUT : Duration = Duration :: from_secs ( 60 ) ;
23+ const DEFAULT_COMPILE_TIMEOUT : Duration = Duration :: from_secs ( 60 ) ;
2424
2525#[ derive( Debug , Error ) ]
2626pub enum CompileError {
@@ -88,6 +88,27 @@ pub async fn compile_tree(
8888 let tempdir = tempdir_in ( cache) . map_err ( CompileError :: TempFile ) ?;
8989 let pip_compileall_py = tempdir. path ( ) . join ( "pip_compileall.py" ) ;
9090
91+ let timeout: Option < Duration > = env:: var ( EnvVars :: UV_COMPILE_BYTECODE_TIMEOUT )
92+ . map ( |value| {
93+ // 0 indicates no timeout
94+ if value == "0" {
95+ None
96+ } else {
97+ Some ( value. parse :: < u64 > ( )
98+ . map ( Duration :: from_secs)
99+ . expect ( "Got invalid value from environment for `UV_COMPILE_BYTECODE_TIMEOUT`. Expected an integer number of seconds, got \" {value}\" ." ) )
100+ }
101+ } )
102+ . unwrap_or ( Some ( DEFAULT_COMPILE_TIMEOUT ) ) ;
103+ if let Some ( duration) = timeout {
104+ debug ! (
105+ "Using bytecode compilation timeout of {}s" ,
106+ duration. as_secs( )
107+ ) ;
108+ } else {
109+ debug ! ( "Disabling bytecode compilation timeout" ) ;
110+ }
111+
91112 debug ! ( "Starting {} bytecode compilation workers" , worker_count) ;
92113 let mut worker_handles = Vec :: new ( ) ;
93114 for _ in 0 ..worker_count {
@@ -98,6 +119,7 @@ pub async fn compile_tree(
98119 python_executable. to_path_buf ( ) ,
99120 pip_compileall_py. clone ( ) ,
100121 receiver. clone ( ) ,
122+ timeout,
101123 ) ;
102124
103125 // Spawn each worker on a dedicated thread.
@@ -189,6 +211,7 @@ async fn worker(
189211 interpreter : PathBuf ,
190212 pip_compileall_py : PathBuf ,
191213 receiver : Receiver < PathBuf > ,
214+ timeout : Option < Duration > ,
192215) -> Result < ( ) , CompileError > {
193216 fs_err:: tokio:: write ( & pip_compileall_py, COMPILEALL_SCRIPT )
194217 . await
@@ -208,12 +231,17 @@ async fn worker(
208231 }
209232 }
210233 } ;
234+
211235 // Handle a broken `python` by using a timeout, one that's higher than any compilation
212236 // should ever take.
213237 let ( mut bytecode_compiler, child_stdin, mut child_stdout, mut child_stderr) =
214- tokio:: time:: timeout ( COMPILE_TIMEOUT , wait_until_ready)
215- . await
216- . map_err ( |_| CompileError :: StartupTimeout ( COMPILE_TIMEOUT ) ) ??;
238+ if let Some ( duration) = timeout {
239+ tokio:: time:: timeout ( duration, wait_until_ready)
240+ . await
241+ . map_err ( |_| CompileError :: StartupTimeout ( timeout. unwrap ( ) ) ) ??
242+ } else {
243+ wait_until_ready. await ?
244+ } ;
217245
218246 let stderr_reader = tokio:: task:: spawn ( async move {
219247 let mut child_stderr_collected: Vec < u8 > = Vec :: new ( ) ;
@@ -223,7 +251,7 @@ async fn worker(
223251 Ok ( child_stderr_collected)
224252 } ) ;
225253
226- let result = worker_main_loop ( receiver, child_stdin, & mut child_stdout) . await ;
254+ let result = worker_main_loop ( receiver, child_stdin, & mut child_stdout, timeout ) . await ;
227255 // Reap the process to avoid zombies.
228256 let _ = bytecode_compiler. kill ( ) . await ;
229257
@@ -340,6 +368,7 @@ async fn worker_main_loop(
340368 receiver : Receiver < PathBuf > ,
341369 mut child_stdin : ChildStdin ,
342370 child_stdout : & mut BufReader < ChildStdout > ,
371+ timeout : Option < Duration > ,
343372) -> Result < ( ) , CompileError > {
344373 let mut out_line = String :: new ( ) ;
345374 while let Ok ( source_file) = receiver. recv ( ) . await {
@@ -372,12 +401,16 @@ async fn worker_main_loop(
372401
373402 // Handle a broken `python` by using a timeout, one that's higher than any compilation
374403 // should ever take.
375- tokio:: time:: timeout ( COMPILE_TIMEOUT , python_handle)
376- . await
377- . map_err ( |_| CompileError :: CompileTimeout {
378- elapsed : COMPILE_TIMEOUT ,
379- source_file : source_file. clone ( ) ,
380- } ) ??;
404+ if let Some ( duration) = timeout {
405+ tokio:: time:: timeout ( duration, python_handle)
406+ . await
407+ . map_err ( |_| CompileError :: CompileTimeout {
408+ elapsed : duration,
409+ source_file : source_file. clone ( ) ,
410+ } ) ??;
411+ } else {
412+ python_handle. await ?;
413+ }
381414
382415 // This is a sanity check, if we don't get the path back something has gone wrong, e.g.
383416 // we're not actually running a python interpreter.
0 commit comments