From 1e8c61a72283c4529395d0054823299f48ec9a43 Mon Sep 17 00:00:00 2001 From: konstin Date: Fri, 1 Mar 2024 20:16:59 +0100 Subject: [PATCH] Review, part 2 --- crates/uv-installer/src/compile.rs | 76 ++++++++++++++++----------- crates/uv/src/commands/pip_install.rs | 2 +- crates/uv/tests/common/mod.rs | 1 + 3 files changed, 47 insertions(+), 32 deletions(-) diff --git a/crates/uv-installer/src/compile.rs b/crates/uv-installer/src/compile.rs index 6a3dd57c10e0c..a4a1ae50d50f2 100644 --- a/crates/uv-installer/src/compile.rs +++ b/crates/uv-installer/src/compile.rs @@ -10,9 +10,11 @@ use thiserror::Error; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::process::{ChildStdin, ChildStdout, Command}; use tokio::task::JoinError; -use uv_warnings::warn_user; +use tracing::instrument; use walkdir::WalkDir; +use uv_warnings::warn_user; + const COMPILEALL_SCRIPT: &str = include_str!("pip_compileall.py"); const MAIN_TIMEOUT: Duration = Duration::from_secs(10); @@ -41,25 +43,37 @@ pub enum CompileError { /// Bytecode compile all file in `dir` using a pool of work-stealing python interpreters running a /// Python script that calls `compileall.compile_file`. /// -/// All compilation errors are muted (like pip). There is a 10s timeout to catch anything gotten -/// stuck; This happens way to easily with channels and subprocesses, e.g. because some pipe is -/// full, we're waiting when it's buffered or we didn't handle channel closing properly. +/// All compilation errors are muted (like pip). There is a 10s timeout to handle the case that +/// the workers have gotten stuck; This happens way to easily with channels and subprocesses, e.g. +/// because some pipe is full, we're waiting when it's buffered, or we didn't handle channel closing +/// properly. /// /// We only compile all files, but we don't update the RECORD, relying on PEP 491: /// > Uninstallers should be smart enough to remove .pyc even if it is not mentioned in RECORD. +#[instrument(skip(python_executable))] pub async fn compile_tree(dir: &Path, python_executable: &Path) -> Result { - let workers = std::thread::available_parallelism().unwrap_or_else(|err| { + debug_assert!( + dir.is_absolute(), + "compileall doesn't work with relative paths" + ); + let worker_count = std::thread::available_parallelism().unwrap_or_else(|err| { warn_user!("Couldn't determine number of cores, compiling with a single thread: {err}"); NonZeroUsize::MIN }); - let (sender, receiver) = async_channel::bounded::(1); + + let (sender, receiver) = async_channel::bounded::(worker_count.get() * 10); + + // Running python with an actual file will produce better error messages. + let tempdir = tempdir().map_err(CompileError::TempFile)?; + let pip_compileall_py = tempdir.path().join("pip_compileall.py"); // Start the workers. let mut worker_handles = Vec::new(); - for _ in 0..workers.get() { + for _ in 0..worker_count.get() { worker_handles.push(tokio::task::spawn(worker( dir.to_path_buf(), python_executable.to_path_buf(), + pip_compileall_py.clone(), receiver.clone(), ))); } @@ -67,13 +81,20 @@ pub async fn compile_tree(dir: &Path, python_executable: &Path) -> Result return Err(CompileError::Timeout(MAIN_TIMEOUT)), // The workers exited. // If e.g. something with the Python interpreter is wrong, the workers have exited @@ -86,7 +107,7 @@ pub async fn compile_tree(dir: &Path, python_executable: &Path) -> Result Result, ) -> Result<(), CompileError> { - // Running python with an actual file will produce better error messages. - let tempdir = tempdir().map_err(CompileError::TempFile)?; - let pip_compileall_py = tempdir.path().join("pip_compileall.py"); fs_err::tokio::write(&pip_compileall_py, COMPILEALL_SCRIPT) .await .map_err(CompileError::TempFile)?; @@ -158,8 +177,7 @@ async fn worker( ); let result = worker_main_loop(receiver, child_stdin, &mut child_stdout).await; - // Avoid `/.venv/bin/python: can't open file` on stderr after success because the tempdir was - // dropped before the Python process even started. + // Reap the process to avoid zombies. let _ = bytecode_compiler.kill().await; result } @@ -174,21 +192,18 @@ async fn worker_main_loop( ) -> Result<(), CompileError> { let mut out_line = String::new(); while let Ok(source_file) = receiver.recv().await { - // Luckily, LF alone works on windows too - let bytes = format!("{}\n", source_file.display()).into_bytes(); - // Ensure we don't get stuck here because some pipe ran full. The Python process isn't doing - // anything at this point except waiting for stdin, so this should be immediate. - let timeout = Duration::from_secs(1); - match tokio::time::timeout(timeout, child_stdin.write_all(&bytes)).await { - // Timeout - Err(_) => { - receiver.close(); - return Err(CompileError::Timeout(timeout)); - } - // Write error - Ok(Err(err)) => return Err(CompileError::ChildStdin(err)), - Ok(Ok(())) => {} + let source_file = source_file.display().to_string(); + if source_file.contains(['\r', '\n']) { + warn_user!("Path contains newline, skipping: {source_file:?}"); + continue; } + // Luckily, LF alone works on windows too + let bytes = format!("{source_file}\n").into_bytes(); + + child_stdin + .write_all(&bytes) + .await + .map_err(CompileError::ChildStdin)?; out_line.clear(); child_stdout @@ -199,9 +214,8 @@ async fn worker_main_loop( // This is a sanity check, if we don't get the path back something has gone wrong, e.g. // we're not actually running a python interpreter. let actual = out_line.trim_end_matches(['\n', '\r']); - let expected = source_file.display().to_string(); - if expected != actual { - return Err(CompileError::WrongPath(expected, actual.to_string())); + if actual != source_file { + return Err(CompileError::WrongPath(source_file, actual.to_string())); } } Ok(()) diff --git a/crates/uv/src/commands/pip_install.rs b/crates/uv/src/commands/pip_install.rs index f39f994508ebb..60504c9a6392c 100644 --- a/crates/uv/src/commands/pip_install.rs +++ b/crates/uv/src/commands/pip_install.rs @@ -683,7 +683,7 @@ async fn install( } if compile { - compile_and_report(&mut printer, &venv).await?; + compile_and_report(&mut printer, venv).await?; } // TODO(konstin): Also check the cache whether any cached or installed dist is already known to diff --git a/crates/uv/tests/common/mod.rs b/crates/uv/tests/common/mod.rs index ca21fd8152495..45c8affefd228 100644 --- a/crates/uv/tests/common/mod.rs +++ b/crates/uv/tests/common/mod.rs @@ -170,6 +170,7 @@ impl TestContext { } /// For when we add pypy to the test suite. + #[allow(clippy::unused_self)] pub fn python_kind(&self) -> &str { "python" }