Skip to content

Commit

Permalink
Remove unnecessary layer of indirection in ActionEnvironment.
Browse files Browse the repository at this point in the history
Make `ActionEnvironment` itself an abstract class to capture the simple and compound cases and remove the internal `EnvironmentVariables` interface.

There's likely more potential memory savings around this class which I am investigating.

PiperOrigin-RevId: 530881470
Change-Id: I7df44faac3e94046d552f9392165dbb5ca6badf0
  • Loading branch information
justinhorvitz authored and copybara-github committed May 10, 2023
1 parent 6945edd commit df728bd
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 185 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -42,125 +42,30 @@
* action cache), such that Bazel knows exactly which actions it needs to rerun, and does not have
* to reanalyze the entire dependency graph.
*/
public final class ActionEnvironment {
public abstract class ActionEnvironment {

/**
* A map of environment variables and their values together with a list of variables whose values
* should be inherited from the client environment.
*/
public interface EnvironmentVariables {

/**
* 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 allocates additional objects if the underlying implementation is a {@link
* CompoundEnvironmentVariables}; use sparingly.
*/
ImmutableSet<String> getInheritedEnvironment();
public static final ActionEnvironment EMPTY = new EmptyActionEnvironment();

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

default int size() {
return getFixedEnvironment().size() + getInheritedEnvironment().size();
}
/** Convenience method for creating an {@link ActionEnvironment} with no inherited variables. */
public static ActionEnvironment create(ImmutableMap<String, String> fixedEnv) {
return create(fixedEnv, /* inheritedEnv= */ ImmutableSet.of());
}

/**
* An {@link EnvironmentVariables} that combines variables from two different environments without
* allocating a new map.
* Creates a new {@link ActionEnvironment}.
*
* <p>If an environment variable is contained both as a key in {@code fixedEnv} and in {@code
* inheritedEnv}, the result of {@link #resolve} will contain the value inherited from the client
* environment.
*/
static class CompoundEnvironmentVariables implements EnvironmentVariables {

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

private final EnvironmentVariables current;
private final EnvironmentVariables base;

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

@Override
public boolean isEmpty() {
return current.isEmpty() && base.isEmpty();
}

@Override
public ImmutableMap<String, String> getFixedEnvironment() {
ImmutableMap.Builder<String, String> result = new ImmutableMap.Builder<>();
result.putAll(base.getFixedEnvironment());
result.putAll(current.getFixedEnvironment());
return result.buildKeepingLast();
}

@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> fixedVars, Set<String> inheritedVars) {
if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
return EMPTY_ENVIRONMENT_VARIABLES;
}
return new SimpleEnvironmentVariables(fixedVars, inheritedVars);
}

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

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 ImmutableSet<String> getInheritedEnvironment() {
return inheritedVars;
public static ActionEnvironment create(
ImmutableMap<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
if (fixedEnv.isEmpty() && inheritedEnv.isEmpty()) {
return EMPTY;
}
return new SimpleActionEnvironment(fixedEnv, inheritedEnv);
}

/** An empty {@link EnvironmentVariables}. */
public static final EnvironmentVariables EMPTY_ENVIRONMENT_VARIABLES =
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);

/**
* Splits the given map into a map of variables with a fixed value, and a set of variables that
* should be inherited, the latter of which are identified by having a {@code null} value in the
Expand All @@ -173,117 +78,147 @@ public static ActionEnvironment split(Map<String, String> env) {
if (entry.getValue() != null) {
fixedEnv.put(entry.getKey(), entry.getValue());
} else {
String key = entry.getKey();
inheritedEnv.add(key);
inheritedEnv.add(entry.getKey());
}
}
return create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
return create(ImmutableMap.copyOf(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
}

private final EnvironmentVariables vars;
private ActionEnvironment() {}

private ActionEnvironment(EnvironmentVariables vars) {
this.vars = vars;
}
/**
* 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 abstract ImmutableMap<String, String> getFixedEnv();

/**
* Creates a new action environment. If an environment variable is contained in both {@link
* EnvironmentVariables#getFixedEnvironment()} and {@link
* EnvironmentVariables#getInheritedEnvironment()}, the result of {@link
* ActionEnvironment#resolve(Map, Map)} will contain the value inherited from the client
* environment.
* Returns the 'inherited' part of the environment, i.e., those environment variables that are
* inherited from the client environment and therefore have no fixed value here. This should only
* be used for testing and to compute the cache keys of actions. Use {@link #resolve} instead to
* get the complete environment.
*/
private static ActionEnvironment create(EnvironmentVariables vars) {
if (vars.isEmpty()) {
return EMPTY;
}
return new ActionEnvironment(vars);
public abstract ImmutableSet<String> getInheritedEnv();

/**
* Returns an upper bound on the combined size of the fixed and inherited environments. A call to
* {@link #resolve} may add fewer entries than this number if environment variables are contained
* in both the fixed and the inherited environment.
*/
public final int size() {
return getFixedEnv().size() + getInheritedEnv().size();
}

/**
* Creates a new action environment. If an environment variable is contained both as a key in
* {@code fixedEnv} and in {@code inheritedEnv}, the result of {@link
* ActionEnvironment#resolve(Map, Map)} will contain the value inherited from the client
* environment.
* Resolves the action environment and adds the resulting entries to the given {@code result} map,
* by looking up any inherited env variables in the given {@code clientEnv}.
*
* <p>We pass in a map to mutate to avoid creating and merging intermediate maps.
*/
public static ActionEnvironment create(
Map<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, inheritedEnv));
public final void resolve(Map<String, String> result, Map<String, String> clientEnv) {
checkNotNull(clientEnv);
result.putAll(getFixedEnv());
for (String var : getInheritedEnv()) {
String value = clientEnv.get(var);
if (value != null) {
result.put(var, value);
}
}
}

public static ActionEnvironment create(Map<String, String> fixedEnv) {
return ActionEnvironment.create(SimpleEnvironmentVariables.create(fixedEnv, ImmutableSet.of()));
public final void addTo(Fingerprint f) {
f.addStringMap(getFixedEnv());
f.addStrings(getInheritedEnv());
}

/**
* 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 withAdditionalFixedVariables(Map<String, String> fixedVars) {
public final ActionEnvironment withAdditionalFixedVariables(Map<String, String> fixedVars) {
return withAdditionalVariables(fixedVars, ImmutableSet.of());
}

/**
* Returns a copy of the environment with the given fixed and inherited variables added to it,
* Returns a copy of this environment with the given fixed and inherited variables added to it,
* <em>overwriting any existing occurrences of those variables</em>.
*/
public ActionEnvironment withAdditionalVariables(
public final ActionEnvironment withAdditionalVariables(
Map<String, String> fixedVars, Set<String> inheritedVars) {
EnvironmentVariables newVars =
CompoundEnvironmentVariables.create(fixedVars, inheritedVars, vars);
if (newVars == vars) {
if (fixedVars.isEmpty() && inheritedVars.isEmpty()) {
return this;
}
return ActionEnvironment.create(newVars);
if (this == EMPTY) {
return new SimpleActionEnvironment(
ImmutableMap.copyOf(fixedVars), ImmutableSet.copyOf(inheritedVars));
}
return new CompoundActionEnvironment(
this, ImmutableMap.copyOf(fixedVars), ImmutableSet.copyOf(inheritedVars));
}

/**
* Returns an upper bound on the combined size of the fixed and inherited environments. A call to
* {@link ActionEnvironment#resolve(Map, Map)} may add less entries than this number if
* environment variables are contained in both the fixed and the inherited environment.
*/
public int size() {
return vars.size();
}
private static final class EmptyActionEnvironment extends ActionEnvironment {

/**
* 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 ImmutableMap<String, String> getFixedEnv() {
return vars.getFixedEnvironment();
}
@Override
public ImmutableMap<String, String> getFixedEnv() {
return ImmutableMap.of();
}

/**
* Returns the 'inherited' part of the environment, i.e., those environment variables that are
* inherited from the client environment and therefore have no fixed value here. 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 ImmutableSet<String> getInheritedEnv() {
return vars.getInheritedEnvironment();
@Override
public ImmutableSet<String> getInheritedEnv() {
return ImmutableSet.of();
}
}

/**
* Resolves the action environment and adds the resulting entries to the given {@code result} map,
* by looking up any inherited env variables in the given {@code clientEnv}.
*
* <p>We pass in a map to mutate to avoid creating and merging intermediate maps.
*/
public void resolve(Map<String, String> result, Map<String, String> clientEnv) {
checkNotNull(clientEnv);
result.putAll(vars.getFixedEnvironment());
for (String var : vars.getInheritedEnvironment()) {
String value = clientEnv.get(var);
if (value != null) {
result.put(var, value);
}
private static final class SimpleActionEnvironment extends ActionEnvironment {
private final ImmutableMap<String, String> fixedEnv;
private final ImmutableSet<String> inheritedEnv;

SimpleActionEnvironment(
ImmutableMap<String, String> fixedEnv, ImmutableSet<String> inheritedEnv) {
this.fixedEnv = fixedEnv;
this.inheritedEnv = inheritedEnv;
}

@Override
public ImmutableMap<String, String> getFixedEnv() {
return fixedEnv;
}

@Override
public ImmutableSet<String> getInheritedEnv() {
return inheritedEnv;
}
}

public void addTo(Fingerprint f) {
f.addStringMap(vars.getFixedEnvironment());
f.addStrings(vars.getInheritedEnvironment());
private static final class CompoundActionEnvironment extends ActionEnvironment {
private final ActionEnvironment base;
private final ImmutableMap<String, String> fixedVars;
private final ImmutableSet<String> inheritedVars;

private CompoundActionEnvironment(
ActionEnvironment base,
ImmutableMap<String, String> fixedVars,
ImmutableSet<String> inheritedVars) {
this.base = base;
this.fixedVars = fixedVars;
this.inheritedVars = inheritedVars;
}

@Override
public ImmutableMap<String, String> getFixedEnv() {
return ImmutableMap.<String, String>builder()
.putAll(base.getFixedEnv())
.putAll(fixedVars)
.buildKeepingLast();
}

@Override
public ImmutableSet<String> getInheritedEnv() {
return ImmutableSet.<String>builder()
.addAll(base.getInheritedEnv())
.addAll(inheritedVars)
.build();
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -519,7 +519,8 @@ private static ActionEnvironment computeActionEnvironment(RuleContext ruleContex
}
}
}
return ActionEnvironment.create(fixedEnv, ImmutableSet.copyOf(inheritedEnv));
return ActionEnvironment.create(
ImmutableMap.copyOf(fixedEnv), ImmutableSet.copyOf(inheritedEnv));
}

/** Returns the path of the input manifest of {@code runfilesDir}. */
Expand Down

0 comments on commit df728bd

Please sign in to comment.