Skip to content

Commit

Permalink
fix: self-clobber when updating/downgrading packages (#893)
Browse files Browse the repository at this point in the history
  • Loading branch information
@wolfv
wolfv authored Oct 7, 2024
1 parent ba30fa9 commit 7d692da
Show file tree
Hide file tree
Showing 2 changed files with 157 additions and 37 deletions.
188 changes: 153 additions & 35 deletions crates/rattler/src/install/clobber_registry.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,12 +103,22 @@ impl ClobberRegistry {
/// Register that all the paths of a package are being removed.
pub fn unregister_paths(&mut self, prefix_paths: &PrefixRecord) {
// Find the name in the registry
let name_idx = PackageNameIdx(
self.package_names
.iter()
.position(|n| n == &prefix_paths.repodata_record.package_record.name)
.expect("Package name not found in registry"),
);
let Some(name_idx) = self
.package_names
.iter()
.position(|n| n == &prefix_paths.repodata_record.package_record.name)
.map(PackageNameIdx)
else {
tracing::warn!(
"Tried to unregister paths for a package ({}) that is not in the registry",
prefix_paths
.repodata_record
.package_record
.name
.as_normalized()
);
return;
};

// Remove this package from any clobbering consideration.
for p in &prefix_paths.paths_data.paths {
Expand All @@ -117,7 +127,11 @@ impl ClobberRegistry {
clobber.retain(|&idx| idx != name_idx);
}

let paths_entry = self.paths_registry.get_mut(path).expect("entry must exist");
let Some(paths_entry) = self.paths_registry.get_mut(path) else {
tracing::warn!("The path {} is not in the registry", path.display());
continue;
};

if *paths_entry == Some(name_idx) {
*paths_entry = None;
}
Expand Down Expand Up @@ -147,16 +161,29 @@ impl ClobberRegistry {
};

for (_, path) in computed_paths {
// if we find an entry, we have a clobbering path!
if let Some(&primary_package_idx) = self.paths_registry.get(path) {
let new_path = clobber_name(path, &self.package_names[name_idx.0]);
self.clobbers
.entry(path.clone())
.or_insert_with(|| primary_package_idx.map(|v| vec![v]).unwrap_or_default())
.push(name_idx);

// We insert the non-renamed path here
clobber_paths.insert(path.clone(), new_path);
if let Some(&entry) = self.paths_registry.get(path) {
if let Some(primary_package_idx) = entry {
// if we find an entry, we have a clobbering path!
// Then we rename the current path to a clobbered path
let new_path = clobber_name(path, &self.package_names[name_idx.0]);
self.clobbers
.entry(path.clone())
.or_insert_with(|| vec![primary_package_idx])
.push(name_idx);

// We insert the non-renamed path here
clobber_paths.insert(path.clone(), new_path);
} else {
// In this case, the path we are looking at was previously
// removed so we need to add it back to the registry
self.paths_registry.insert(path.clone(), Some(name_idx));

// If we previously had clobbers with this path, we need to
// add the re-installed package back to the clobbers
if let Some(entry) = self.clobbers.get_mut(path) {
entry.push(name_idx);
}
}
} else {
self.paths_registry.insert(path.clone(), Some(name_idx));
}
Expand Down Expand Up @@ -185,6 +212,7 @@ impl ClobberRegistry {
let mut prefix_records_to_rewrite = HashSet::new();
let mut result = HashMap::new();

tracing::info!("Unclobbering {} files", self.clobbers.len());
for (path, clobbered_by) in self.clobbers.iter() {
let clobbered_by_names = clobbered_by
.iter()
Expand All @@ -199,11 +227,15 @@ impl ClobberRegistry {
.filter(|(_, n)| clobbered_by_names.contains(n))
.collect::<Vec<_>>();

let current_winner = self
.paths_registry
.get(path)
.expect("if a file is clobbered it must also be in the registry")
.map(|idx| &self.package_names[idx.0]);
let Some(current_winner_entry) = self.paths_registry.get(path) else {
tracing::warn!(
"The path {} is clobbered but not in the registry",
path.display()
);
continue;
};

let current_winner = current_winner_entry.map(|idx| &self.package_names[idx.0]);

// Determine which package should write to the file
let winner = match sorted_clobbered_by.last() {
Expand Down Expand Up @@ -264,25 +296,25 @@ impl ClobberRegistry {
},
)?;

let loser_idx = sorted_clobbered_by
if let Some(loser_idx) = sorted_clobbered_by
.iter()
.find(|(_, n)| n == loser_name)
.expect("loser not found")
.0;

rename_path_in_prefix_record(
&mut prefix_records[loser_idx],
path,
&loser_path,
true,
);
prefix_records_to_rewrite.insert(loser_idx);
.map(|(idx, _)| *idx)
{
rename_path_in_prefix_record(
&mut prefix_records[loser_idx],
path,
&loser_path,
true,
);
prefix_records_to_rewrite.insert(loser_idx);
}
}
}

// Rename the winner
let winner_path = clobber_name(path, &winner.1);
tracing::trace!("renaming {} to {}", winner_path.display(), path.display());
tracing::debug!("renaming {} to {}", winner_path.display(), path.display());
fs::rename(target_prefix.join(&winner_path), target_prefix.join(path)).map_err(
|e| {
ClobberError::IoError(
Expand Down Expand Up @@ -835,7 +867,7 @@ mod tests {
.with_prefix_records(&prefix_records)
.finish();

execute_transaction(
let result = execute_transaction(
transaction,
target_prefix.path(),
&reqwest_middleware::ClientWithMiddleware::from(reqwest::Client::new()),
Expand All @@ -845,6 +877,8 @@ mod tests {
)
.await;

println!("== RESULT: {:?}", result.clobbered_paths);

assert_check_files(
target_prefix.path(),
&[
Expand All @@ -867,6 +901,90 @@ mod tests {
);
}

#[tokio::test]
async fn test_self_clobber_update() {
// Create a transaction
let repodata_record_1 = get_repodata_record(
get_test_data_dir().join("clobber/clobber-1-0.1.0-h4616a5c_0.tar.bz2"),
);

let transaction = transaction::Transaction::<PrefixRecord, RepoDataRecord> {
operations: vec![TransactionOperation::Install(repodata_record_1.clone())],
python_info: None,
current_python_info: None,
platform: Platform::current(),
};

// execute transaction
let target_prefix = tempfile::tempdir().unwrap();

let packages_dir = tempfile::tempdir().unwrap();
let cache = PackageCache::new(packages_dir.path());

execute_transaction(
transaction,
target_prefix.path(),
&reqwest_middleware::ClientWithMiddleware::from(reqwest::Client::new()),
&cache,
&InstallDriver::default(),
&InstallOptions::default(),
)
.await;

// check that the files are there
assert_check_files(
target_prefix.path(),
&["clobber.txt", "another-clobber.txt"],
);

let mut prefix_records = PrefixRecord::collect_from_prefix(target_prefix.path()).unwrap();
prefix_records.sort_by(|a, b| {
a.repodata_record
.package_record
.name
.as_normalized()
.cmp(b.repodata_record.package_record.name.as_normalized())
});

// Reinstall the same package
let transaction = transaction::Transaction::<PrefixRecord, RepoDataRecord> {
operations: vec![TransactionOperation::Change {
old: prefix_records[0].clone(),
new: repodata_record_1,
}],
python_info: None,
current_python_info: None,
platform: Platform::current(),
};

let install_driver = InstallDriver::builder()
.with_prefix_records(&prefix_records)
.finish();

install_driver
.pre_process(&transaction, target_prefix.path())
.unwrap();
let dl_client = reqwest_middleware::ClientWithMiddleware::from(reqwest::Client::new());
for op in &transaction.operations {
execute_operation(
target_prefix.path(),
&dl_client,
&cache,
&install_driver,
op.clone(),
&InstallOptions::default(),
)
.await;
}

// Check what files are in the prefix now (note that unclobbering wasn't run yet)
// But also, this is a reinstall so the files should just be overwritten.
assert_check_files(
target_prefix.path(),
&["clobber.txt", "another-clobber.txt"],
);
}

#[tokio::test]
async fn test_clobber_update_and_remove() {
// Create a transaction
Expand Down
6 changes: 4 additions & 2 deletions crates/rattler/src/install/test_utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ use crate::{
package_cache::PackageCache,
};

use super::driver::PostProcessResult;

/// Install a package into the environment and write a `conda-meta` file that
/// contains information about how the file was linked.
pub async fn install_package_to_environment(
Expand Down Expand Up @@ -124,7 +126,7 @@ pub async fn execute_transaction(
package_cache: &PackageCache,
install_driver: &InstallDriver,
install_options: &InstallOptions,
) {
) -> PostProcessResult {
install_driver
.pre_process(&transaction, target_prefix)
.unwrap();
Expand All @@ -143,7 +145,7 @@ pub async fn execute_transaction(

install_driver
.post_process(&transaction, target_prefix)
.unwrap();
.unwrap()
}

pub fn find_prefix_record<'a>(
Expand Down

0 comments on commit 7d692da

Please sign in to comment.