Add locks

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<OsString, OsString>,
+    mutex: Arc<Mutex<()>>,
 }
 
 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<I, K, V>(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<Mutex<()>> {
+        self.mutex.clone()
+    }
+
+    fn lock(&self) -> impl Future<Output = Option<MutexGuard<'_, ()>>> + 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" },

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<Option<Metadata21>, 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(

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<Mutex<()>>,
 }
 
 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<Mutex<()>> {
+        self.lock.clone()
+    }
+
+    /// Grab a file lock for the virtual environment to prevent concurrent writes across processes.
     pub fn lock(&self) -> Result<LockedFile, std::io::Error> {
         if self.interpreter.is_virtualenv() {
             // If the environment a virtualenv, use a virtualenv-specific lock file.

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<Mutex<()>>;
+
+    fn lock(&self) -> impl Future<Output = Option<MutexGuard<'_, ()>>> + 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 {