1+ use std:: num:: ParseIntError ;
12use std:: panic:: AssertUnwindSafe ;
23use std:: path:: { Path , PathBuf } ;
34use std:: process:: Stdio ;
45use std:: time:: Duration ;
5- use std:: { io, panic} ;
6+ use std:: { env , io, panic} ;
67
78use async_channel:: { Receiver , SendError } ;
89use tempfile:: tempdir_in;
@@ -16,11 +17,11 @@ use walkdir::WalkDir;
1617use uv_configuration:: Concurrency ;
1718use uv_fs:: Simplified ;
1819use uv_static:: EnvVars ;
19- use uv_warnings:: warn_user;
20+ use uv_warnings:: { warn_user, warn_user_once } ;
2021
2122const COMPILEALL_SCRIPT : & str = include_str ! ( "pip_compileall.py" ) ;
2223/// This is longer than any compilation should ever take.
23- const COMPILE_TIMEOUT : Duration = Duration :: from_secs ( 60 ) ;
24+ const DEFAULT_COMPILE_TIMEOUT : Duration = Duration :: from_secs ( 60 ) ;
2425
2526#[ derive( Debug , Error ) ]
2627pub enum CompileError {
@@ -88,6 +89,32 @@ pub async fn compile_tree(
8889 let tempdir = tempdir_in ( cache) . map_err ( CompileError :: TempFile ) ?;
8990 let pip_compileall_py = tempdir. path ( ) . join ( "pip_compileall.py" ) ;
9091
92+ let timeout: Option < Duration > = env:: var ( EnvVars :: UV_COMPILE_BYTECODE_TIMEOUT )
93+ . and_then ( |value| {
94+ // 0 indicates no timeout
95+ if value == "0" {
96+ Ok ( None )
97+ } else {
98+ Ok ( Some ( value. parse :: < u64 > ( )
99+ . map ( Duration :: from_secs)
100+ . or_else ( |_| {
101+ // On parse error, warn and use the default timeout
102+ warn_user_once ! ( "Ignoring invalid value from environment for `UV_COMPILE_BYTECODE_TIMEOUT`. Expected an integer number of seconds, got \" {value}\" ." ) ;
103+ Ok :: < Duration , ParseIntError > ( DEFAULT_COMPILE_TIMEOUT )
104+ } )
105+ . unwrap_or ( DEFAULT_COMPILE_TIMEOUT ) ) )
106+ }
107+ } )
108+ . unwrap_or ( Some ( DEFAULT_COMPILE_TIMEOUT ) ) ;
109+ if let Some ( duration) = timeout {
110+ debug ! (
111+ "Using bytecode compilation timeout of {}s" ,
112+ duration. as_secs( )
113+ ) ;
114+ } else {
115+ debug ! ( "Disabling bytecode compilation timeout" ) ;
116+ }
117+
91118 debug ! ( "Starting {} bytecode compilation workers" , worker_count) ;
92119 let mut worker_handles = Vec :: new ( ) ;
93120 for _ in 0 ..worker_count {
@@ -98,6 +125,7 @@ pub async fn compile_tree(
98125 python_executable. to_path_buf ( ) ,
99126 pip_compileall_py. clone ( ) ,
100127 receiver. clone ( ) ,
128+ timeout,
101129 ) ;
102130
103131 // Spawn each worker on a dedicated thread.
@@ -189,6 +217,7 @@ async fn worker(
189217 interpreter : PathBuf ,
190218 pip_compileall_py : PathBuf ,
191219 receiver : Receiver < PathBuf > ,
220+ timeout : Option < Duration > ,
192221) -> Result < ( ) , CompileError > {
193222 fs_err:: tokio:: write ( & pip_compileall_py, COMPILEALL_SCRIPT )
194223 . await
@@ -208,12 +237,17 @@ async fn worker(
208237 }
209238 }
210239 } ;
240+
211241 // Handle a broken `python` by using a timeout, one that's higher than any compilation
212242 // should ever take.
213243 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 ) ) ??;
244+ if let Some ( duration) = timeout {
245+ tokio:: time:: timeout ( duration, wait_until_ready)
246+ . await
247+ . map_err ( |_| CompileError :: StartupTimeout ( timeout. unwrap ( ) ) ) ??
248+ } else {
249+ wait_until_ready. await ?
250+ } ;
217251
218252 let stderr_reader = tokio:: task:: spawn ( async move {
219253 let mut child_stderr_collected: Vec < u8 > = Vec :: new ( ) ;
@@ -223,7 +257,7 @@ async fn worker(
223257 Ok ( child_stderr_collected)
224258 } ) ;
225259
226- let result = worker_main_loop ( receiver, child_stdin, & mut child_stdout) . await ;
260+ let result = worker_main_loop ( receiver, child_stdin, & mut child_stdout, timeout ) . await ;
227261 // Reap the process to avoid zombies.
228262 let _ = bytecode_compiler. kill ( ) . await ;
229263
@@ -340,6 +374,7 @@ async fn worker_main_loop(
340374 receiver : Receiver < PathBuf > ,
341375 mut child_stdin : ChildStdin ,
342376 child_stdout : & mut BufReader < ChildStdout > ,
377+ timeout : Option < Duration > ,
343378) -> Result < ( ) , CompileError > {
344379 let mut out_line = String :: new ( ) ;
345380 while let Ok ( source_file) = receiver. recv ( ) . await {
@@ -372,12 +407,16 @@ async fn worker_main_loop(
372407
373408 // Handle a broken `python` by using a timeout, one that's higher than any compilation
374409 // 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- } ) ??;
410+ if let Some ( duration) = timeout {
411+ tokio:: time:: timeout ( duration, python_handle)
412+ . await
413+ . map_err ( |_| CompileError :: CompileTimeout {
414+ elapsed : duration,
415+ source_file : source_file. clone ( ) ,
416+ } ) ??;
417+ } else {
418+ python_handle. await ?;
419+ }
381420
382421 // This is a sanity check, if we don't get the path back something has gone wrong, e.g.
383422 // we're not actually running a python interpreter.
0 commit comments