Skip to content

Commit

Permalink
Let Starlark tests inherit env variables (bazelbuild#15217)
Browse files Browse the repository at this point in the history
Adds an inherited_environment field to testing.TestEnvironment to allow
Starlark rules to implement the equivalent of the native rules'
`env_inherit` attribute.

Work towards bazelbuild#7364. To fully resolve that issue, it remains to handle
executable non-test rules.

RELNOTES: Starlark test rules can use the new inherited_environment
parameter of testing.TestEnvironment to specify environment variables
whose values should be inherited from the shell environment.

Closes bazelbuild#14849.

PiperOrigin-RevId: 439277689

Cherry-pick contains parts of:

Delete non-interning, non-singleton @AutoCodec.

PiperOrigin-RevId: 411683398

Cherry-pick makes the following additional changes:

- Replace use of ImmutableMap.buildKeepingLast with .copyOf

Co-authored-by: Chenchu Kolli <[email protected]>
  • Loading branch information
fmeum and ckolli5 authored May 12, 2022
1 parent 707e8d4 commit 6e54699
Show file tree
Hide file tree
Showing 18 changed files with 269 additions and 72 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -701,7 +701,7 @@ public Dict<String, String> getExecutionInfoDict() {

@Override
public Dict<String, String> getEnv() {
return Dict.immutableCopyOf(env.getFixedEnv().toMap());
return Dict.immutableCopyOf(env.getFixedEnv());
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
import com.google.devtools.build.lib.util.Fingerprint;
import java.util.LinkedHashMap;
import java.util.Map;
Expand All @@ -44,26 +43,36 @@
* action cache), such that Bazel knows exactly which actions it needs to rerun, and does not have
* to reanalyze the entire dependency graph.
*/
@AutoCodec
public final class ActionEnvironment {

/** A map of environment variables. */
/**
* A map of environment variables together with a list of variables to inherit from the shell
* environment.
*/
public interface EnvironmentVariables {

/**
* Returns the environment variables as a map.
* Returns the fixed environment variables as a map.
*
* <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
* CompoundEnvironmentVariables}; use sparingly.
*/
ImmutableMap<String, String> getFixedEnvironment();

/**
* Returns the inherited environment variables as a set.
*
* <p>WARNING: this allocations additional objects if the underlying implementation is a {@link
* <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
* CompoundEnvironmentVariables}; use sparingly.
*/
ImmutableMap<String, String> toMap();
ImmutableSet<String> getInheritedEnvironment();

default boolean isEmpty() {
return toMap().isEmpty();
return getFixedEnvironment().isEmpty() && getInheritedEnvironment().isEmpty();
}

default int size() {
return toMap().size();
return getFixedEnvironment().size() + getInheritedEnvironment().size();
}
}

Expand All @@ -72,11 +81,21 @@ default int size() {
* allocation a new map.
*/
static class CompoundEnvironmentVariables implements EnvironmentVariables {

static EnvironmentVariables create(
Map<String, String> fixedVars, Set<String> inheritedVars, EnvironmentVariables base) {
if (fixedVars.isEmpty() && inheritedVars.isEmpty() && base.isEmpty()) {
return EMPTY_ENVIRONMENT_VARIABLES;
}
return new CompoundEnvironmentVariables(fixedVars, inheritedVars, base);
}

private final EnvironmentVariables current;
private final EnvironmentVariables base;

CompoundEnvironmentVariables(Map<String, String> vars, EnvironmentVariables base) {
this.current = new SimpleEnvironmentVariables(vars);
private CompoundEnvironmentVariables(
Map<String, String> fixedVars, Set<String> inheritedVars, EnvironmentVariables base) {
this.current = SimpleEnvironmentVariables.create(fixedVars, inheritedVars);
this.base = base;
}

Expand All @@ -86,48 +105,62 @@ public boolean isEmpty() {
}

@Override
public ImmutableMap<String, String> toMap() {
Map<String, String> result = new LinkedHashMap<>();
result.putAll(base.toMap());
result.putAll(current.toMap());
public ImmutableMap<String, String> getFixedEnvironment() {
LinkedHashMap<String, String> result = new LinkedHashMap<>();
result.putAll(base.getFixedEnvironment());
result.putAll(current.getFixedEnvironment());
return ImmutableMap.copyOf(result);
}

@Override
public ImmutableSet<String> getInheritedEnvironment() {
ImmutableSet.Builder<String> result = new ImmutableSet.Builder<>();
result.addAll(base.getInheritedEnvironment());
result.addAll(current.getInheritedEnvironment());
return result.build();
}
}

/** A simple {@link EnvironmentVariables}. */
static class SimpleEnvironmentVariables implements EnvironmentVariables {

static EnvironmentVariables create(Map<String, String> vars) {
if (vars.isEmpty()) {
static EnvironmentVariables create(Map<String, String> fixedVars, Set<String> inheritedVars) {
if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
return EMPTY_ENVIRONMENT_VARIABLES;
}
return new SimpleEnvironmentVariables(vars);
return new SimpleEnvironmentVariables(fixedVars, inheritedVars);
}

private final ImmutableMap<String, String> vars;
private final ImmutableMap<String, String> fixedVars;
private final ImmutableSet<String> inheritedVars;

private SimpleEnvironmentVariables(Map<String, String> vars) {
this.vars = ImmutableMap.copyOf(vars);
private SimpleEnvironmentVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
this.fixedVars = ImmutableMap.copyOf(fixedVars);
this.inheritedVars = ImmutableSet.copyOf(inheritedVars);
}

@Override
public ImmutableMap<String, String> getFixedEnvironment() {
return fixedVars;
}

@Override
public ImmutableMap<String, String> toMap() {
return vars;
public ImmutableSet<String> getInheritedEnvironment() {
return inheritedVars;
}
}

/** An empty {@link EnvironmentVariables}. */
public static final EnvironmentVariables EMPTY_ENVIRONMENT_VARIABLES =
new SimpleEnvironmentVariables(ImmutableMap.of());
new SimpleEnvironmentVariables(ImmutableMap.of(), ImmutableSet.of());

/**
* An empty environment, mainly for testing. Production code should never use this, but instead
* get the proper environment from the current configuration.
*/
// TODO(ulfjack): Migrate all production code to use the proper action environment, and then make
// this @VisibleForTesting or rename it to clarify.
public static final ActionEnvironment EMPTY =
new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES, ImmutableSet.of());
public static final ActionEnvironment EMPTY = new ActionEnvironment(EMPTY_ENVIRONMENT_VARIABLES);

/**
* Splits the given map into a map of variables with a fixed value, and a set of variables that
Expand All @@ -147,60 +180,66 @@ public static ActionEnvironment split(Map<String, String> env) {
inheritedEnv.add(key);
}
}
return create(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
return create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
}

private final EnvironmentVariables fixedEnv;
private final ImmutableSet<String> inheritedEnv;
private final EnvironmentVariables vars;

private ActionEnvironment(EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
this.fixedEnv = fixedEnv;
this.inheritedEnv = inheritedEnv;
private ActionEnvironment(EnvironmentVariables vars) {
this.vars = vars;
}

/**
* Creates a new action environment. The order in which the environments are combined is
* undefined, so callers need to take care that the key set of the {@code fixedEnv} map and the
* set of {@code inheritedEnv} elements are disjoint.
*/
@AutoCodec.Instantiator
public static ActionEnvironment create(
EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
if (fixedEnv.isEmpty() && inheritedEnv.isEmpty()) {
private static ActionEnvironment create(EnvironmentVariables vars) {
if (vars.isEmpty()) {
return EMPTY;
}
return new ActionEnvironment(fixedEnv, inheritedEnv);
return new ActionEnvironment(vars);
}

public static ActionEnvironment create(
Map<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
return new ActionEnvironment(SimpleEnvironmentVariables.create(fixedEnv), inheritedEnv);
return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
}

public static ActionEnvironment create(Map<String, String> fixedEnv) {
return new ActionEnvironment(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.of());
return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, ImmutableSet.of()));
}

/**
* Returns a copy of the environment with the given fixed variables added to it, <em>overwriting
* any existing occurrences of those variables</em>.
*/
public ActionEnvironment addFixedVariables(Map<String, String> vars) {
return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), inheritedEnv);
public ActionEnvironment addFixedVariables(Map<String, String> fixedVars) {
return ActionEnvironment.create(
CompoundEnvironmentVariables.create(fixedVars, ImmutableSet.of(), vars));
}

/**
* Returns a copy of the environment with the given fixed and inherited variables added to it,
* <em>overwriting any existing occurrences of those variables</em>.
*/
public ActionEnvironment addVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
return ActionEnvironment.create(
CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars));
}

/** Returns the combined size of the fixed and inherited environments. */
public int size() {
return fixedEnv.size() + inheritedEnv.size();
return vars.size();
}

/**
* Returns the 'fixed' part of the environment, i.e., those environment variables that are set to
* fixed values and their values. This should only be used for testing and to compute the cache
* keys of actions. Use {@link #resolve} instead to get the complete environment.
*/
public EnvironmentVariables getFixedEnv() {
return fixedEnv;
public ImmutableMap<String, String> getFixedEnv() {
return vars.getFixedEnvironment();
}

/**
Expand All @@ -210,7 +249,7 @@ public EnvironmentVariables getFixedEnv() {
* get the complete environment.
*/
public ImmutableSet<String> getInheritedEnv() {
return inheritedEnv;
return vars.getInheritedEnvironment();
}

/**
Expand All @@ -221,8 +260,8 @@ public ImmutableSet<String> getInheritedEnv() {
*/
public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
checkNotNull(clientEnv);
result.putAll(fixedEnv.toMap());
for (String var : inheritedEnv) {
result.putAll(vars.getFixedEnvironment());
for (String var : vars.getInheritedEnvironment()) {
String value = clientEnv.get(var);
if (value != null) {
result.put(var, value);
Expand All @@ -231,7 +270,7 @@ public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
}

public void addTo(Fingerprint f) {
f.addStringMap(fixedEnv.toMap());
f.addStrings(inheritedEnv);
f.addStringMap(vars.getFixedEnvironment());
f.addStrings(vars.getInheritedEnvironment());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -470,6 +470,7 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
providersBuilder.getProvider(TestEnvironmentInfo.PROVIDER.getKey());
if (environmentProvider != null) {
testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
testActionBuilder.addExtraInheritedEnv(environmentProvider.getInheritedEnvironment());
}

TestParams testParams =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -442,7 +442,7 @@ public String describeKey() {
StringBuilder message = new StringBuilder();
message.append(getProgressMessage());
message.append('\n');
for (Map.Entry<String, String> entry : env.getFixedEnv().toMap().entrySet()) {
for (Map.Entry<String, String> entry : env.getFixedEnv().entrySet()) {
message.append(" Environment variable: ");
message.append(ShellEscaper.escapeString(entry.getKey()));
message.append('=');
Expand Down Expand Up @@ -551,7 +551,7 @@ public final ImmutableMap<String, String> getIncompleteEnvironmentForTesting() {
// ActionEnvironment to avoid developers misunderstanding the purpose of this method. That
// requires first updating all subclasses and callers to actually handle environments correctly,
// so it's not a small change.
return env.getFixedEnv().toMap();
return env.getFixedEnv();
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -442,7 +442,7 @@ public boolean isSiblingRepositoryLayout() {
*/
@Override
public ImmutableMap<String, String> getLocalShellEnvironment() {
return actionEnv.getFixedEnv().toMap();
return actionEnv.getFixedEnv();
}

/**
Expand Down Expand Up @@ -579,7 +579,7 @@ public boolean legacyExternalRunfiles() {
*/
@Override
public ImmutableMap<String, String> getTestEnv() {
return testEnv.getFixedEnv().toMap();
return testEnv.getFixedEnv();
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,9 @@
import com.google.devtools.build.lib.vfs.PathFragment;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;
import javax.annotation.Nullable;

/**
Expand Down Expand Up @@ -96,10 +98,12 @@ public NestedSet<PackageGroupContents> getPackageSpecifications() {
private InstrumentedFilesInfo instrumentedFiles;
private int explicitShardCount;
private final Map<String, String> extraEnv;
private final Set<String> extraInheritedEnv;

public TestActionBuilder(RuleContext ruleContext) {
this.ruleContext = ruleContext;
this.extraEnv = new TreeMap<>();
this.extraInheritedEnv = new TreeSet<>();
this.additionalTools = new ImmutableList.Builder<>();
}

Expand Down Expand Up @@ -154,6 +158,11 @@ public TestActionBuilder addExtraEnv(Map<String, String> extraEnv) {
return this;
}

public TestActionBuilder addExtraInheritedEnv(List<String> extraInheritedEnv) {
this.extraInheritedEnv.addAll(extraInheritedEnv);
return this;
}

/**
* Set the explicit shard count. Note that this may be overridden by the sharding strategy.
*/
Expand Down Expand Up @@ -442,7 +451,9 @@ private TestParams createTestAction(int shards)
coverageArtifact,
coverageDirectory,
testProperties,
runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
runfilesSupport
.getActionEnvironment()
.addVariables(extraTestEnv, extraInheritedEnv),
executionSettings,
shard,
run,
Expand Down
Loading

0 comments on commit 6e54699

Please sign in to comment.