From 53239a62db8dcf5d9a376e3e1985bfa97f0313d6 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 6 Mar 2024 20:57:44 -0500 Subject: [PATCH] Add locks --- crates/uv-dispatch/src/lib.rs | 25 +++++++++++++++++-- crates/uv-distribution/src/source/mod.rs | 19 ++++++++++++++ .../uv-interpreter/src/python_environment.rs | 18 +++++++++++-- crates/uv-traits/src/lib.rs | 15 +++++++++++ 4 files changed, 73 insertions(+), 4 deletions(-) diff --git a/crates/uv-dispatch/src/lib.rs b/crates/uv-dispatch/src/lib.rs index 9a7f35e98e90..7d6b32dbefc2 100644 --- a/crates/uv-dispatch/src/lib.rs +++ b/crates/uv-dispatch/src/lib.rs @@ -4,12 +4,14 @@ use std::ffi::OsStr; use std::path::Path; +use std::sync::Arc; use std::{ffi::OsString, future::Future}; use anyhow::{bail, Context, Result}; use futures::FutureExt; use itertools::Itertools; use rustc_hash::FxHashMap; +use tokio::sync::{Mutex, MutexGuard}; use tracing::{debug, instrument}; use distribution_types::{IndexLocations, Name, Resolution, SourceDist}; @@ -42,6 +44,7 @@ pub struct BuildDispatch<'a> { source_build_context: SourceBuildContext, options: Options, build_extra_env_vars: FxHashMap, + mutex: Arc>, } impl<'a> BuildDispatch<'a> { @@ -76,6 +79,7 @@ impl<'a> BuildDispatch<'a> { source_build_context: SourceBuildContext::default(), options: Options::default(), build_extra_env_vars: FxHashMap::default(), + mutex: Arc::new(Mutex::new(())), } } @@ -85,6 +89,12 @@ impl<'a> BuildDispatch<'a> { self } + #[must_use] + pub fn with_build_isolation(mut self, build_isolation: BuildIsolation<'a>) -> Self { + self.build_isolation = build_isolation; + self + } + /// Set the environment variables to be used when building a source distribution. #[must_use] pub fn with_build_extra_env_vars(mut self, sdist_build_env_variables: I) -> Self @@ -104,6 +114,17 @@ impl<'a> BuildDispatch<'a> { impl<'a> BuildContext for BuildDispatch<'a> { type SourceDistBuilder = SourceBuild; + fn mutex(&self) -> Arc> { + self.mutex.clone() + } + + fn lock(&self) -> impl Future>> + Send { + async move { + let guard = self.mutex.lock().await; + Some(guard) + } + } + fn cache(&self) -> &Cache { self.cache } @@ -144,7 +165,7 @@ impl<'a> BuildContext for BuildDispatch<'a> { self.client, self.flat_index, self.index, - self, + self )?; let graph = resolver.resolve().await.with_context(|| { format!( @@ -222,7 +243,7 @@ impl<'a> BuildContext for BuildDispatch<'a> { vec![] } else { // TODO(konstin): Check that there is no endless recursion. - let downloader = Downloader::new(self.cache(), tags, self.client, self); + let downloader = Downloader::new(self.cache, tags, self.client, self); debug!( "Downloading and building requirement{} for build: {}", if remote.len() == 1 { "" } else { "s" }, diff --git a/crates/uv-distribution/src/source/mod.rs b/crates/uv-distribution/src/source/mod.rs index 078283e4e0a5..99d498e19a7b 100644 --- a/crates/uv-distribution/src/source/mod.rs +++ b/crates/uv-distribution/src/source/mod.rs @@ -857,6 +857,11 @@ impl<'a, T: BuildContext> SourceDistCachedBuilder<'a, T> { return Err(Error::NoBuild); } + // Lock the build context. + println!("Locking build context"); + let mutex = self.build_context.mutex(); + let lock = mutex.lock().await; + // Build the wheel. fs::create_dir_all(&cache_shard) .await @@ -876,6 +881,8 @@ impl<'a, T: BuildContext> SourceDistCachedBuilder<'a, T> { .await .map_err(|err| Error::Build(dist.to_string(), err))?; + drop(lock); + // Read the metadata from the wheel. let filename = WheelFilename::from_str(&disk_filename)?; let metadata = read_wheel_metadata(&filename, cache_shard.join(&disk_filename))?; @@ -902,6 +909,16 @@ impl<'a, T: BuildContext> SourceDistCachedBuilder<'a, T> { ) -> Result, Error> { debug!("Preparing metadata for: {dist}"); + // Lock the build context. + // let mutex; + // let _lock; + // if self.build_context.build_isolation().is_shared() { + // mutex = self.build_context.mutex(); + // _lock = mutex.lock().await; + // } + + // let _lock = self.build_context.lock().await; + // Setup the builder. let mut builder = self .build_context @@ -949,6 +966,8 @@ impl<'a, T: BuildContext> SourceDistCachedBuilder<'a, T> { editable_wheel_dir: &Path, ) -> Result<(Dist, String, WheelFilename, Metadata21), Error> { debug!("Building (editable) {editable}"); + + // Build the wheel. let disk_filename = self .build_context .setup_build( diff --git a/crates/uv-interpreter/src/python_environment.rs b/crates/uv-interpreter/src/python_environment.rs index 3987b0fa79c2..7a52a62b8608 100644 --- a/crates/uv-interpreter/src/python_environment.rs +++ b/crates/uv-interpreter/src/python_environment.rs @@ -1,5 +1,6 @@ use std::env; use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; use tracing::debug; @@ -15,6 +16,7 @@ use crate::{find_default_python, find_requested_python, Error, Interpreter}; pub struct PythonEnvironment { root: PathBuf, interpreter: Interpreter, + lock: Arc>, } impl PythonEnvironment { @@ -37,6 +39,7 @@ impl PythonEnvironment { Ok(Self { root: venv, interpreter, + lock: Arc::new(Mutex::default()), }) } @@ -52,6 +55,7 @@ impl PythonEnvironment { Ok(Self { root: interpreter.prefix().to_path_buf(), interpreter, + lock: Arc::new(Mutex::default()), }) } @@ -61,12 +65,17 @@ impl PythonEnvironment { Ok(Self { root: interpreter.prefix().to_path_buf(), interpreter, + lock: Arc::new(Mutex::default()), }) } /// Create a [`PythonEnvironment`] from an existing [`Interpreter`] and root directory. pub fn from_interpreter(interpreter: Interpreter, root: PathBuf) -> Self { - Self { root, interpreter } + Self { + root, + interpreter, + lock: Arc::new(Mutex::default()), + } } /// Returns the location of the Python interpreter. @@ -100,7 +109,12 @@ impl PythonEnvironment { self.interpreter.scripts() } - /// Lock the virtual environment to prevent concurrent writes. + /// Grab an in-memory lock for the virtual environment to prevent concurrent writes across threads. + pub fn acquire(&self) -> Arc> { + self.lock.clone() + } + + /// Grab a file lock for the virtual environment to prevent concurrent writes across processes. pub fn lock(&self) -> Result { if self.interpreter.is_virtualenv() { // If the environment a virtualenv, use a virtualenv-specific lock file. diff --git a/crates/uv-traits/src/lib.rs b/crates/uv-traits/src/lib.rs index d4198ed0af86..7e08e82d6e81 100644 --- a/crates/uv-traits/src/lib.rs +++ b/crates/uv-traits/src/lib.rs @@ -6,8 +6,10 @@ use std::fmt::{Display, Formatter}; use std::future::Future; use std::path::{Path, PathBuf}; use std::str::FromStr; +use std::sync::Arc; use anyhow::Result; +use tokio::sync::{Mutex, MutexGuard}; use distribution_types::{CachedDist, DistributionId, IndexLocations, Resolution, SourceDist}; use once_map::OnceMap; @@ -58,6 +60,12 @@ use uv_normalize::PackageName; pub trait BuildContext: Sync { type SourceDistBuilder: SourceBuildTrait + Send + Sync; + /// Return a mutex that can be used to lock the build context. + fn mutex(&self) -> Arc>; + + fn lock(&self) -> impl Future>> + Send; + + /// Return a reference to the cache. fn cache(&self) -> &Cache; /// All (potentially nested) source distribution builds use the same base python and can reuse @@ -147,6 +155,13 @@ pub enum BuildIsolation<'a> { Shared(&'a PythonEnvironment), } +impl<'a> BuildIsolation<'a> { + /// Returns `true` if build isolation is not enforced. + pub fn is_shared(&self) -> bool { + matches!(self, Self::Shared(_)) + } +} + /// The strategy to use when building source distributions that lack a `pyproject.toml`. #[derive(Copy, Clone, Debug, Default, PartialEq, Eq)] pub enum SetupPyStrategy {