Auto merge of rust-lang#13058 - dpaoliello:extraenv, r=Veykril

Add a new configuration settings to set env vars when running cargo, rustc, etc. commands: cargo.extraEnv and checkOnSave.extraEnv

It can be extremely useful to be able to set environment variables when rust-analyzer is running various cargo or rustc commands (such as `cargo check`, `cargo --print cfg` or `cargo metadata`): users may want to set custom `RUSTFLAGS`, change `PATH` to use a custom toolchain or set a different `CARGO_HOME`.

There is the existing `server.extraEnv` setting that allows env vars to be set when the rust-analyzer server is launched, but using this as the recommended mechanism to also configure cargo/rust has some drawbacks:
- It convolutes configuring the rust-analyzer server with configuring cargo/rustc (one may want to change the `PATH` for cargo/rustc without affecting the rust-analyzer server).
- The name `server.extraEnv` doesn't indicate that cargo/rustc will be affected but renaming it to `cargo.extraEnv` doesn't indicate that the rust-analyzer server would be affected.
- To make the setting useful, it needs to be dynamically reloaded without requiring that the entire extension is reloaded. It might be possible to do this, but it would require the client communicating to the server what the overwritten env vars were at first launch, which isn't easy to do.

This change adds two new configuration settings: `cargo.extraEnv` and `checkOnSave.extraEnv` that can be used to change the environment for the rust-analyzer server after launch (thus affecting any process that rust-analyzer invokes) and the `cargo check` command respectively. `cargo.extraEnv` supports dynamic changes by keeping track of the pre-change values of environment variables, thus it can undo changes made previously before applying the new configuration (and then requesting a workspace reload).

crates/flycheck/Cargo.toml

@@ -13,6 +13,7 @@ doctest = false
 crossbeam-channel = "0.5.5"
 tracing = "0.1.35"
 cargo_metadata = "0.15.0"
+rustc-hash = "1.1.0"
 serde = { version = "1.0.137", features = ["derive"] }
 serde_json = "1.0.81"
 jod-thread = "0.1.2"

crates/flycheck/src/lib.rs

@@ -12,6 +12,7 @@ use std::{
 
 use crossbeam_channel::{never, select, unbounded, Receiver, Sender};
 use paths::AbsPathBuf;
+use rustc_hash::FxHashMap;
 use serde::Deserialize;
 use stdx::{process::streaming_output, JodChild};
 
@@ -30,18 +31,20 @@ pub enum FlycheckConfig {
         all_features: bool,
         features: Vec<String>,
         extra_args: Vec<String>,
+        extra_env: FxHashMap<String, String>,
     },
     CustomCommand {
         command: String,
         args: Vec<String>,
+        extra_env: FxHashMap<String, String>,
     },
 }
 
 impl fmt::Display for FlycheckConfig {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             FlycheckConfig::CargoCommand { command, .. } => write!(f, "cargo {}", command),
-            FlycheckConfig::CustomCommand { command, args } => {
+            FlycheckConfig::CustomCommand { command, args, .. } => {
                 write!(f, "{} {}", command, args.join(" "))
             }
         }
@@ -256,6 +259,7 @@ impl FlycheckActor {
                 all_features,
                 extra_args,
                 features,
+                extra_env,
             } => {
                 let mut cmd = Command::new(toolchain::cargo());
                 cmd.arg(command);
@@ -281,11 +285,13 @@ impl FlycheckActor {
                     }
                 }
                 cmd.args(extra_args);
+                cmd.envs(extra_env);
                 cmd
             }
-            FlycheckConfig::CustomCommand { command, args } => {
+            FlycheckConfig::CustomCommand { command, args, extra_env } => {
                 let mut cmd = Command::new(command);
                 cmd.args(args);
+                cmd.envs(extra_env);
                 cmd
             }
         };

crates/project-model/src/build_scripts.rs

@@ -43,10 +43,12 @@ impl WorkspaceBuildScripts {
         if let Some([program, args @ ..]) = config.run_build_script_command.as_deref() {
             let mut cmd = Command::new(program);
             cmd.args(args);
+            cmd.envs(&config.extra_env);
             return cmd;
         }
 
         let mut cmd = Command::new(toolchain::cargo());
+        cmd.envs(&config.extra_env);
 
         cmd.args(&["check", "--quiet", "--workspace", "--message-format=json"]);
 

crates/project-model/src/cargo_workspace.rs

@@ -2,6 +2,7 @@
 
 use std::iter;
 use std::path::PathBuf;
+use std::str::from_utf8;
 use std::{ops, process::Command};
 
 use anyhow::{Context, Result};
@@ -98,6 +99,8 @@ pub struct CargoConfig {
     pub wrap_rustc_in_build_scripts: bool,
 
     pub run_build_script_command: Option<Vec<String>>,
+
+    pub extra_env: FxHashMap<String, String>,
 }
 
 impl CargoConfig {
@@ -263,8 +266,8 @@ impl CargoWorkspace {
         let target = config
             .target
             .clone()
-            .or_else(|| cargo_config_build_target(cargo_toml))
-            .or_else(|| rustc_discover_host_triple(cargo_toml));
+            .or_else(|| cargo_config_build_target(cargo_toml, config))
+            .or_else(|| rustc_discover_host_triple(cargo_toml, config));
 
         let mut meta = MetadataCommand::new();
         meta.cargo_path(toolchain::cargo());
@@ -292,8 +295,27 @@ impl CargoWorkspace {
         // unclear whether cargo itself supports it.
         progress("metadata".to_string());
 
-        let meta =
-            meta.exec().with_context(|| format!("Failed to run `{:?}`", meta.cargo_command()))?;
+        fn exec_with_env(
+            command: &cargo_metadata::MetadataCommand,
+            extra_env: &FxHashMap<String, String>,
+        ) -> Result<cargo_metadata::Metadata, cargo_metadata::Error> {
+            let mut command = command.cargo_command();
+            command.envs(extra_env);
+            let output = command.output()?;
+            if !output.status.success() {
+                return Err(cargo_metadata::Error::CargoMetadata {
+                    stderr: String::from_utf8(output.stderr)?,
+                });
+            }
+            let stdout = from_utf8(&output.stdout)?
+                .lines()
+                .find(|line| line.starts_with('{'))
+                .ok_or(cargo_metadata::Error::NoJson)?;
+            cargo_metadata::MetadataCommand::parse(stdout)
+        }
+
+        let meta = exec_with_env(&meta, &config.extra_env)
+            .with_context(|| format!("Failed to run `{:?}`", meta.cargo_command()))?;
 
         Ok(meta)
     }
@@ -463,8 +485,9 @@ impl CargoWorkspace {
     }
 }
 
-fn rustc_discover_host_triple(cargo_toml: &ManifestPath) -> Option<String> {
+fn rustc_discover_host_triple(cargo_toml: &ManifestPath, config: &CargoConfig) -> Option<String> {
     let mut rustc = Command::new(toolchain::rustc());
+    rustc.envs(&config.extra_env);
     rustc.current_dir(cargo_toml.parent()).arg("-vV");
     tracing::debug!("Discovering host platform by {:?}", rustc);
     match utf8_stdout(rustc) {
@@ -486,8 +509,9 @@ fn rustc_discover_host_triple(cargo_toml: &ManifestPath) -> Option<String> {
     }
 }
 
-fn cargo_config_build_target(cargo_toml: &ManifestPath) -> Option<String> {
+fn cargo_config_build_target(cargo_toml: &ManifestPath, config: &CargoConfig) -> Option<String> {
     let mut cargo_config = Command::new(toolchain::cargo());
+    cargo_config.envs(&config.extra_env);
     cargo_config
         .current_dir(cargo_toml.parent())
         .args(&["-Z", "unstable-options", "config", "get", "build.target"])

crates/project-model/src/rustc_cfg.rs

@@ -4,9 +4,13 @@ use std::process::Command;
 
 use anyhow::Result;
 
-use crate::{cfg_flag::CfgFlag, utf8_stdout, ManifestPath};
+use crate::{cfg_flag::CfgFlag, utf8_stdout, CargoConfig, ManifestPath};
 
-pub(crate) fn get(cargo_toml: Option<&ManifestPath>, target: Option<&str>) -> Vec<CfgFlag> {
+pub(crate) fn get(
+    cargo_toml: Option<&ManifestPath>,
+    target: Option<&str>,
+    config: &CargoConfig,
+) -> Vec<CfgFlag> {
     let _p = profile::span("rustc_cfg::get");
     let mut res = Vec::with_capacity(6 * 2 + 1);
 
@@ -18,7 +22,7 @@ pub(crate) fn get(cargo_toml: Option<&ManifestPath>, target: Option<&str>) -> Ve
         }
     }
 
-    match get_rust_cfgs(cargo_toml, target) {
+    match get_rust_cfgs(cargo_toml, target, config) {
         Ok(rustc_cfgs) => {
             tracing::debug!(
                 "rustc cfgs found: {:?}",
@@ -35,9 +39,14 @@ pub(crate) fn get(cargo_toml: Option<&ManifestPath>, target: Option<&str>) -> Ve
     res
 }
 
-fn get_rust_cfgs(cargo_toml: Option<&ManifestPath>, target: Option<&str>) -> Result<String> {
+fn get_rust_cfgs(
+    cargo_toml: Option<&ManifestPath>,
+    target: Option<&str>,
+    config: &CargoConfig,
+) -> Result<String> {
     if let Some(cargo_toml) = cargo_toml {
         let mut cargo_config = Command::new(toolchain::cargo());
+        cargo_config.envs(&config.extra_env);
         cargo_config
             .current_dir(cargo_toml.parent())
             .args(&["-Z", "unstable-options", "rustc", "--print", "cfg"])
@@ -52,6 +61,7 @@ fn get_rust_cfgs(cargo_toml: Option<&ManifestPath>, target: Option<&str>) -> Res
     }
     // using unstable cargo features failed, fall back to using plain rustc
     let mut cmd = Command::new(toolchain::rustc());
+    cmd.envs(&config.extra_env);
     cmd.args(&["--print", "cfg", "-O"]);
     if let Some(target) = target {
         cmd.args(&["--target", target]);

crates/project-model/src/sysroot.rs

@@ -10,7 +10,7 @@ use anyhow::{format_err, Result};
 use la_arena::{Arena, Idx};
 use paths::{AbsPath, AbsPathBuf};
 
-use crate::{utf8_stdout, ManifestPath};
+use crate::{utf8_stdout, CargoConfig, ManifestPath};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Sysroot {
@@ -67,18 +67,20 @@ impl Sysroot {
         self.crates.iter().map(|(id, _data)| id)
     }
 
-    pub fn discover(dir: &AbsPath) -> Result<Sysroot> {
+    pub fn discover(dir: &AbsPath, config: &CargoConfig) -> Result<Sysroot> {
         tracing::debug!("Discovering sysroot for {}", dir.display());
-        let sysroot_dir = discover_sysroot_dir(dir)?;
-        let sysroot_src_dir = discover_sysroot_src_dir(&sysroot_dir, dir)?;
+        let sysroot_dir = discover_sysroot_dir(dir, config)?;
+        let sysroot_src_dir = discover_sysroot_src_dir(&sysroot_dir, dir, config)?;
         let res = Sysroot::load(sysroot_dir, sysroot_src_dir)?;
         Ok(res)
     }
 
-    pub fn discover_rustc(cargo_toml: &ManifestPath) -> Option<ManifestPath> {
+    pub fn discover_rustc(cargo_toml: &ManifestPath, config: &CargoConfig) -> Option<ManifestPath> {
         tracing::debug!("Discovering rustc source for {}", cargo_toml.display());
         let current_dir = cargo_toml.parent();
-        discover_sysroot_dir(current_dir).ok().and_then(|sysroot_dir| get_rustc_src(&sysroot_dir))
+        discover_sysroot_dir(current_dir, config)
+            .ok()
+            .and_then(|sysroot_dir| get_rustc_src(&sysroot_dir))
     }
 
     pub fn load(sysroot_dir: AbsPathBuf, sysroot_src_dir: AbsPathBuf) -> Result<Sysroot> {
@@ -144,8 +146,9 @@ impl Sysroot {
     }
 }
 
-fn discover_sysroot_dir(current_dir: &AbsPath) -> Result<AbsPathBuf> {
+fn discover_sysroot_dir(current_dir: &AbsPath, config: &CargoConfig) -> Result<AbsPathBuf> {
     let mut rustc = Command::new(toolchain::rustc());
+    rustc.envs(&config.extra_env);
     rustc.current_dir(current_dir).args(&["--print", "sysroot"]);
     tracing::debug!("Discovering sysroot by {:?}", rustc);
     let stdout = utf8_stdout(rustc)?;
@@ -155,6 +158,7 @@ fn discover_sysroot_dir(current_dir: &AbsPath) -> Result<AbsPathBuf> {
 fn discover_sysroot_src_dir(
     sysroot_path: &AbsPathBuf,
     current_dir: &AbsPath,
+    config: &CargoConfig,
 ) -> Result<AbsPathBuf> {
     if let Ok(path) = env::var("RUST_SRC_PATH") {
         let path = AbsPathBuf::try_from(path.as_str())
@@ -170,6 +174,7 @@ fn discover_sysroot_src_dir(
     get_rust_src(sysroot_path)
         .or_else(|| {
             let mut rustup = Command::new(toolchain::rustup());
+            rustup.envs(&config.extra_env);
             rustup.current_dir(current_dir).args(&["component", "add", "rust-src"]);
             utf8_stdout(rustup).ok()?;
             get_rust_src(sysroot_path)

crates/project-model/src/tests.rs

@@ -10,8 +10,8 @@ use paths::{AbsPath, AbsPathBuf};
 use serde::de::DeserializeOwned;
 
 use crate::{
-    CargoWorkspace, CfgOverrides, ProjectJson, ProjectJsonData, ProjectWorkspace, Sysroot,
-    WorkspaceBuildScripts,
+    CargoConfig, CargoWorkspace, CfgOverrides, ProjectJson, ProjectJsonData, ProjectWorkspace,
+    Sysroot, WorkspaceBuildScripts,
 };
 
 fn load_cargo(file: &str) -> CrateGraph {
@@ -92,13 +92,17 @@ fn rooted_project_json(data: ProjectJsonData) -> ProjectJson {
 }
 
 fn to_crate_graph(project_workspace: ProjectWorkspace) -> CrateGraph {
-    project_workspace.to_crate_graph(&mut |_, _| Ok(Vec::new()), &mut {
-        let mut counter = 0;
-        move |_path| {
-            counter += 1;
-            Some(FileId(counter))
-        }
-    })
+    project_workspace.to_crate_graph(
+        &mut |_, _| Ok(Vec::new()),
+        &mut {
+            let mut counter = 0;
+            move |_path| {
+                counter += 1;
+                Some(FileId(counter))
+            }
+        },
+        &CargoConfig::default(),
+    )
 }
 
 fn check_crate_graph(crate_graph: CrateGraph, expect: Expect) {