feat: add run_exports cache

crates/rattler_cache/Cargo.toml

@@ -23,6 +23,7 @@ rattler_digest = { version = "1.0.6", path = "../rattler_digest", default-featur
 rattler_networking = { version = "0.22.4", path = "../rattler_networking", default-features = false }
 rattler_package_streaming = { version = "0.22.28", path = "../rattler_package_streaming", default-features = false, features = ["reqwest"] }
 reqwest.workspace = true
+tempfile.workspace = true
 tokio = { workspace = true, features = ["macros"] }
 tracing.workspace = true
 url.workspace = true
@@ -32,14 +33,14 @@ digest.workspace = true
 fs4 = { workspace = true, features = ["fs-err3-tokio", "tokio"] }
 simple_spawn_blocking = { version = "1.0.0", path = "../simple_spawn_blocking", features = ["tokio"] }
 rayon = { workspace = true }
+serde_json = { workspace = true }
 
 [dev-dependencies]
 assert_matches.workspace = true
 axum.workspace = true
 bytes.workspace = true
 futures.workspace = true
 rstest.workspace = true
-tempfile.workspace = true
 tokio-stream.workspace = true
 tower-http = { workspace = true, features = ["fs"] }
 tools = { path = "../tools" }

crates/rattler_cache/src/consts.rs

@@ -1,4 +1,5 @@
 /// The location in the main cache folder where the conda package cache is stored.
 pub const PACKAGE_CACHE_DIR: &str = "pkgs";
+pub const RUN_EXPORTS_CACHE_DIR: &str = "run_exports";
 /// The location in the main cache folder where the repodata cache is stored.
 pub const REPODATA_CACHE_DIR: &str = "repodata";

crates/rattler_cache/src/lib.rs

@@ -1,11 +1,12 @@
 use std::path::PathBuf;
 
 pub mod package_cache;
+pub mod run_exports_cache;
 
 pub mod validation;
 
 mod consts;
-pub use consts::{PACKAGE_CACHE_DIR, REPODATA_CACHE_DIR};
+pub use consts::{PACKAGE_CACHE_DIR, REPODATA_CACHE_DIR, RUN_EXPORTS_CACHE_DIR};
 
 /// Determines the default cache directory for rattler.
 /// It first checks the environment variable `RATTLER_CACHE_DIR`.

crates/rattler_cache/src/run_exports_cache/cache_key.rs

@@ -0,0 +1,86 @@
+use rattler_conda_types::{package::ArchiveIdentifier, PackageRecord};
+use rattler_digest::{Md5Hash, Sha256Hash};
+use std::fmt::{Display, Formatter};
+
+/// Provides a unique identifier for packages in the cache.
+#[derive(Debug, Hash, Clone, Eq, PartialEq)]
+pub struct CacheKey {
+    pub(crate) name: String,
+    pub(crate) version: String,
+    pub(crate) build_string: String,
+    pub(crate) sha256: Option<Sha256Hash>,
+    pub(crate) md5: Option<Md5Hash>,
+    pub(crate) extension: String,
+}
+
+impl CacheKey {
+    /// Potentially adds a sha256 hash of the archive.
+    pub fn with_opt_sha256(mut self, sha256: Option<Sha256Hash>) -> Self {
+        self.sha256 = sha256;
+        self
+    }
+
+    /// Potentially adds a md5 hash of the archive.
+    pub fn with_opt_md5(mut self, md5: Option<Md5Hash>) -> Self {
+        self.md5 = md5;
+        self
+    }
+}
+
+impl CacheKey {
+    /// Return the sha256 hash of the package if it is known.
+    pub fn sha256(&self) -> Option<Sha256Hash> {
+        self.sha256
+    }
+
+    /// Return the md5 hash of the package if it is known.
+    pub fn md5(&self) -> Option<Md5Hash> {
+        self.md5
+    }
+
+    /// Return the sha256 hash string of the package if it is known.
+    pub fn sha256_str(&self) -> String {
+        self.sha256()
+            .map(|hash| format!("{hash:x}"))
+            .unwrap_or_default()
+    }
+
+    /// Try to create a new cache key from a package record and a filename.
+    pub fn create(record: &PackageRecord, filename: &str) -> Result<Self, CacheKeyError> {
+        let archive_identifier = ArchiveIdentifier::try_from_filename(filename)
+            .ok_or_else(|| CacheKeyError::InvalidArchiveIdentifier(filename.to_string()))?;
+
+        Ok(Self {
+            name: record.name.as_normalized().to_string(),
+            version: record.version.to_string(),
+            build_string: record.build.clone(),
+            sha256: record.sha256,
+            md5: record.md5,
+            extension: archive_identifier.archive_type.extension().to_string(),
+        })
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum CacheKeyError {
+    #[error("could not identify the archive type from the name: {0}")]
+    InvalidArchiveIdentifier(String),
+}
+
+impl Display for CacheKey {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        // we need to use either sha256 or md5 hash to display the key
+        // if both are none, we ignore them
+        let display_key = match (self.sha256(), self.md5()) {
+            (Some(sha256), _) => format!("-{sha256:x}"),
+            (_, Some(md5)) => format!("-{md5:x}"),
+            _ => "".to_string(),
+        };
+
+        write!(
+            f,
+            "{}-{}-{}{}{}",
+            &self.name, &self.version, &self.build_string, display_key, self.extension
+        )
+    }
+}

crates/rattler_cache/src/run_exports_cache/download.rs

@@ -0,0 +1,66 @@
+use std::sync::Arc;
+
+use ::tokio::io::{AsyncSeekExt, AsyncWriteExt};
+use fs_err::tokio;
+use futures::StreamExt;
+use rattler_package_streaming::DownloadReporter;
+use tempfile::NamedTempFile;
+use url::Url;
+
+/// Download the contents of the archive from the specified remote location
+/// and store it in a temporary file.
+pub(crate) async fn download(
+    client: reqwest_middleware::ClientWithMiddleware,
+    url: Url,
+    suffix: &str,
+    reporter: Option<Arc<dyn DownloadReporter>>,
+) -> Result<NamedTempFile, DownloadError> {
+    let temp_file = NamedTempFile::with_suffix(suffix)?;
+
+    // Send the request for the file
+    let response = client.get(url.clone()).send().await?.error_for_status()?;
+
+    if let Some(reporter) = &reporter {
+        reporter.on_download_start();
+    }
+
+    let total_bytes = response.content_length();
+    let (tmp_file_handle, tmp_path) = temp_file.into_parts();
+    // Convert the named temp file into a tokio file
+    let mut file = tokio::File::from_std(fs_err::File::from_parts(tmp_file_handle, &tmp_path));
+
+    let mut stream = response.bytes_stream();
+
+    let mut bytes_received = 0;
+    while let Some(chunk_result) = stream.next().await {
+        let chunk = chunk_result?;
+
+        if let Some(reporter) = &reporter {
+            bytes_received += chunk.len() as u64;
+            reporter.on_download_progress(bytes_received, total_bytes);
+        }
+        file.write_all(&chunk).await?;
+    }
+
+    file.flush().await?;
+
+    file.rewind().await?;
+
+    let file_handle = file.into_parts().0.into_std().await;
+
+    Ok(NamedTempFile::from_parts(file_handle, tmp_path))
+}
+
+/// An error that can occur when downloading an archive.
+#[derive(thiserror::Error, Debug)]
+#[allow(missing_docs)]
+pub enum DownloadError {
+    #[error("an io error occurred: {0}")]
+    Io(#[from] std::io::Error),
+
+    #[error(transparent)]
+    ReqwestMiddleware(#[from] ::reqwest_middleware::Error),
+
+    #[error(transparent)]
+    Reqwest(#[from] ::reqwest::Error),
+}