install_wheel_rs/
install_location.rs

1//! Multiplexing between venv install and monotrail install
2
3use fs2::FileExt;
4use fs_err as fs;
5use fs_err::File;
6use std::io;
7use std::ops::Deref;
8use std::path::{Path, PathBuf};
9use tracing::{error, warn};
10
11const INSTALL_LOCKFILE: &str = "install-wheel-rs.lock";
12
13/// I'm not sure that's the right way to normalize here, but it's a single place to change
14/// everything.
15///
16/// For displaying to the user, `-` is better, and it's also what poetry lockfile 2.0 does
17///
18/// Keep in sync with `find_distributions`
19pub fn normalize_name(dep_name: &str) -> String {
20    dep_name.to_lowercase().replace(['.', '_'], "-")
21}
22
23/// A directory for which we acquired a install-wheel-rs.lock lockfile
24pub struct LockedDir {
25    /// The directory to lock
26    path: PathBuf,
27    /// handle on the install-wheel-rs.lock that drops the lock
28    lockfile: File,
29}
30
31impl LockedDir {
32    /// Tries to lock the directory, returns Ok(None) if it is already locked
33    pub fn try_acquire(path: &Path) -> io::Result<Option<Self>> {
34        let lockfile = File::create(path.join(INSTALL_LOCKFILE))?;
35        if lockfile.file().try_lock_exclusive().is_ok() {
36            Ok(Some(Self {
37                path: path.to_path_buf(),
38                lockfile,
39            }))
40        } else {
41            Ok(None)
42        }
43    }
44
45    /// Locks the directory, if necessary blocking until the lock becomes free
46    pub fn acquire(path: &Path) -> io::Result<Self> {
47        let lockfile = File::create(path.join(INSTALL_LOCKFILE))?;
48        lockfile.file().lock_exclusive()?;
49        Ok(Self {
50            path: path.to_path_buf(),
51            lockfile,
52        })
53    }
54}
55
56impl Drop for LockedDir {
57    fn drop(&mut self) {
58        if let Err(err) = self.lockfile.file().unlock() {
59            error!(
60                "Failed to unlock {}: {}",
61                self.lockfile.path().display(),
62                err
63            );
64        }
65    }
66}
67
68impl Deref for LockedDir {
69    type Target = Path;
70
71    fn deref(&self) -> &Self::Target {
72        &self.path
73    }
74}
75
76/// Multiplexing between venv install and monotrail install
77///
78/// For monotrail, we have a structure that is {monotrail}/{normalized(name)}/{version}/tag
79///
80/// We use a lockfile to prevent multiple instance writing stuff on the same time
81/// As of pip 22.0, e.g. `pip install numpy; pip install numpy; pip install numpy` will
82/// nondeterministically fail
83///
84/// I was also thinking about making a shared lock on the import side, but monotrail install
85/// is supposedly atomic (by directory renaming), while for venv installation there can't be
86/// atomicity (we need to add lots of different file without a top level directory / key-turn
87/// file we could rename) and the locking would also need to happen in the import mechanism
88/// itself to ensure
89pub enum InstallLocation<T: Deref<Target = Path>> {
90    Venv {
91        /// absolute path
92        venv_base: T,
93        python_version: (u8, u8),
94    },
95    Monotrail {
96        monotrail_root: T,
97        python: PathBuf,
98        python_version: (u8, u8),
99    },
100}
101
102impl<T: Deref<Target = Path>> InstallLocation<T> {
103    /// Returns the location of the python interpreter
104    pub fn get_python(&self) -> PathBuf {
105        match self {
106            InstallLocation::Venv { venv_base, .. } => {
107                if cfg!(windows) {
108                    venv_base.join("Scripts").join("python.exe")
109                } else {
110                    // canonicalize on python would resolve the symlink
111                    venv_base.join("bin").join("python")
112                }
113            }
114            // TODO: For monotrail use the monotrail launcher
115            InstallLocation::Monotrail { python, .. } => python.clone(),
116        }
117    }
118
119    pub fn get_python_version(&self) -> (u8, u8) {
120        match self {
121            InstallLocation::Venv { python_version, .. } => *python_version,
122            InstallLocation::Monotrail { python_version, .. } => *python_version,
123        }
124    }
125
126    /// TODO: This function is unused?
127    pub fn is_installed(&self, normalized_name: &str, version: &str) -> bool {
128        match self {
129            InstallLocation::Venv {
130                venv_base,
131                python_version,
132            } => {
133                let site_packages = if cfg!(target_os = "windows") {
134                    venv_base.join("Lib").join("site-packages")
135                } else {
136                    venv_base
137                        .join("lib")
138                        .join(format!("python{}.{}", python_version.0, python_version.1))
139                        .join("site-packages")
140                };
141                site_packages
142                    .join(format!("{}-{}.dist-info", normalized_name, version))
143                    .is_dir()
144            }
145            InstallLocation::Monotrail { monotrail_root, .. } => monotrail_root
146                .join(format!("{}-{}", normalized_name, version))
147                .is_dir(),
148        }
149    }
150}
151
152impl InstallLocation<PathBuf> {
153    pub fn acquire_lock(&self) -> io::Result<InstallLocation<LockedDir>> {
154        let root = match self {
155            Self::Venv { venv_base, .. } => venv_base,
156            Self::Monotrail { monotrail_root, .. } => monotrail_root,
157        };
158
159        // If necessary, create monotrail dir
160        fs::create_dir_all(root)?;
161
162        let locked_dir = if let Some(locked_dir) = LockedDir::try_acquire(root)? {
163            locked_dir
164        } else {
165            warn!(
166                "Could not acquire exclusive lock for installing, is another installation process \
167                running? Sleeping until lock becomes free"
168            );
169            LockedDir::acquire(root)?
170        };
171
172        Ok(match self {
173            Self::Venv { python_version, .. } => InstallLocation::Venv {
174                venv_base: locked_dir,
175                python_version: *python_version,
176            },
177            Self::Monotrail {
178                python_version,
179                python,
180                ..
181            } => InstallLocation::Monotrail {
182                monotrail_root: locked_dir,
183                python: python.clone(),
184                python_version: *python_version,
185            },
186        })
187    }
188}
install_wheel_rs/install_location.rs

install_wheel_rs/
install_location.rs
