Add cache key migration

Author: nilskch

crates/rattler_cache/src/package_cache/cache_key.rs

@@ -43,6 +43,11 @@ impl CacheKey {
         self.location_hash = Some(format!("{path_hash:x}"));
         self
     }
+
+    /// Return the old cache key format without the location hash
+    pub fn key_without_location(&self) -> String {
+        format!("{}-{}-{}", &self.name, &self.version, &self.build_string)
+    }
 }
 
 impl CacheKey {
@@ -84,7 +89,7 @@ impl Display for CacheKey {
                 "{}-{}-{}-{}",
                 &self.name, &self.version, &self.build_string, url_hash
             ),
-            None => write!(f, "{}-{}-{}", &self.name, &self.version, &self.build_string),
+            None => write!(f, "{}", self.key_without_location()),
         }
     }
 }

crates/rattler_cache/src/package_cache/mod.rs

@@ -2,6 +2,7 @@
 //! [`PackageCache`].
 
 use std::{
+    collections::VecDeque,
     error::Error,
     fmt::Debug,
     future::Future,
@@ -89,6 +90,10 @@ pub enum PackageCacheError {
     #[error("{0}")]
     LockError(String, #[source] std::io::Error),
 
+    /// A cache key migration error.
+    #[error("failed to migrate the cache key: {0}")]
+    CacheKeyMigrationError(tokio::io::Error),
+
     /// The operation was cancelled
     #[error("operation was cancelled")]
     Cancelled,
@@ -133,7 +138,6 @@ impl PackageCache {
         E: std::error::Error + Send + Sync + 'static,
     {
         let cache_key: CacheKey = pkg.into();
-        let cache_path = self.inner.path.join(cache_key.to_string());
         let cache_entry = self
             .inner
             .packages
@@ -147,7 +151,8 @@ impl PackageCache {
 
         // Validate the cache entry or fetch the package if it is not valid.
         let cache_lock = validate_or_fetch_to_cache(
-            cache_path,
+            &cache_key,
+            &self.inner.path,
             fetch,
             cache_entry.last_revision,
             cache_key.sha256.as_ref(),
@@ -297,10 +302,82 @@ impl PackageCache {
     }
 }
 
+/// Adds the .lock extension to a given path.
+fn get_path_with_lock_extension(path: &Path) -> PathBuf {
+    // Append the `.lock` extension to the cache path to create the lock file path.
+    // `Path::with_extension` strips too much from the filename if it contains one
+    // or more dots.
+    let mut path_str = path.as_os_str().to_owned();
+    path_str.push(".lock");
+    PathBuf::from(path_str)
+}
+
+/// Moves a directory with all nested contents.
+async fn move_dir_all(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> tokio::io::Result<()> {
+    let src = src.as_ref().to_path_buf();
+    let dst = dst.as_ref().to_path_buf();
+    let mut dir_queue = VecDeque::new();
+    dir_queue.push_back((src.clone(), dst.clone()));
+
+    while let Some((src_dir, dst_dir)) = dir_queue.pop_front() {
+        tokio_fs::create_dir_all(&dst_dir).await?;
+        let mut entries = tokio_fs::read_dir(&src_dir).await?;
+
+        while let Some(entry) = entries.next_entry().await? {
+            let src_path = entry.path();
+            let dst_path = dst_dir.join(entry.file_name());
+
+            if entry.file_type().await?.is_dir() {
+                dir_queue.push_back((src_path.clone(), dst_path));
+            } else {
+                tokio_fs::rename(&src_path, &dst_path).await?;
+            }
+        }
+    }
+
+    tokio_fs::remove_dir_all(src).await?;
+
+    Ok(())
+}
+
+/// Migrates old cache keys that do not contain the location hash.
+async fn migrate_cache_entry_to_location_hash(
+    cache_dir: &Path,
+    cache_key: &CacheKey,
+) -> Result<(), PackageCacheError> {
+    let old_cache_entry_path = cache_dir.join(cache_key.key_without_location());
+
+    // Check if the old cache entry exists or not, because it is not necessary to acquire a lock
+    // if the entry does not exist.
+    if old_cache_entry_path.exists() {
+        let new_cache_entry_path = cache_dir.join(cache_key.to_string());
+        let new_cache_lock_path = get_path_with_lock_extension(&new_cache_entry_path);
+        let old_cache_lock_path = get_path_with_lock_extension(&old_cache_entry_path);
+        let _old_lock = CacheRwLock::acquire_write(&old_cache_lock_path).await?;
+        let _new_lock = CacheRwLock::acquire_write(&new_cache_lock_path).await?;
+
+        // We check if the entries exists again, because there could have been a race condition and
+        // another process might have already renamed to old cache entries.
+        if old_cache_entry_path.exists() {
+            move_dir_all(old_cache_entry_path, new_cache_entry_path)
+                .await
+                .map_err(PackageCacheError::CacheKeyMigrationError)?;
+        }
+        if old_cache_lock_path.exists() {
+            tokio_fs::rename(old_cache_lock_path, new_cache_lock_path)
+                .await
+                .map_err(PackageCacheError::CacheKeyMigrationError)?;
+        }
+    }
+
+    Ok(())
+}
+
 /// Validates that the package that is currently stored is a valid package and
 /// otherwise calls the `fetch` method to populate the cache.
 async fn validate_or_fetch_to_cache<F, Fut, E>(
-    path: PathBuf,
+    cache_key: &CacheKey,
+    cache_dir: &Path,
     fetch: F,
     known_valid_revision: Option<u64>,
     given_sha: Option<&Sha256Hash>,
@@ -311,16 +388,10 @@ where
     Fut: Future<Output = Result<(), E>> + 'static,
     E: Error + Send + Sync + 'static,
 {
+    let path = cache_dir.join(cache_key.to_string());
     // Acquire a read lock on the cache entry. This ensures that no other process is
     // currently writing to the cache.
-    let lock_file_path = {
-        // Append the `.lock` extension to the cache path to create the lock file path.
-        // `Path::with_extension` strips too much from the filename if it contains one
-        // or more dots.
-        let mut path_str = path.as_os_str().to_owned();
-        path_str.push(".lock");
-        PathBuf::from(path_str)
-    };
+    let lock_file_path = get_path_with_lock_extension(&path);
 
     // Ensure the directory containing the lock-file exists.
     if let Some(root_dir) = lock_file_path.parent() {
@@ -334,6 +405,9 @@ where
             .await?;
     }
 
+    // Ensure that all cache keys are migrated to contain the location hash.
+    migrate_cache_entry_to_location_hash(cache_dir, cache_key).await?;
+
     // The revision of the cache entry that we already know is valid.
     let mut validated_revision = known_valid_revision;
 
@@ -510,12 +584,13 @@ mod test {
     use reqwest_retry::RetryTransientMiddleware;
     use tempfile::tempdir;
     use tokio::sync::Mutex;
+    use tokio::{fs, io::AsyncWriteExt};
     use tokio_stream::StreamExt;
     use url::Url;
 
     use super::PackageCache;
     use crate::{
-        package_cache::CacheKey,
+        package_cache::{migrate_cache_entry_to_location_hash, CacheKey},
         validation::{validate_package_directory, ValidationMode},
     };
 
@@ -871,4 +946,92 @@ mod test {
             "expected sha256 to be different"
         );
     }
+
+    #[tokio::test]
+    async fn test_migrate_cache_entry_to_location_hash() {
+        let package_url: Url ="https://conda.anaconda.org/robostack/linux-64/ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14.tar.bz2".parse().unwrap();
+        let tar_archive_path = tools::download_and_cache_file_async(
+            package_url.clone(),
+            "4dd9893f1eee45e1579d1a4f5533ef67a84b5e4b7515de7ed0db1dd47adc6bc8",
+        )
+        .await
+        .unwrap();
+
+        // Create a temporary directory to store the packages
+        let packages_dir = tempdir().unwrap().into_path();
+
+        let new_cache_key = "ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14-c5e00a13f8127a70ed3004ecb6f83eafa98ced9e5602bca19e8d52ae3c9fc881";
+        let old_cache_key = "ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14";
+        let new_cache_key_lock = "ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14-c5e00a13f8127a70ed3004ecb6f83eafa98ced9e5602bca19e8d52ae3c9fc881.lock";
+        let old_cache_key_lock = "ros-noetic-rosbridge-suite-0.11.14-py39h6fdeb60_14.lock";
+
+        // Set the sha256 of the package
+        let cache_key: CacheKey = ArchiveIdentifier::try_from_path(&tar_archive_path)
+            .unwrap()
+            .into();
+        let cache_key = cache_key
+            .with_sha256(
+                parse_digest_from_hex::<Sha256>(
+                    "4dd9893f1eee45e1579d1a4f5533ef67a84b5e4b7515de7ed0db1dd47adc6bc8",
+                )
+                .unwrap(),
+            )
+            .with_url(package_url);
+
+        assert_eq!(&cache_key.key_without_location(), old_cache_key);
+        assert_eq!(&cache_key.to_string(), new_cache_key);
+
+        fs::create_dir(packages_dir.join(old_cache_key))
+            .await
+            .unwrap();
+        fs::create_dir(packages_dir.join(old_cache_key).join("foobar"))
+            .await
+            .unwrap();
+        fs::create_dir(packages_dir.join(old_cache_key).join("baz"))
+            .await
+            .unwrap();
+        let mut file = fs::File::create(
+            packages_dir
+                .join(old_cache_key)
+                .join("foobar")
+                .join("some_file"),
+        )
+        .await
+        .unwrap();
+        file.write(b"Some random file content").await.unwrap();
+        let mut file = fs::File::create(packages_dir.join(old_cache_key_lock))
+            .await
+            .unwrap();
+        file.write(b"Fake lockfile content").await.unwrap();
+
+        migrate_cache_entry_to_location_hash(&packages_dir, &cache_key)
+            .await
+            .unwrap();
+
+        assert!(!packages_dir.join(old_cache_key).exists());
+        assert!(!packages_dir.join(old_cache_key_lock).exists());
+        assert!(packages_dir.join(new_cache_key).exists());
+        assert!(packages_dir.join(new_cache_key).is_dir());
+        assert!(packages_dir.join(new_cache_key).join("foobar").exists());
+        assert!(packages_dir.join(new_cache_key).join("baz").exists());
+        assert!(packages_dir.join(new_cache_key_lock).is_file());
+        assert!(packages_dir.join(new_cache_key_lock).exists());
+        assert_eq!(
+            fs::read_to_string(packages_dir.join(new_cache_key_lock))
+                .await
+                .unwrap(),
+            "Fake lockfile content"
+        );
+        assert_eq!(
+            fs::read_to_string(
+                packages_dir
+                    .join(new_cache_key)
+                    .join("foobar")
+                    .join("some_file")
+            )
+            .await
+            .unwrap(),
+            "Some random file content"
+        );
+    }
 }