perf: cache activation environment variables

crates/pixi_consts/src/consts.rs

@@ -18,6 +18,7 @@ pub const ENVIRONMENTS_DIR: &str = "envs";
 pub const SOLVE_GROUP_ENVIRONMENTS_DIR: &str = "solve-group-envs";
 pub const PYPI_DEPENDENCIES: &str = "pypi-dependencies";
 pub const TASK_CACHE_DIR: &str = "task-cache-v0";
+pub const ACTIVATION_ENV_CACHE_DIR: &str = "activation-env-v0";
 pub const PIXI_UV_INSTALLER: &str = "uv-pixi";
 pub const PYPI_CACHE_DIR: &str = "uv-cache";
 pub const CONDA_INSTALLER: &str = "conda";

src/activation.rs

@@ -1,4 +1,6 @@
 use indexmap::IndexMap;
+use rattler_lock::LockFile;
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
 use itertools::Itertools;
@@ -12,8 +14,8 @@ use rattler_shell::{
     shell::ShellEnum,
 };
 
-use crate::project::HasProjectRef;
 use crate::{project::Environment, Project};
+use crate::{project::HasProjectRef, task::EnvironmentHash};
 use pixi_manifest::EnvironmentName;
 use pixi_manifest::FeaturesExt;
 
@@ -30,6 +32,14 @@ pub enum CurrentEnvVarBehavior {
     Exclude,
 }
 
+#[derive(Serialize, Deserialize)]
+struct ActivationCache {
+    /// The hash of the environment which produced the activation's environment variables.
+    hash: EnvironmentHash,
+    /// The environment variables set by the activation.
+    environment_variables: HashMap<String, String>,
+}
+
 impl Project {
     /// Returns environment variables and their values that should be injected when running a command.
     pub(crate) fn get_metadata_env(&self) -> HashMap<String, String> {
@@ -146,7 +156,29 @@ pub(crate) fn get_activator<'p>(
 pub async fn run_activation(
     environment: &Environment<'_>,
     env_var_behavior: &CurrentEnvVarBehavior,
+    lock_file: Option<&LockFile>,
 ) -> miette::Result<HashMap<String, String>> {
+    // If a lock file is provided, we can check whether there exists an environment cache.
+    if let Some(lock_file) = lock_file {
+        let cache_file = environment
+            .project()
+            .activation_env_cache_folder()
+            .join(environment.cache_name());
+        if cache_file.exists() {
+            let cache = tokio::fs::read_to_string(&cache_file)
+                .await
+                .into_diagnostic()?;
+            let cache: ActivationCache = serde_json::from_str(&cache).into_diagnostic()?;
+            let hash = EnvironmentHash::from_environment(environment, lock_file);
+
+            // If the cache's hash matches the environment hash, we can return the cached
+            // environment variables.
+            if cache.hash == hash {
+                return Ok(cache.environment_variables);
+            }
+        }
+    }
+
     let activator = get_activator(environment, ShellEnum::default()).map_err(|e| {
         miette::miette!(format!(
             "failed to create activator for {:?}\n{}",
@@ -210,6 +242,27 @@ pub async fn run_activation(
         }
     };
 
+    // If the lock file is provided and we can compute the environment hash, let's rewrite the
+    // cache file.
+    if let Some(lock_file) = lock_file {
+        let cache_file = environment
+            .project()
+            .activation_env_cache_folder()
+            .join(environment.cache_name());
+        let cache = ActivationCache {
+            hash: EnvironmentHash::from_environment(environment, lock_file),
+            environment_variables: activator_result.clone(),
+        };
+        let cache = serde_json::to_string(&cache).into_diagnostic()?;
+
+        tokio::fs::create_dir_all(environment.project().activation_env_cache_folder())
+            .await
+            .into_diagnostic()?;
+        tokio::fs::write(&cache_file, cache)
+            .await
+            .into_diagnostic()?;
+    }
+
     Ok(activator_result)
 }
 
@@ -289,8 +342,9 @@ pub(crate) fn get_clean_environment_variables() -> HashMap<String, String> {
 pub(crate) async fn initialize_env_variables(
     environment: &Environment<'_>,
     env_var_behavior: CurrentEnvVarBehavior,
+    lock_file: Option<&LockFile>,
 ) -> miette::Result<HashMap<String, String>> {
-    let activation_env = run_activation(environment, &env_var_behavior).await?;
+    let activation_env = run_activation(environment, &env_var_behavior, lock_file).await?;
 
     // Get environment variables from the currently activated shell.
     let current_shell_env_vars = match env_var_behavior {

src/cli/clean.rs

@@ -94,11 +94,11 @@ pub async fn execute(args: Args) -> miette::Result<()> {
                         false,
                     )
                     .await?;
-                    remove_folder_with_progress(project.task_cache_folder(), false).await?;
                 }
                 remove_folder_with_progress(project.environments_dir(), true).await?;
                 remove_folder_with_progress(project.solve_group_environments_dir(), false).await?;
                 remove_folder_with_progress(project.task_cache_folder(), false).await?;
+                remove_folder_with_progress(project.activation_env_cache_folder(), false).await?;
             }
 
             Project::warn_on_discovered_from_env(args.project_config.manifest_path.as_deref())

src/cli/init.rs

@@ -15,7 +15,7 @@ use rattler_conda_types::{NamedChannelOrUrl, Platform};
 use url::Url;
 
 use crate::{
-    environment::{get_up_to_date_prefix, LockFileUsage},
+    environment::{get_up_to_date_lock_file_and_prefix, LockFileUsage},
     Project,
 };
 use pixi_config::{get_default_author, Config};
@@ -243,7 +243,12 @@ pub async fn execute(args: Args) -> miette::Result<()> {
         }
         project.save()?;
 
-        get_up_to_date_prefix(&project.default_environment(), LockFileUsage::Update, false).await?;
+        get_up_to_date_lock_file_and_prefix(
+            &project.default_environment(),
+            LockFileUsage::Update,
+            false,
+        )
+        .await?;
     } else {
         let channels = if let Some(channels) = args.channels {
             channels

src/cli/install.rs

@@ -1,5 +1,5 @@
 use crate::cli::cli_config::ProjectConfig;
-use crate::environment::get_up_to_date_prefix;
+use crate::environment::get_up_to_date_lock_file_and_prefix;
 use crate::Project;
 use clap::Parser;
 use fancy_display::FancyDisplay;
@@ -50,7 +50,8 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     let mut installed_envs = Vec::with_capacity(envs.len());
     for env in envs {
         let environment = project.environment_from_name_or_env_var(Some(env))?;
-        get_up_to_date_prefix(&environment, args.lock_file_usage.into(), false).await?;
+        get_up_to_date_lock_file_and_prefix(&environment, args.lock_file_usage.into(), false)
+            .await?;
         installed_envs.push(environment.name().clone());
     }
 

src/cli/project/channel/add.rs

@@ -1,5 +1,5 @@
 use crate::{
-    environment::{get_up_to_date_prefix, LockFileUsage},
+    environment::{get_up_to_date_lock_file_and_prefix, LockFileUsage},
     Project,
 };
 
@@ -12,7 +12,7 @@ pub async fn execute(mut project: Project, args: AddRemoveArgs) -> miette::Resul
         .add_channels(args.prioritized_channels(), &args.feature_name())?;
 
     // TODO: Update all environments touched by the features defined.
-    get_up_to_date_prefix(
+    get_up_to_date_lock_file_and_prefix(
         &project.default_environment(),
         LockFileUsage::Update,
         args.no_install,

src/cli/project/channel/remove.rs

@@ -1,5 +1,5 @@
 use crate::{
-    environment::{get_up_to_date_prefix, LockFileUsage},
+    environment::{get_up_to_date_lock_file_and_prefix, LockFileUsage},
     Project,
 };
 
@@ -12,7 +12,7 @@ pub async fn execute(mut project: Project, args: AddRemoveArgs) -> miette::Resul
         .remove_channels(args.prioritized_channels(), &args.feature_name())?;
 
     // Try to update the lock-file without the removed channels
-    get_up_to_date_prefix(
+    get_up_to_date_lock_file_and_prefix(
         &project.default_environment(),
         LockFileUsage::Update,
         args.no_install,

src/cli/project/platform/add.rs

@@ -1,7 +1,7 @@
 use std::str::FromStr;
 
 use crate::{
-    environment::{get_up_to_date_prefix, LockFileUsage},
+    environment::{get_up_to_date_lock_file_and_prefix, LockFileUsage},
     Project,
 };
 use clap::Parser;
@@ -44,7 +44,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> {
         .add_platforms(platforms.iter(), &feature_name)?;
 
     // Try to update the lock-file with the new channels
-    get_up_to_date_prefix(
+    get_up_to_date_lock_file_and_prefix(
         &project.default_environment(),
         LockFileUsage::Update,
         args.no_install,

src/cli/project/platform/remove.rs

@@ -1,7 +1,7 @@
 use std::str::FromStr;
 
 use crate::{
-    environment::{get_up_to_date_prefix, LockFileUsage},
+    environment::{get_up_to_date_lock_file_and_prefix, LockFileUsage},
     Project,
 };
 use clap::Parser;
@@ -43,7 +43,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> {
         .manifest
         .remove_platforms(platforms.clone(), &feature_name)?;
 
-    get_up_to_date_prefix(
+    get_up_to_date_lock_file_and_prefix(
         &project.default_environment(),
         LockFileUsage::Update,
         args.no_install,

src/cli/remove.rs

@@ -1,6 +1,6 @@
 use clap::Parser;
 
-use crate::environment::get_up_to_date_prefix;
+use crate::environment::get_up_to_date_lock_file_and_prefix;
 use crate::DependencyType;
 use crate::Project;
 
@@ -65,7 +65,7 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     // TODO: update all environments touched by this feature defined.
     // updating prefix after removing from toml
     if !prefix_update_config.no_lockfile_update {
-        get_up_to_date_prefix(
+        get_up_to_date_lock_file_and_prefix(
             &project.default_environment(),
             prefix_update_config.lock_file_usage(),
             prefix_update_config.no_install,

src/cli/run.rs

@@ -253,9 +253,11 @@ pub async fn get_task_env<'p>(
         CurrentEnvVarBehavior::Include
     };
     let activation_env = await_in_progress("activating environment", |_| {
-        environment
-            .project()
-            .get_activated_environment_variables(environment, env_var_behavior)
+        environment.project().get_activated_environment_variables(
+            environment,
+            env_var_behavior,
+            Some(&lock_file_derived_data.lock_file),
+        )
     })
     .await
     .wrap_err("failed to activate environment")?;

src/cli/shell.rs

@@ -10,7 +10,8 @@ use rattler_shell::{
 
 use crate::cli::cli_config::ProjectConfig;
 use crate::{
-    activation::CurrentEnvVarBehavior, cli::LockFileUsageArgs, environment::get_up_to_date_prefix,
+    activation::CurrentEnvVarBehavior, cli::LockFileUsageArgs,
+    environment::get_up_to_date_lock_file_and_prefix,
     project::virtual_packages::verify_current_platform_has_required_virtual_packages, prompt,
     Project,
 };
@@ -215,11 +216,17 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     };
 
     // Make sure environment is up-to-date, default to install, users can avoid this with frozen or locked.
-    get_up_to_date_prefix(&environment, args.lock_file_usage.into(), false).await?;
+    let (lock_file, _) =
+        get_up_to_date_lock_file_and_prefix(&environment, args.lock_file_usage.into(), false)
+            .await?;
 
     // Get the environment variables we need to set activate the environment in the shell.
     let env = project
-        .get_activated_environment_variables(&environment, CurrentEnvVarBehavior::Exclude)
+        .get_activated_environment_variables(
+            &environment,
+            CurrentEnvVarBehavior::Exclude,
+            Some(&lock_file.lock_file),
+        )
         .await?;
 
     tracing::debug!("Pixi environment activation:\n{:?}", env);

src/cli/shell_hook.rs

@@ -3,6 +3,7 @@ use std::{collections::HashMap, default::Default};
 use clap::Parser;
 use miette::IntoDiagnostic;
 use pixi_config::ConfigCliPrompt;
+use rattler_lock::LockFile;
 use rattler_shell::{
     activation::{ActivationVariables, PathModificationBehavior},
     shell::ShellEnum,
@@ -15,7 +16,7 @@ use crate::project::HasProjectRef;
 use crate::{
     activation::{get_activator, CurrentEnvVarBehavior},
     cli::LockFileUsageArgs,
-    environment::get_up_to_date_prefix,
+    environment::get_up_to_date_lock_file_and_prefix,
     project::Environment,
     Project,
 };
@@ -88,10 +89,17 @@ async fn generate_activation_script(
 
 /// Generates a JSON object describing the changes to the shell environment when
 /// activating the provided pixi environment.
-async fn generate_environment_json(environment: &Environment<'_>) -> miette::Result<String> {
+async fn generate_environment_json(
+    environment: &Environment<'_>,
+    lock_file: &LockFile,
+) -> miette::Result<String> {
     let environment_variables = environment
         .project()
-        .get_activated_environment_variables(environment, CurrentEnvVarBehavior::Exclude)
+        .get_activated_environment_variables(
+            environment,
+            CurrentEnvVarBehavior::Exclude,
+            Some(lock_file),
+        )
         .await?;
 
     let shell_env = ShellEnv {
@@ -107,10 +115,12 @@ pub async fn execute(args: Args) -> miette::Result<()> {
         .with_cli_config(args.config);
     let environment = project.environment_from_name_or_env_var(args.environment)?;
 
-    get_up_to_date_prefix(&environment, args.lock_file_usage.into(), false).await?;
+    let (lock_file, _) =
+        get_up_to_date_lock_file_and_prefix(&environment, args.lock_file_usage.into(), false)
+            .await?;
 
     let output = match args.json {
-        true => generate_environment_json(&environment).await?,
+        true => generate_environment_json(&environment, &lock_file.lock_file).await?,
         false => generate_activation_script(args.shell, &environment).await?,
     };
 
@@ -125,6 +135,8 @@ mod tests {
     use rattler_conda_types::Platform;
     use rattler_shell::shell::{Bash, CmdExe, Fish, NuShell, PowerShell, Shell, Xonsh, Zsh};
 
+    use crate::UpdateLockFileOptions;
+
     use super::*;
 
     #[tokio::test]
@@ -185,7 +197,13 @@ mod tests {
         let path_var_name = default_shell.path_var(&Platform::current());
         let project = Project::discover().unwrap();
         let environment = project.default_environment();
-        let json_env = generate_environment_json(&environment).await.unwrap();
+        let lock_file = project
+            .up_to_date_lock_file(UpdateLockFileOptions::default())
+            .await
+            .unwrap();
+        let json_env = generate_environment_json(&environment, &lock_file.lock_file)
+            .await
+            .unwrap();
         assert!(json_env.contains("\"PIXI_ENVIRONMENT_NAME\":\"default\""));
         assert!(json_env.contains("\"CONDA_PREFIX\":"));
         assert!(json_env.contains(&format!("\"{path_var_name}\":")));