perf: enable using sharded repodata for custom channels (#910)

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

@@ -1,29 +1,35 @@
-use crate::global_multi_progress;
+use std::{
+    borrow::Cow,
+    env,
+    future::IntoFuture,
+    path::PathBuf,
+    str::FromStr,
+    sync::Arc,
+    time::{Duration, Instant},
+};
+
 use anyhow::Context;
 use clap::ValueEnum;
 use indicatif::{ProgressBar, ProgressStyle};
 use itertools::Itertools;
-use rattler::install::{IndicatifReporter, Installer};
-use rattler::package_cache::PackageCache;
 use rattler::{
     default_cache_dir,
-    install::{Transaction, TransactionOperation},
+    install::{IndicatifReporter, Installer, Transaction, TransactionOperation},
+    package_cache::PackageCache,
 };
 use rattler_conda_types::{
     Channel, ChannelConfig, GenericVirtualPackage, MatchSpec, ParseStrictness, Platform,
     PrefixRecord, RepoDataRecord, Version,
 };
 use rattler_networking::{AuthenticationMiddleware, AuthenticationStorage};
-use rattler_repodata_gateway::{Gateway, RepoData};
+use rattler_repodata_gateway::{Gateway, RepoData, SourceConfig};
 use rattler_solve::{
     libsolv_c::{self},
     resolvo, SolverImpl, SolverTask,
 };
 use reqwest::Client;
-use std::future::IntoFuture;
-use std::sync::Arc;
-use std::time::Instant;
-use std::{borrow::Cow, env, path::PathBuf, str::FromStr, time::Duration};
+
+use crate::global_multi_progress;
 
 #[derive(Debug, clap::Parser)]
 pub struct Opt {
@@ -105,8 +111,9 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
 
     println!("Installing for platform: {install_platform:?}");
 
-    // Parse the specs from the command line. We do this explicitly instead of allow clap to deal
-    // with this because we need to parse the `channel_config` when parsing matchspecs.
+    // Parse the specs from the command line. We do this explicitly instead of allow
+    // clap to deal with this because we need to parse the `channel_config` when
+    // parsing matchspecs.
     let specs = opt
         .specs
         .iter()
@@ -118,8 +125,9 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
     std::fs::create_dir_all(&cache_dir)
         .map_err(|e| anyhow::anyhow!("could not create cache directory: {}", e))?;
 
-    // Determine the channels to use from the command line or select the default. Like matchspecs
-    // this also requires the use of the `channel_config` so we have to do this manually.
+    // Determine the channels to use from the command line or select the default.
+    // Like matchspecs this also requires the use of the `channel_config` so we
+    // have to do this manually.
     let channels = opt
         .channels
         .unwrap_or_else(|| vec![String::from("conda-forge")])
@@ -130,9 +138,10 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
     // Determine the packages that are currently installed in the environment.
     let installed_packages = PrefixRecord::collect_from_prefix(&target_prefix)?;
 
-    // For each channel/subdirectory combination, download and cache the `repodata.json` that should
-    // be available from the corresponding Url. The code below also displays a nice CLI progress-bar
-    // to give users some more information about what is going on.
+    // For each channel/subdirectory combination, download and cache the
+    // `repodata.json` that should be available from the corresponding Url. The
+    // code below also displays a nice CLI progress-bar to give users some more
+    // information about what is going on.
     let download_client = Client::builder()
         .no_gzip()
         .build()
@@ -147,13 +156,21 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
         .with(rattler_networking::GCSMiddleware)
         .build();
 
-    // Get the package names from the matchspecs so we can only load the package records that we need.
+    // Get the package names from the matchspecs so we can only load the package
+    // records that we need.
     let gateway = Gateway::builder()
         .with_cache_dir(cache_dir.join(rattler_cache::REPODATA_CACHE_DIR))
         .with_package_cache(PackageCache::new(
             cache_dir.join(rattler_cache::PACKAGE_CACHE_DIR),
         ))
         .with_client(download_client.clone())
+        .with_channel_config(rattler_repodata_gateway::ChannelConfig {
+            default: SourceConfig {
+                sharded_enabled: false,
+                ..SourceConfig::default()
+            },
+            ..rattler_repodata_gateway::ChannelConfig::default()
+        })
         .finish();
 
     let start_load_repo_data = Instant::now();
@@ -178,9 +195,9 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
         start_load_repo_data.elapsed()
     );
 
-    // Determine virtual packages of the system. These packages define the capabilities of the
-    // system. Some packages depend on these virtual packages to indicate compatibility with the
-    // hardware of the system.
+    // Determine virtual packages of the system. These packages define the
+    // capabilities of the system. Some packages depend on these virtual
+    // packages to indicate compatibility with the hardware of the system.
     let virtual_packages = wrap_in_progress("determining virtual packages", move || {
         if let Some(virtual_packages) = opt.virtual_package {
             Ok(virtual_packages
@@ -218,9 +235,10 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
             .format_with("\n", |i, f| f(&format_args!("  - {i}",)))
     );
 
-    // Now that we parsed and downloaded all information, construct the packaging problem that we
-    // need to solve. We do this by constructing a `SolverProblem`. This encapsulates all the
-    // information required to be able to solve the problem.
+    // Now that we parsed and downloaded all information, construct the packaging
+    // problem that we need to solve. We do this by constructing a
+    // `SolverProblem`. This encapsulates all the information required to be
+    // able to solve the problem.
     let locked_packages = installed_packages
         .iter()
         .map(|record| record.repodata_record.clone())
@@ -235,8 +253,9 @@ pub async fn create(opt: Opt) -> anyhow::Result<()> {
         ..SolverTask::from_iter(&repo_data)
     };
 
-    // Next, use a solver to solve this specific problem. This provides us with all the operations
-    // we need to apply to our environment to bring it up to date.
+    // Next, use a solver to solve this specific problem. This provides us with all
+    // the operations we need to apply to our environment to bring it up to
+    // date.
     let required_packages =
         wrap_in_progress("solving", move || match opt.solver.unwrap_or_default() {
             Solver::Resolvo => resolvo::Solver.solve(solver_task),
@@ -340,7 +359,8 @@ fn print_transaction(transaction: &Transaction<PrefixRecord, RepoDataRecord>) {
     }
 }
 
-/// Displays a spinner with the given message while running the specified function to completion.
+/// Displays a spinner with the given message while running the specified
+/// function to completion.
 fn wrap_in_progress<T, F: FnOnce() -> T>(msg: impl Into<Cow<'static, str>>, func: F) -> T {
     let pb = ProgressBar::new_spinner();
     pb.enable_steady_tick(Duration::from_millis(100));
@@ -351,7 +371,8 @@ fn wrap_in_progress<T, F: FnOnce() -> T>(msg: impl Into<Cow<'static, str>>, func
     result
 }
 
-/// Displays a spinner with the given message while running the specified function to completion.
+/// Displays a spinner with the given message while running the specified
+/// function to completion.
 async fn wrap_in_async_progress<T, F: IntoFuture<Output = T>>(
     msg: impl Into<Cow<'static, str>>,
     fut: F,
@@ -365,7 +386,8 @@ async fn wrap_in_async_progress<T, F: IntoFuture<Output = T>>(
     result
 }
 
-/// Returns the style to use for a progressbar that is indeterminate and simply shows a spinner.
+/// Returns the style to use for a progressbar that is indeterminate and simply
+/// shows a spinner.
 fn long_running_progress_style() -> indicatif::ProgressStyle {
     ProgressStyle::with_template("{spinner:.green} {msg}").unwrap()
 }

crates/rattler_conda_types/src/channel/channel_url.rs

@@ -0,0 +1,65 @@
+use std::fmt::{Debug, Display, Formatter};
+
+use serde::{Deserialize, Serialize};
+use url::Url;
+
+use crate::{utils::url_with_trailing_slash::UrlWithTrailingSlash, Platform};
+
+/// Represents a channel base url. This is a wrapper around an url that is
+/// normalized:
+///
+/// * The URL always contains a trailing `/`.
+///
+/// This is useful to be able to compare different channels.
+#[derive(Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(transparent)]
+pub struct ChannelUrl(UrlWithTrailingSlash);
+
+impl ChannelUrl {
+    /// Returns the base Url of the channel.
+    pub fn url(&self) -> &Url {
+        &self.0
+    }
+
+    /// Returns the string representation of the url.
+    pub fn as_str(&self) -> &str {
+        self.0.as_str()
+    }
+
+    /// Append the platform to the base url.
+    pub fn platform_url(&self, platform: Platform) -> Url {
+        self.0
+            .join(&format!("{}/", platform.as_str())) // trailing slash is important here as this signifies a directory
+            .expect("platform is a valid url fragment")
+    }
+}
+
+impl Debug for ChannelUrl {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0.as_str())
+    }
+}
+
+impl From<Url> for ChannelUrl {
+    fn from(url: Url) -> Self {
+        Self(UrlWithTrailingSlash::from(url))
+    }
+}
+
+impl From<ChannelUrl> for Url {
+    fn from(value: ChannelUrl) -> Self {
+        value.0.into()
+    }
+}
+
+impl AsRef<Url> for ChannelUrl {
+    fn as_ref(&self) -> &Url {
+        &self.0
+    }
+}
+
+impl Display for ChannelUrl {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", &self.0)
+    }
+}