remove drop-bomb, move empty folder removal to post_process

Author: wolfv

crates/rattler-bin/src/commands/create.rs

@@ -353,7 +353,8 @@ async fn execute_transaction(
     link_pb.enable_steady_tick(Duration::from_millis(100));
 
     // Perform all transactions operations in parallel.
-    stream::iter(transaction.operations)
+    let operations = transaction.operations.clone();
+    stream::iter(operations)
         .map(Ok)
         .try_for_each_concurrent(50, |op| {
             let target_prefix = target_prefix.clone();
@@ -380,11 +381,8 @@ async fn execute_transaction(
         .await?;
 
     // Perform any post processing that is required.
-    let prefix_records = find_installed_packages(&target_prefix, 100)
-        .await
-        .context("failed to determine currently installed packages")?;
     install_driver
-        .post_process(&prefix_records, &target_prefix)
+        .post_process(&transaction, &target_prefix)
         .expect("bla");
 
     Ok(())

crates/rattler/src/install/clobber_registry.rs

@@ -6,31 +6,16 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use drop_bomb::DropBomb;
 use rattler_conda_types::{package::PathsJson, PackageName, PrefixRecord};
 
 /// A registry for clobbering files
 /// The registry keeps track of all files that are installed by a package and
 /// can be used to rename files that are already installed by another package.
-#[derive(Debug)]
+#[derive(Debug, Default, Clone)]
 pub struct ClobberRegistry {
     paths_registry: HashMap<PathBuf, usize>,
     clobbers: HashMap<PathBuf, Vec<usize>>,
     package_names: Vec<PackageName>,
-    drop_bomb: DropBomb,
-}
-
-impl Default for ClobberRegistry {
-    fn default() -> Self {
-        Self {
-            paths_registry: HashMap::new(),
-            clobbers: HashMap::new(),
-            package_names: Vec::new(),
-            drop_bomb: DropBomb::new(
-                "did not call post_process on InstallDriver / ClobberRegistry",
-            ),
-        }
-    }
 }
 
 static CLOBBER_TEMPLATE: &str = "__clobber-from-";
@@ -147,7 +132,6 @@ impl ClobberRegistry {
         sorted_prefix_records: &[&PrefixRecord],
         target_prefix: &Path,
     ) -> Result<(), std::io::Error> {
-        self.drop_bomb.defuse();
         let sorted_names = sorted_prefix_records
             .iter()
             .map(|p| p.repodata_record.package_record.name.clone())
@@ -426,21 +410,20 @@ mod tests {
         install_driver: &InstallDriver,
         install_options: &InstallOptions,
     ) {
-        for op in transaction.operations {
+        for op in &transaction.operations {
             execute_operation(
                 target_prefix,
                 download_client,
                 package_cache,
                 install_driver,
-                op,
+                op.clone(),
                 install_options,
             )
             .await;
         }
 
-        let prefix_records = PrefixRecord::collect_from_prefix(target_prefix).unwrap();
         install_driver
-            .post_process(&prefix_records, target_prefix)
+            .post_process(&transaction, target_prefix)
             .unwrap();
     }
 

crates/rattler/src/install/driver.rs

@@ -1,8 +1,13 @@
 use super::clobber_registry::ClobberRegistry;
-use super::InstallError;
+use super::unlink::{recursively_remove_empty_directories, UnlinkError};
+use super::{InstallError, Transaction};
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
-use rattler_conda_types::{PackageRecord, PrefixRecord};
+use indexmap::IndexSet;
+use itertools::Itertools;
+use rattler_conda_types::prefix_record::PathType;
+use rattler_conda_types::{PackageRecord, PrefixRecord, RepoDataRecord};
+use std::collections::HashSet;
 use std::future::pending;
 use std::path::Path;
 use std::sync::MutexGuard;
@@ -152,18 +157,84 @@ impl InstallDriver {
     /// and will also execute any `post-link.sh/bat` scripts
     pub fn post_process(
         &self,
-        prefix_records: &[PrefixRecord],
+        transaction: &Transaction<PrefixRecord, RepoDataRecord>,
         target_prefix: &Path,
     ) -> Result<(), InstallError> {
+        let prefix_records = PrefixRecord::collect_from_prefix(target_prefix)
+            .map_err(InstallError::PostProcessFailed)?;
+
         let required_packages =
             PackageRecord::sort_topologically(prefix_records.iter().collect::<Vec<_>>());
 
+        self.remove_empty_directories(transaction, &prefix_records, target_prefix)
+            .unwrap_or_else(|e| {
+                tracing::warn!("Failed to remove empty directories: {} (ignored)", e);
+            });
+
         self.clobber_registry()
             .unclobber(&required_packages, target_prefix)
             .map_err(InstallError::PostProcessFailed)?;
 
         Ok(())
     }
+
+    /// Remove all empty directories that are not part of the new prefix records.
+    pub fn remove_empty_directories(
+        &self,
+        transaction: &Transaction<PrefixRecord, RepoDataRecord>,
+        new_prefix_records: &[PrefixRecord],
+        target_prefix: &Path,
+    ) -> Result<(), UnlinkError> {
+        let mut keep_directories = HashSet::new();
+
+        // find all forced directories in the prefix records
+        for record in new_prefix_records {
+            for paths in record.paths_data.paths.iter() {
+                if paths.path_type == PathType::Directory {
+                    let path = target_prefix.join(&paths.relative_path);
+                    keep_directories.insert(path);
+                }
+            }
+        }
+
+        // find all removed directories
+        for record in transaction.removed_packages() {
+            let mut removed_directories = HashSet::new();
+
+            for paths in record.paths_data.paths.iter() {
+                if paths.path_type != PathType::Directory {
+                    if let Some(parent) = paths.relative_path.parent() {
+                        removed_directories.insert(parent);
+                    }
+                }
+            }
+
+            let is_python_noarch = record.repodata_record.package_record.noarch.is_python();
+
+            // Sort the directories by length, so that we delete the deepest directories first.
+            let mut directories: IndexSet<_> = removed_directories.into_iter().sorted().collect();
+
+            while let Some(directory) = directories.pop() {
+                let directory_path = target_prefix.join(directory);
+                let removed_until = recursively_remove_empty_directories(
+                    &directory_path,
+                    target_prefix,
+                    is_python_noarch,
+                    &keep_directories,
+                )?;
+
+                // The directory is not empty which means our parent directory is also not empty,
+                // recursively remove the parent directory from the set as well.
+                while let Some(parent) = removed_until.parent() {
+                    if !directories.shift_remove(parent) {
+                        break;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 impl Drop for InstallDriverInner {

crates/rattler/src/install/mod.rs

@@ -722,12 +722,6 @@ mod test {
             })
             .await;
 
-        // test doesn't write conda-meta, so we ignore the post processing
-        let prefix_records = vec![];
-        install_driver
-            .post_process(&prefix_records, target_dir.path())
-            .unwrap();
-
         // Run the python command and validate the version it outputs
         let python_path = if Platform::current().is_windows() {
             "python.exe"
@@ -771,10 +765,6 @@ mod test {
         .await
         .unwrap();
 
-        install_driver
-            .post_process(&vec![], environment_dir.path())
-            .unwrap();
-
         insta::assert_yaml_snapshot!(paths);
     }
 }

crates/rattler/src/install/transaction.rs

@@ -13,7 +13,7 @@ pub enum TransactionError {
 }
 
 /// Describes an operation to perform
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub enum TransactionOperation<Old, New> {
     /// The given package record should be installed
     Install(New),
@@ -77,6 +77,24 @@ pub struct Transaction<Old, New> {
     pub platform: Platform,
 }
 
+impl<Old, New> Transaction<Old, New> {
+    /// Return an iterator over the prefix records of all packages that are going to be removed.
+    pub fn removed_packages(&self) -> impl Iterator<Item = &Old> + '_ {
+        self.operations
+            .iter()
+            .filter_map(TransactionOperation::record_to_remove)
+    }
+}
+
+impl<Old: AsRef<New>, New> Transaction<Old, New> {
+    /// Return an iterator over the prefix records of all packages that are going to be installed.
+    pub fn installed_packages(&self) -> impl Iterator<Item = &New> + '_ {
+        self.operations
+            .iter()
+            .filter_map(TransactionOperation::record_to_install)
+    }
+}
+
 impl<Old: AsRef<PackageRecord>, New: AsRef<PackageRecord>> Transaction<Old, New> {
     /// Constructs a [`Transaction`] by taking the current situation and diffing that against the
     /// desired situation.