Rust test targets now create a test launcher allow for setting env vars

rust/private/rust.bzl

@@ -15,7 +15,7 @@
 # buildifier: disable=module-docstring
 load("//rust/private:common.bzl", "rust_common")
 load("//rust/private:rustc.bzl", "rustc_compile_action")
-load("//rust/private:utils.bzl", "determine_output_hash", "find_toolchain")
+load("//rust/private:utils.bzl", "determine_output_hash", "expand_locations", "find_toolchain")
 
 # TODO(marco): Separate each rule into its own file.
 
@@ -219,6 +219,99 @@ def _rust_binary_impl(ctx):
         ),
     )
 
+def _create_test_launcher(ctx, toolchain, output, providers):
+    """Create a process wrapper to ensure runtime environment variables are defined for the test binary
+
+    Args:
+        ctx (ctx): The rule's context object
+        toolchain (rust_toolchain): The current rust toolchain
+        output (File): The output File that will be produced, depends on crate type.
+        providers (list): Providers from a rust compile action. See `rustc_compile_action`
+
+    Returns:
+        list: A list of providers similar to `rustc_compile_action` but with modified default info
+    """
+
+    args = ctx.actions.args()
+
+    # TODO: It's unclear if the toolchain is in the same configuration as the `_launcher` attribute
+    # This should be investigated but for now, we generally assume if the target environment is windows,
+    # the execution environment is windows.
+    if toolchain.os == "windows":
+        launcher = ctx.actions.declare_file(name_to_crate_name(ctx.label.name + ".launcher.exe"))
+        # Because the windows target is a batch file, it expects native windows paths (with backslashes)
+        args.add_all([
+            ctx.executable._launcher.path.replace("/", "\\"),
+            launcher.path.replace("/", "\\"),
+        ])
+    else:
+        launcher = ctx.actions.declare_file(name_to_crate_name(ctx.label.name + ".launcher"))
+        args.add_all([
+            ctx.executable._launcher,
+            launcher,
+        ])
+
+    # Because returned executables must be created from the same rule, the 
+    # launcher target is simply copied and exposed. 
+    ctx.actions.run(
+        outputs = [launcher],
+        tools = [ctx.executable._launcher],
+        mnemonic = "GeneratingLauncher",
+        executable = ctx.executable._launcher_installer,
+        arguments = [args],
+    )
+
+    # Get data attribute
+    data = getattr(ctx.attr, "data", [])
+
+    # Expand the environment variables and write them to a file
+    environ_file = ctx.actions.declare_file(launcher.basename + ".launchfiles/env")
+    environ = expand_locations(
+        ctx,
+        getattr(ctx.attr, "env", {}),
+        data,
+    )
+
+    # Convert the environment variables into a list to be written into a file.
+    environ_list = []
+    for key, value in sorted(environ.items()):
+        environ_list.extend([key, value])
+
+    ctx.actions.write(
+        output = environ_file,
+        content = "\n".join(environ_list)
+    )
+
+    launcher_files = [environ_file]
+
+    # Replace the `DefaultInfo` provider in the returned list
+    default_info = None
+    for i in range(len(providers)):
+        if type(providers[i]) == "DefaultInfo":
+            default_info = providers[i]
+            providers.pop(i)
+            break
+
+    if not default_info:
+        fail("No DefaultInfo provider returned from `rustc_compile_action`")
+
+    providers.extend([
+        DefaultInfo(
+            files = default_info.files,
+            runfiles = default_info.default_runfiles.merge(
+                # The output is now also considered a runfile
+                ctx.runfiles(files = launcher_files + [output]),
+            ),
+            executable = launcher,
+        ),
+        OutputGroupInfo(
+            launcher_files = depset(launcher_files),
+            output = depset([output]),
+        ),
+    ])
+
+    return providers
+
 def _rust_test_common(ctx, toolchain, output):
     """Builds a Rust test binary.
 
@@ -267,15 +360,16 @@ def _rust_test_common(ctx, toolchain, output):
             is_test = True,
         )
 
-    return rustc_compile_action(
+    providers = rustc_compile_action(
         ctx = ctx,
         toolchain = toolchain,
         crate_type = crate_type,
         crate_info = target,
         rust_flags = ["--test"],
-        environ = ctx.attr.env,
     )
 
+    return _create_test_launcher(ctx, toolchain, output, providers)
+
 def _rust_test_impl(ctx):
     """The implementation of the `rust_test` rule
 
@@ -511,6 +605,23 @@ _rust_test_attrs = {
             ["Make variable"](https://docs.bazel.build/versions/master/be/make-variables.html) substitution.
         """),
     ),
+    "_launcher": attr.label(
+        executable = True,
+        default = Label("//util/launcher:launcher"),
+        cfg = "exec",
+        doc = _tidy("""
+            A launcher executable for loading environment and argument files passed in via the `env` attribute
+            and ensuring the variables are set for the underlying test executable.
+        """),
+    ),
+    "_launcher_installer": attr.label(
+        executable = True,
+        default = Label("//util/launcher:installer"),
+        cfg = "exec",
+        doc = _tidy("""
+            A helper script for creating an installer within the test rule.
+        """),
+    ),
 }
 
 rust_library = rule(

rust/private/rustc.bzl

@@ -586,13 +586,6 @@ def rustc_compile_action(
         build_flags_files,
     )
 
-    # Make the user defined enviroment variables available to the action
-    expanded_env = dict()
-    data = getattr(ctx.attr, "data", [])
-    for key in environ:
-        expanded_env[key] = ctx.expand_location(environ[key], data)
-    env.update(expanded_env)
-
     if hasattr(ctx.attr, "version") and ctx.attr.version != "0.0.0":
         formatted_version = " v{}".format(ctx.attr.version)
     else:

test/test_env/tests/run.rs

@@ -4,7 +4,7 @@ fn run() {
     let output = std::process::Command::new(path).output().expect("Failed to run process");
     assert_eq!(&b"Hello world\n"[..], output.stdout.as_slice());
 
-    // Test the `env` attribute of `rust_test`
-    assert_eq!(env!("FERRIS_SAYS"), "Hello fellow Rustaceans!");
-    assert_eq!(env!("HELLO_WORLD_BIN"), "test/test_env/hello-world");
+    // Test the `env` attribute of `rust_test` at run time
+    assert_eq!(std::env::var("FERRIS_SAYS").unwrap(), "Hello fellow Rustaceans!");
+    assert_eq!(std::env::var("HELLO_WORLD_BIN").unwrap(), "test/test_env/hello-world");
 }

util/launcher/BUILD.bazel

@@ -0,0 +1,18 @@
+load("//rust:rust.bzl", "rust_binary")
+load(":installer.bzl", "installer")
+
+package(default_visibility = ["//visibility:public"])
+
+rust_binary(
+    name = "launcher",
+    edition = "2018",
+    srcs = ["launcher_main.rs"]
+)
+
+installer(
+    name = "installer",
+    src = select({
+        "//rust/platform:windows": "launcher_installer.bat",
+        "//conditions:default": "launcher_installer.sh",
+    }),
+)

util/launcher/installer.bzl

@@ -0,0 +1,36 @@
+"""A module defining the installer rule for the rules_rust test launcher"""
+
+def _installer_impl(ctx):
+    """The `installer` rule's implementation
+
+    Args:
+        ctx (ctx): The rule's context object
+
+    Returns:
+        list: A list a DefaultInfo provider
+    """
+
+    installer = ctx.actions.declare_file(ctx.file.src.basename)
+
+    ctx.actions.expand_template(
+        template = ctx.file.src,
+        output = installer,
+        substitutions = {},
+        is_executable = True,
+    )
+
+    return [DefaultInfo(
+        files = depset([installer]),
+        executable = installer,
+    )]
+
+installer = rule(
+    doc = "A rule which makes a native executable script available to other rules",
+    implementation = _installer_impl,
+    attrs = {
+        "src": attr.label(
+            allow_single_file = [".sh", ".bat"],
+        ),
+    },
+    executable = True,
+)

util/launcher/launcher_installer.bat

@@ -0,0 +1,5 @@
+@ECHO OFF
+@REM A native windows script for creating a `run` action output of the
+@REM the rules_rust launcher binary
+
+copy /v /y /b "%1" "%2"

util/launcher/launcher_installer.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+# A simple script for creating a `run` action output of the
+# the rules_rust launcher binary
+
+cp "$1" "$2"

util/launcher/launcher_main.rs

@@ -0,0 +1,102 @@
+use std::collections::BTreeMap;
+use std::fs::File;
+use std::io::{BufReader, BufRead};
+use std::path::PathBuf;
+use std::process::Command;
+use std::vec::Vec;
+
+#[cfg(target_family = "unix")]
+use std::os::unix::process::CommandExt;
+
+/// This string must match the one found in `_create_test_launcher`
+const LAUNCHFILES_ENV_PATH: &'static str = ".launchfiles/env";
+
+/// Load environment variables from a uniquly formatted 
+fn environ() -> BTreeMap<String, String> {
+    let mut environ = BTreeMap::new();
+
+    let mut key: Option<String> = None;
+
+    // Load the environment file into a map
+    let env_path = std::env::args().nth(0).expect("arg 0 was not set") + LAUNCHFILES_ENV_PATH;
+    let file = File::open(env_path).expect("Failed to load the environment file");
+
+    // Find all environment variables by reading pairs of lines as key/value pairs
+    for line in BufReader::new(file).lines() {
+        if key.is_none() {
+            key = Some(line.expect("Failed to read line"));
+            continue;
+        }
+
+        environ.insert(
+            key.expect("Key is not set"), 
+            line.expect("Failed to read line"),
+        );
+
+        key = None;
+    }
+
+    environ
+}
+
+/// Locate the executable based on the name of the launcher executable
+fn executable() -> PathBuf {
+    let mut exec_path = std::env::args().nth(0).expect("arg 0 was not set");
+    let stem_index = exec_path.rfind(".launcher").expect("This executable should always contain `.launcher`");
+
+    // Remove the substring from the exec path
+    for _char in ".launcher".chars() {
+        exec_path.remove(stem_index);
+    }
+
+    PathBuf::from(exec_path)
+}
+
+/// Parse the command line arguments but skip the first element which
+/// is the path to the test runner executable.
+fn args() -> Vec<String> {
+    std::env::args().skip(1).collect()
+}
+
+/// Simply replace the current process with our test
+#[cfg(target_family = "unix")]
+fn exec(environ: BTreeMap<String, String>, executable: PathBuf, args: Vec<String>) {
+    let error = Command::new(&executable)
+    .envs(environ.iter())
+    .args(args)
+    .exec();
+
+    panic!("Process failed to start: {:?} with {:?}", executable, error)
+}
+
+/// On windows, there is no way to replace the current process
+/// so instead we allow the command to run in a subprocess.
+#[cfg(target_family = "windows")]
+fn exec(environ: BTreeMap<String, String>, executable: PathBuf, args: Vec<String>) {
+    let output = Command::new(executable)
+    .envs(environ.iter())
+    .args(args)
+    .output()
+    .expect("Failed to run process");
+
+    std::process::exit(output.status.code().unwrap_or(1));
+}
+
+/// Main entrypoint
+fn main() {
+    // Gather environment variables
+    let environ = environ();
+
+    // Gather arguments
+    let args = args();
+
+    // Find executable
+    let executable = executable();
+
+    // Replace the current process with the test target
+    exec(environ, executable, args);
+
+    // The call to exec should have exited the application.
+    // This code should be unreachable.
+    panic!("Process did not exit");
+}