Skip to content

Commit 6e54699

Browse files
fmeumckolli5
fmeum
and
ckolli5
authored
Let Starlark tests inherit env variables (bazelbuild#15217)
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]>
1 parent 707e8d4 commit 6e54699

File tree

18 files changed

+269
-72
lines changed

18 files changed

+269
-72
lines changed

‎src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -701,7 +701,7 @@ public Dict<String, String> getExecutionInfoDict() {
701701

702702
@Override
703703
public Dict<String, String> getEnv() {
704-
return Dict.immutableCopyOf(env.getFixedEnv().toMap());
704+
return Dict.immutableCopyOf(env.getFixedEnv());
705705
}
706706

707707
@Override

‎src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java‎

Lines changed: 87 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

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

50-
/** A map of environment variables. */
48+
/**
49+
* A map of environment variables together with a list of variables to inherit from the shell
50+
* environment.
51+
*/
5152
public interface EnvironmentVariables {
5253

5354
/**
54-
* Returns the environment variables as a map.
55+
* Returns the fixed environment variables as a map.
56+
*
57+
* <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
58+
* CompoundEnvironmentVariables}; use sparingly.
59+
*/
60+
ImmutableMap<String, String> getFixedEnvironment();
61+
62+
/**
63+
* Returns the inherited environment variables as a set.
5564
*
56-
* <p>WARNING: this allocations additional objects if the underlying implementation is a {@link
65+
* <p>WARNING: this allocates additional objects if the underlying implementation is a {@link
5766
* CompoundEnvironmentVariables}; use sparingly.
5867
*/
59-
ImmutableMap<String, String> toMap();
68+
ImmutableSet<String> getInheritedEnvironment();
6069

6170
default boolean isEmpty() {
62-
return toMap().isEmpty();
71+
return getFixedEnvironment().isEmpty() && getInheritedEnvironment().isEmpty();
6372
}
6473

6574
default int size() {
66-
return toMap().size();
75+
return getFixedEnvironment().size() + getInheritedEnvironment().size();
6776
}
6877
}
6978

@@ -72,11 +81,21 @@ default int size() {
7281
* allocation a new map.
7382
*/
7483
static class CompoundEnvironmentVariables implements EnvironmentVariables {
84+
85+
static EnvironmentVariables create(
86+
Map<String, String> fixedVars, Set<String> inheritedVars, EnvironmentVariables base) {
87+
if (fixedVars.isEmpty() && inheritedVars.isEmpty() && base.isEmpty()) {
88+
return EMPTY_ENVIRONMENT_VARIABLES;
89+
}
90+
return new CompoundEnvironmentVariables(fixedVars, inheritedVars, base);
91+
}
92+
7593
private final EnvironmentVariables current;
7694
private final EnvironmentVariables base;
7795

78-
CompoundEnvironmentVariables(Map<String, String> vars, EnvironmentVariables base) {
79-
this.current = new SimpleEnvironmentVariables(vars);
96+
private CompoundEnvironmentVariables(
97+
Map<String, String> fixedVars, Set<String> inheritedVars, EnvironmentVariables base) {
98+
this.current = SimpleEnvironmentVariables.create(fixedVars, inheritedVars);
8099
this.base = base;
81100
}
82101

@@ -86,48 +105,62 @@ public boolean isEmpty() {
86105
}
87106

88107
@Override
89-
public ImmutableMap<String, String> toMap() {
90-
Map<String, String> result = new LinkedHashMap<>();
91-
result.putAll(base.toMap());
92-
result.putAll(current.toMap());
108+
public ImmutableMap<String, String> getFixedEnvironment() {
109+
LinkedHashMap<String, String> result = new LinkedHashMap<>();
110+
result.putAll(base.getFixedEnvironment());
111+
result.putAll(current.getFixedEnvironment());
93112
return ImmutableMap.copyOf(result);
94113
}
114+
115+
@Override
116+
public ImmutableSet<String> getInheritedEnvironment() {
117+
ImmutableSet.Builder<String> result = new ImmutableSet.Builder<>();
118+
result.addAll(base.getInheritedEnvironment());
119+
result.addAll(current.getInheritedEnvironment());
120+
return result.build();
121+
}
95122
}
96123

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

100-
static EnvironmentVariables create(Map<String, String> vars) {
101-
if (vars.isEmpty()) {
127+
static EnvironmentVariables create(Map<String, String> fixedVars, Set<String> inheritedVars) {
128+
if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
102129
return EMPTY_ENVIRONMENT_VARIABLES;
103130
}
104-
return new SimpleEnvironmentVariables(vars);
131+
return new SimpleEnvironmentVariables(fixedVars, inheritedVars);
105132
}
106133

107-
private final ImmutableMap<String, String> vars;
134+
private final ImmutableMap<String, String> fixedVars;
135+
private final ImmutableSet<String> inheritedVars;
108136

109-
private SimpleEnvironmentVariables(Map<String, String> vars) {
110-
this.vars = ImmutableMap.copyOf(vars);
137+
private SimpleEnvironmentVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
138+
this.fixedVars = ImmutableMap.copyOf(fixedVars);
139+
this.inheritedVars = ImmutableSet.copyOf(inheritedVars);
140+
}
141+
142+
@Override
143+
public ImmutableMap<String, String> getFixedEnvironment() {
144+
return fixedVars;
111145
}
112146

113147
@Override
114-
public ImmutableMap<String, String> toMap() {
115-
return vars;
148+
public ImmutableSet<String> getInheritedEnvironment() {
149+
return inheritedVars;
116150
}
117151
}
118152

119153
/** An empty {@link EnvironmentVariables}. */
120154
public static final EnvironmentVariables EMPTY_ENVIRONMENT_VARIABLES =
121-
new SimpleEnvironmentVariables(ImmutableMap.of());
155+
new SimpleEnvironmentVariables(ImmutableMap.of(), ImmutableSet.of());
122156

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

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

153-
private final EnvironmentVariables fixedEnv;
154-
private final ImmutableSet<String> inheritedEnv;
186+
private final EnvironmentVariables vars;
155187

156-
private ActionEnvironment(EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
157-
this.fixedEnv = fixedEnv;
158-
this.inheritedEnv = inheritedEnv;
188+
private ActionEnvironment(EnvironmentVariables vars) {
189+
this.vars = vars;
159190
}
160191

161192
/**
162193
* Creates a new action environment. The order in which the environments are combined is
163194
* undefined, so callers need to take care that the key set of the {@code fixedEnv} map and the
164195
* set of {@code inheritedEnv} elements are disjoint.
165196
*/
166-
@AutoCodec.Instantiator
167-
public static ActionEnvironment create(
168-
EnvironmentVariables fixedEnv, ImmutableSet<String> inheritedEnv) {
169-
if (fixedEnv.isEmpty() && inheritedEnv.isEmpty()) {
197+
private static ActionEnvironment create(EnvironmentVariables vars) {
198+
if (vars.isEmpty()) {
170199
return EMPTY;
171200
}
172-
return new ActionEnvironment(fixedEnv, inheritedEnv);
201+
return new ActionEnvironment(vars);
173202
}
174203

175204
public static ActionEnvironment create(
176205
Map<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
177-
return new ActionEnvironment(SimpleEnvironmentVariables.create(fixedEnv), inheritedEnv);
206+
return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
178207
}
179208

180209
public static ActionEnvironment create(Map<String, String> fixedEnv) {
181-
return new ActionEnvironment(new SimpleEnvironmentVariables(fixedEnv), ImmutableSet.of());
210+
return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, ImmutableSet.of()));
182211
}
183212

184213
/**
185214
* Returns a copy of the environment with the given fixed variables added to it, <em>overwriting
186215
* any existing occurrences of those variables</em>.
187216
*/
188-
public ActionEnvironment addFixedVariables(Map<String, String> vars) {
189-
return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), inheritedEnv);
217+
public ActionEnvironment addFixedVariables(Map<String, String> fixedVars) {
218+
return ActionEnvironment.create(
219+
CompoundEnvironmentVariables.create(fixedVars, ImmutableSet.of(), vars));
220+
}
221+
222+
/**
223+
* Returns a copy of the environment with the given fixed and inherited variables added to it,
224+
* <em>overwriting any existing occurrences of those variables</em>.
225+
*/
226+
public ActionEnvironment addVariables(Map<String, String> fixedVars, Set<String> inheritedVars) {
227+
return ActionEnvironment.create(
228+
CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars));
190229
}
191230

192231
/** Returns the combined size of the fixed and inherited environments. */
193232
public int size() {
194-
return fixedEnv.size() + inheritedEnv.size();
233+
return vars.size();
195234
}
196235

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

206245
/**
@@ -210,7 +249,7 @@ public EnvironmentVariables getFixedEnv() {
210249
* get the complete environment.
211250
*/
212251
public ImmutableSet<String> getInheritedEnv() {
213-
return inheritedEnv;
252+
return vars.getInheritedEnvironment();
214253
}
215254

216255
/**
@@ -221,8 +260,8 @@ public ImmutableSet<String> getInheritedEnv() {
221260
*/
222261
public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
223262
checkNotNull(clientEnv);
224-
result.putAll(fixedEnv.toMap());
225-
for (String var : inheritedEnv) {
263+
result.putAll(vars.getFixedEnvironment());
264+
for (String var : vars.getInheritedEnvironment()) {
226265
String value = clientEnv.get(var);
227266
if (value != null) {
228267
result.put(var, value);
@@ -231,7 +270,7 @@ public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
231270
}
232271

233272
public void addTo(Fingerprint f) {
234-
f.addStringMap(fixedEnv.toMap());
235-
f.addStrings(inheritedEnv);
273+
f.addStringMap(vars.getFixedEnvironment());
274+
f.addStrings(vars.getInheritedEnvironment());
236275
}
237276
}

‎src/main/java/com/google/devtools/build/lib/analysis/RuleConfiguredTargetBuilder.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,7 @@ private TestProvider initializeTestProvider(FilesToRunProvider filesToRunProvide
470470
providersBuilder.getProvider(TestEnvironmentInfo.PROVIDER.getKey());
471471
if (environmentProvider != null) {
472472
testActionBuilder.addExtraEnv(environmentProvider.getEnvironment());
473+
testActionBuilder.addExtraInheritedEnv(environmentProvider.getInheritedEnvironment());
473474
}
474475

475476
TestParams testParams =

‎src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,7 @@ public String describeKey() {
442442
StringBuilder message = new StringBuilder();
443443
message.append(getProgressMessage());
444444
message.append('\n');
445-
for (Map.Entry<String, String> entry : env.getFixedEnv().toMap().entrySet()) {
445+
for (Map.Entry<String, String> entry : env.getFixedEnv().entrySet()) {
446446
message.append(" Environment variable: ");
447447
message.append(ShellEscaper.escapeString(entry.getKey()));
448448
message.append('=');
@@ -551,7 +551,7 @@ public final ImmutableMap<String, String> getIncompleteEnvironmentForTesting() {
551551
// ActionEnvironment to avoid developers misunderstanding the purpose of this method. That
552552
// requires first updating all subclasses and callers to actually handle environments correctly,
553553
// so it's not a small change.
554-
return env.getFixedEnv().toMap();
554+
return env.getFixedEnv();
555555
}
556556

557557
/**

‎src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfiguration.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -442,7 +442,7 @@ public boolean isSiblingRepositoryLayout() {
442442
*/
443443
@Override
444444
public ImmutableMap<String, String> getLocalShellEnvironment() {
445-
return actionEnv.getFixedEnv().toMap();
445+
return actionEnv.getFixedEnv();
446446
}
447447

448448
/**
@@ -579,7 +579,7 @@ public boolean legacyExternalRunfiles() {
579579
*/
580580
@Override
581581
public ImmutableMap<String, String> getTestEnv() {
582-
return testEnv.getFixedEnv().toMap();
582+
return testEnv.getFixedEnv();
583583
}
584584

585585
/**

‎src/main/java/com/google/devtools/build/lib/analysis/test/TestActionBuilder.java‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,9 @@
5555
import com.google.devtools.build.lib.vfs.PathFragment;
5656
import java.util.List;
5757
import java.util.Map;
58+
import java.util.Set;
5859
import java.util.TreeMap;
60+
import java.util.TreeSet;
5961
import javax.annotation.Nullable;
6062

6163
/**
@@ -96,10 +98,12 @@ public NestedSet<PackageGroupContents> getPackageSpecifications() {
9698
private InstrumentedFilesInfo instrumentedFiles;
9799
private int explicitShardCount;
98100
private final Map<String, String> extraEnv;
101+
private final Set<String> extraInheritedEnv;
99102

100103
public TestActionBuilder(RuleContext ruleContext) {
101104
this.ruleContext = ruleContext;
102105
this.extraEnv = new TreeMap<>();
106+
this.extraInheritedEnv = new TreeSet<>();
103107
this.additionalTools = new ImmutableList.Builder<>();
104108
}
105109

@@ -154,6 +158,11 @@ public TestActionBuilder addExtraEnv(Map<String, String> extraEnv) {
154158
return this;
155159
}
156160

161+
public TestActionBuilder addExtraInheritedEnv(List<String> extraInheritedEnv) {
162+
this.extraInheritedEnv.addAll(extraInheritedEnv);
163+
return this;
164+
}
165+
157166
/**
158167
* Set the explicit shard count. Note that this may be overridden by the sharding strategy.
159168
*/
@@ -442,7 +451,9 @@ private TestParams createTestAction(int shards)
442451
coverageArtifact,
443452
coverageDirectory,
444453
testProperties,
445-
runfilesSupport.getActionEnvironment().addFixedVariables(extraTestEnv),
454+
runfilesSupport
455+
.getActionEnvironment()
456+
.addVariables(extraTestEnv, extraInheritedEnv),
446457
executionSettings,
447458
shard,
448459
run,

0 commit comments

Comments
 (0)