From dbb08efcff1daa0f6fe133c3a09b351a5b00b034 Mon Sep 17 00:00:00 2001 From: Fabian Meumertzheim Date: Wed, 16 Mar 2022 22:58:17 +0100 Subject: [PATCH] Make adding inherited vars to ActionEnvironment memory efficient Along the way, this commit optimizes away a few potential allocations of empty ActionEnvironment instances. --- .../build/lib/actions/AbstractAction.java | 2 +- .../build/lib/actions/ActionEnvironment.java | 149 +++++++++++------- .../lib/analysis/actions/SpawnAction.java | 4 +- .../config/BuildConfigurationValue.java | 4 +- .../lib/analysis/test/TestRunnerAction.java | 2 +- ...ctionGraphTextOutputFormatterCallback.java | 2 +- .../lib/rules/java/JavaCompileAction.java | 2 +- .../actiongraph/v2/ActionGraphDump.java | 2 +- .../lib/actions/ActionEnvironmentTest.java | 4 +- .../rules/BazelRuleClassProviderTest.java | 4 +- 10 files changed, 104 insertions(+), 71 deletions(-) diff --git a/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java b/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java index 408371b8dcf84f..f54952b196c7e3 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java +++ b/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java @@ -699,7 +699,7 @@ public Dict getExecutionInfoDict() { @Override public Dict getEnv() { - return Dict.immutableCopyOf(env.getFixedEnv().toMap()); + return Dict.immutableCopyOf(env.getFixedEnv()); } @Override diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java b/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java index 0fd9ba1c7b5e09..016eef25497d08 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ActionEnvironment.java @@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableSet; import com.google.devtools.build.lib.util.Fingerprint; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.Map; import java.util.Set; import java.util.TreeMap; @@ -45,23 +46,34 @@ */ 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. + * + *

WARNING: this allocates additional objects if the underlying implementation is a {@link + * CompoundEnvironmentVariables}; use sparingly. + */ + ImmutableMap getFixedEnvironment(); + + /** + * Returns the inherited environment variables as a set. * - *

WARNING: this allocations additional objects if the underlying implementation is a {@link + *

WARNING: this allocated additional objects if the underlying implementation is a {@link * CompoundEnvironmentVariables}; use sparingly. */ - ImmutableMap toMap(); + ImmutableSet getInheritedEnvironment(); default boolean isEmpty() { - return toMap().isEmpty(); + return getFixedEnvironment().isEmpty() && getInheritedEnvironment().isEmpty(); } default int size() { - return toMap().size(); + return getFixedEnvironment().size() + getInheritedEnvironment().size(); } } @@ -70,11 +82,21 @@ default int size() { * allocation a new map. */ static class CompoundEnvironmentVariables implements EnvironmentVariables { + + static EnvironmentVariables create(Map fixedVars, Set 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 vars, EnvironmentVariables base) { - this.current = new SimpleEnvironmentVariables(vars); + private CompoundEnvironmentVariables(Map fixedVars, Set inheritedVars, + EnvironmentVariables base) { + this.current = SimpleEnvironmentVariables.create(fixedVars, inheritedVars); this.base = base; } @@ -84,39 +106,58 @@ public boolean isEmpty() { } @Override - public ImmutableMap toMap() { + public ImmutableMap getFixedEnvironment() { Map result = new LinkedHashMap<>(); - result.putAll(base.toMap()); - result.putAll(current.toMap()); + result.putAll(base.getFixedEnvironment()); + result.putAll(current.getFixedEnvironment()); return ImmutableMap.copyOf(result); } + + @Override + public ImmutableSet getInheritedEnvironment() { + Set result = new LinkedHashSet<>(); + result.addAll(base.getInheritedEnvironment()); + result.addAll(current.getInheritedEnvironment()); + return ImmutableSet.copyOf(result); + } } - /** A simple {@link EnvironmentVariables}. */ + /** + * A simple {@link EnvironmentVariables}. + */ static class SimpleEnvironmentVariables implements EnvironmentVariables { - static EnvironmentVariables create(Map vars) { - if (vars.isEmpty()) { + static EnvironmentVariables create(Map fixedVars, Set inheritedVars) { + if (fixedVars.isEmpty() && inheritedVars.isEmpty()) { return EMPTY_ENVIRONMENT_VARIABLES; } - return new SimpleEnvironmentVariables(vars); + return new SimpleEnvironmentVariables(fixedVars, inheritedVars); } - private final ImmutableMap vars; + private final ImmutableMap fixedVars; + private final ImmutableSet inheritedVars; + + private SimpleEnvironmentVariables(Map fixedVars, Set inheritedVars) { + this.fixedVars = ImmutableMap.copyOf(fixedVars); + this.inheritedVars = ImmutableSet.copyOf(inheritedVars); + } - private SimpleEnvironmentVariables(Map vars) { - this.vars = ImmutableMap.copyOf(vars); + @Override + public ImmutableMap getFixedEnvironment() { + return fixedVars; } @Override - public ImmutableMap toMap() { - return vars; + public ImmutableSet getInheritedEnvironment() { + return inheritedVars; } } - /** An empty {@link EnvironmentVariables}. */ + /** + * 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 @@ -124,8 +165,7 @@ public ImmutableMap toMap() { */ // 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 @@ -145,15 +185,13 @@ public static ActionEnvironment split(Map 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 inheritedEnv; + private final EnvironmentVariables vars; - private ActionEnvironment(EnvironmentVariables fixedEnv, ImmutableSet inheritedEnv) { - this.fixedEnv = fixedEnv; - this.inheritedEnv = inheritedEnv; + private ActionEnvironment(EnvironmentVariables vars) { + this.vars = vars; } /** @@ -161,51 +199,46 @@ private ActionEnvironment(EnvironmentVariables fixedEnv, ImmutableSet in * 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. */ - private static ActionEnvironment create( - EnvironmentVariables fixedEnv, ImmutableSet 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 fixedEnv, ImmutableSet inheritedEnv) { - return new ActionEnvironment(SimpleEnvironmentVariables.create(fixedEnv), inheritedEnv); + return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv)); } public static ActionEnvironment create(Map 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, overwriting * any existing occurrences of those variables. */ - public ActionEnvironment addFixedVariables(Map vars) { - return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), inheritedEnv); + public ActionEnvironment addFixedVariables(Map 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, * overwriting any existing occurrences of those variables. */ - public ActionEnvironment addVariables(Map vars, Set inheritedVars) { - if (inheritedVars.isEmpty()) { - return addFixedVariables(vars); - } else { - // TODO: inheritedEnv is currently not optimized for allocation avoidance in the same way as - // fixedEnv. - ImmutableSet newInheritedEnv = ImmutableSet.builder().addAll(inheritedEnv) - .addAll(inheritedVars).build(); - return new ActionEnvironment(new CompoundEnvironmentVariables(vars, fixedEnv), - newInheritedEnv); - } + public ActionEnvironment addVariables(Map fixedVars, Set inheritedVars) { + return ActionEnvironment.create( + CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars)); } - /** Returns the combined size of the fixed and inherited environments. */ + /** + * Returns the combined size of the fixed and inherited environments. + */ public int size() { - return fixedEnv.size() + inheritedEnv.size(); + return vars.size(); } /** @@ -213,8 +246,8 @@ public int size() { * 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 getFixedEnv() { + return vars.getFixedEnvironment(); } /** @@ -224,7 +257,7 @@ public EnvironmentVariables getFixedEnv() { * get the complete environment. */ public ImmutableSet getInheritedEnv() { - return inheritedEnv; + return vars.getInheritedEnvironment(); } /** @@ -235,8 +268,8 @@ public ImmutableSet getInheritedEnv() { */ public void resolve(Map result, Map 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); @@ -245,7 +278,7 @@ public void resolve(Map result, Map clientEnv) { } public void addTo(Fingerprint f) { - f.addStringMap(fixedEnv.toMap()); - f.addStrings(inheritedEnv); + f.addStringMap(vars.getFixedEnvironment()); + f.addStrings(vars.getInheritedEnvironment()); } } diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java index 00690adcad19a6..2c7b93febc290a 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java @@ -465,7 +465,7 @@ public String describeKey() { StringBuilder message = new StringBuilder(); message.append(getProgressMessage()); message.append('\n'); - for (Map.Entry entry : env.getFixedEnv().toMap().entrySet()) { + for (Map.Entry entry : env.getFixedEnv().entrySet()) { message.append(" Environment variable: "); message.append(ShellEscaper.escapeString(entry.getKey())); message.append('='); @@ -574,7 +574,7 @@ public final ImmutableMap 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(); } /** Returns the out-of-band execution data for this action. */ diff --git a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationValue.java b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationValue.java index 1495aeed74b857..bcaff069ad8a9e 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationValue.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/config/BuildConfigurationValue.java @@ -485,7 +485,7 @@ public boolean isSiblingRepositoryLayout() { */ @Override public ImmutableMap getLocalShellEnvironment() { - return actionEnv.getFixedEnv().toMap(); + return actionEnv.getFixedEnv(); } /** @@ -629,7 +629,7 @@ public boolean legacyExternalRunfiles() { */ @Override public ImmutableMap getTestEnv() { - return testEnv.getFixedEnv().toMap(); + return testEnv.getFixedEnv(); } /** diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java index c68157faf7843b..4cce0a6323303b 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java @@ -1051,7 +1051,7 @@ public List getArguments() throws CommandLineExpansionException, Interru @Override public ImmutableMap getIncompleteEnvironmentForTesting() throws ActionExecutionException { - return getEnvironment().getFixedEnv().toMap(); + return getEnvironment().getFixedEnv(); } @Override diff --git a/src/main/java/com/google/devtools/build/lib/query2/aquery/ActionGraphTextOutputFormatterCallback.java b/src/main/java/com/google/devtools/build/lib/query2/aquery/ActionGraphTextOutputFormatterCallback.java index 91ec8a19144911..201ec27f0d8425 100644 --- a/src/main/java/com/google/devtools/build/lib/query2/aquery/ActionGraphTextOutputFormatterCallback.java +++ b/src/main/java/com/google/devtools/build/lib/query2/aquery/ActionGraphTextOutputFormatterCallback.java @@ -221,7 +221,7 @@ private void writeAction(ActionAnalysisMetadata action, PrintStream printStream) // TODO(twerth): This handles the fixed environment. We probably want to output the inherited // environment as well. Iterable> fixedEnvironment = - abstractAction.getEnvironment().getFixedEnv().toMap().entrySet(); + abstractAction.getEnvironment().getFixedEnv().entrySet(); if (!Iterables.isEmpty(fixedEnvironment)) { stringBuilder .append(" Environment: [") diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java index 16e289b0b90b70..eea2bea665c89d 100644 --- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java +++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java @@ -614,7 +614,7 @@ public final ImmutableMap 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(); } @Override diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java b/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java index 5b326c56a9ddf8..5ff47004347ebc 100644 --- a/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java +++ b/src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java @@ -161,7 +161,7 @@ private void dumpSingleAction(ConfiguredTarget configuredTarget, ActionAnalysisM SpawnAction spawnAction = (SpawnAction) action; // TODO(twerth): This handles the fixed environment. We probably want to output the inherited // environment as well. - Map fixedEnvironment = spawnAction.getEnvironment().getFixedEnv().toMap(); + Map fixedEnvironment = spawnAction.getEnvironment().getFixedEnv(); for (Map.Entry environmentVariable : fixedEnvironment.entrySet()) { actionBuilder.addEnvironmentVariables( AnalysisProtosV2.KeyValuePair.newBuilder() diff --git a/src/test/java/com/google/devtools/build/lib/actions/ActionEnvironmentTest.java b/src/test/java/com/google/devtools/build/lib/actions/ActionEnvironmentTest.java index 8ff8e445cc7e97..d48d266db2f373 100644 --- a/src/test/java/com/google/devtools/build/lib/actions/ActionEnvironmentTest.java +++ b/src/test/java/com/google/devtools/build/lib/actions/ActionEnvironmentTest.java @@ -34,10 +34,10 @@ public void compoundEnvOrdering() { // entries added by env2 override the existing entries ActionEnvironment env2 = env1.addFixedVariables(ImmutableMap.of("FOO", "foo2")); - assertThat(env1.getFixedEnv().toMap()).containsExactly("FOO", "foo1", "BAR", "bar"); + assertThat(env1.getFixedEnv()).containsExactly("FOO", "foo1", "BAR", "bar"); assertThat(env1.getInheritedEnv()).containsExactly("baz"); - assertThat(env2.getFixedEnv().toMap()).containsExactly("FOO", "foo2", "BAR", "bar"); + assertThat(env2.getFixedEnv()).containsExactly("FOO", "foo2", "BAR", "bar"); assertThat(env2.getInheritedEnv()).containsExactly("baz"); } } diff --git a/src/test/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProviderTest.java b/src/test/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProviderTest.java index 769b5ae740e77b..18e1179a0b78e9 100644 --- a/src/test/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProviderTest.java +++ b/src/test/java/com/google/devtools/build/lib/bazel/rules/BazelRuleClassProviderTest.java @@ -176,8 +176,8 @@ public void strictActionEnv() throws Exception { "--action_env=FOO=bar"); ActionEnvironment env = BazelRuleClassProvider.SHELL_ACTION_ENV.apply(options); - assertThat(env.getFixedEnv().toMap()).containsEntry("PATH", "/bin:/usr/bin:/usr/local/bin"); - assertThat(env.getFixedEnv().toMap()).containsEntry("FOO", "bar"); + assertThat(env.getFixedEnv()).containsEntry("PATH", "/bin:/usr/bin:/usr/local/bin"); + assertThat(env.getFixedEnv()).containsEntry("FOO", "bar"); } @Test