Allow rules and targets to specify extra execution platform constraints.

RuleClass.Builder now allows authors to specify whether a rule's targets
can add additional constraints on the execution platform, and to declare
additional constraints for all targets of that rule.

Targets which support this now have an attribute,
"exec_compatible_with", which supports specifying additional constraints
that the execution platform used must match.

This attribute is non-configurable. It will only affect execution
platforms used during toolchain resolution.

Part of #5217.

Change-Id: Id2400dbf869a00aa2be3e3d2f085c2850cd6dc00

Closes #5227.

Change-Id: If7d55f08f7f44bc7d7f6dfec86a3e6bcd68574b9
PiperOrigin-RevId: 199326255

6 files changed, 273 insertions, 5 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/PlatformSemantics.java b/src/main/java/com/google/devtools/build/lib/analysis/PlatformSemantics.java
index bd3e89f97b..65c1047187 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/PlatformSemantics.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/PlatformSemantics.java
@@ -24,6 +24,7 @@ import com.google.devtools.build.lib.packages.RuleClass;
 public class PlatformSemantics {
 
   public static final String TOOLCHAINS_ATTR = "$toolchains";
+  public static final String EXEC_COMPATIBLE_WITH_ATTR = "exec_compatible_with";
 
   public static RuleClass.Builder platformAttributes(RuleClass.Builder builder) {
     return builder
diff --git a/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java b/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
index d83246871d..330278f1aa 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/RuleClass.java
@@ -202,6 +202,23 @@ public class RuleClass {
   }
 
   /**
+   * Describes in which way a rule implementation allows additional execution platform constraints.
+   */
+  public enum ExecutionPlatformConstraintsAllowed {
+    /**
+     * Allows additional execution platform constraints to be added in the rule definition, which
+     * apply to all targets of that rule.
+     */
+    PER_RULE,
+    /**
+     * Users are allowed to specify additional execution platform constraints for each target, using
+     * the 'exec_compatible_with' attribute. This also allows setting constraints in the rule
+     * definition, like PER_RULE.
+     */
+    PER_TARGET;
+  }
+
+  /**
    * For Bazel's constraint system: the attribute that declares the set of environments a rule
    * supports, overriding the defaults for their respective groups.
    */
@@ -615,6 +632,9 @@ public class RuleClass {
     private final Map<String, Attribute> attributes = new LinkedHashMap<>();
     private final Set<Label> requiredToolchains = new HashSet<>();
     private boolean supportsPlatforms = true;
+    private ExecutionPlatformConstraintsAllowed executionPlatformConstraintsAllowed =
+        ExecutionPlatformConstraintsAllowed.PER_RULE;
+    private Set<Label> executionPlatformConstraints = new HashSet<>();
     private OutputFile.Kind outputFileKind = OutputFile.Kind.FILE;
 
     /**
@@ -648,6 +668,8 @@ public class RuleClass {
 
         addRequiredToolchains(parent.getRequiredToolchains());
         supportsPlatforms = parent.supportsPlatforms;
+        // executionPlatformConstraintsAllowed is not inherited and takes the default.
+        addExecutionPlatformConstraints(parent.getExecutionPlatformConstraints());
 
         for (Attribute attribute : parent.getAttributes()) {
           String attrName = attribute.getName();
@@ -656,6 +678,13 @@ public class RuleClass {
               "Attribute %s is inherited multiple times in %s ruleclass",
               attrName,
               name);
+          if (attrName.equals("exec_compatible_with")
+              && parent.executionPlatformConstraintsAllowed
+                  == ExecutionPlatformConstraintsAllowed.PER_TARGET) {
+            // This attribute should not be inherited because executionPlatformConstraintsAllowed is
+            // not inherited.
+            continue;
+          }
           attributes.put(attrName, attribute);
         }
 
@@ -708,6 +737,19 @@ public class RuleClass {
       if (type == RuleClassType.PLACEHOLDER) {
         Preconditions.checkNotNull(ruleDefinitionEnvironmentHashCode, this.name);
       }
+      if (executionPlatformConstraintsAllowed == ExecutionPlatformConstraintsAllowed.PER_TARGET) {
+        // Only rules that allow per target execution constraints need this attribute.
+        Preconditions.checkState(
+            !this.attributes.containsKey("exec_compatible_with"),
+            "Rule should not already define the attribute \"exec_compatible_with\""
+                + " if executionPlatformConstraintsAllowed is set to PER_TARGET");
+        this.add(
+            attr("exec_compatible_with", BuildType.LABEL_LIST)
+                .allowedFileTypes()
+                .nonconfigurable("Used in toolchain resolution")
+                .value(ImmutableList.of()));
+      }
+
       return new RuleClass(
           name,
           key,
@@ -735,6 +777,8 @@ public class RuleClass {
           supportsConstraintChecking,
           requiredToolchains,
           supportsPlatforms,
+          executionPlatformConstraintsAllowed,
+          executionPlatformConstraints,
           outputFileKind,
           attributes.values());
     }
@@ -912,8 +956,8 @@ public class RuleClass {
     /**
      * Applies the given transition factory to all incoming edges for this rule class.
      *
-     * <p>Unlike{@link #cfg(PatchTransition)}, the factory can examine the rule when
-     * deciding what transition to use.
+     * <p>Unlike {@link #cfg(PatchTransition)}, the factory can examine the rule when deciding what
+     * transition to use.
      */
     public Builder cfg(RuleTransitionFactory transitionFactory) {
       Preconditions.checkState(type != RuleClassType.ABSTRACT,
@@ -1161,22 +1205,66 @@ public class RuleClass {
       return this;
     }
 
+    /**
+     * Causes rules of this type to require the specified toolchains be available via toolchain
+     * resolution when a target is configured.
+     */
     public Builder addRequiredToolchains(Iterable<Label> toolchainLabels) {
       Iterables.addAll(this.requiredToolchains, toolchainLabels);
       return this;
     }
 
+    /**
+     * Causes rules of this type to require the specified toolchains be available via toolchain
+     * resolution when a target is configured.
+     */
     public Builder addRequiredToolchains(Label... toolchainLabels) {
-      Iterables.addAll(this.requiredToolchains, Lists.newArrayList(toolchainLabels));
-      return this;
+      return this.addRequiredToolchains(Lists.newArrayList(toolchainLabels));
     }
 
+    /**
+     * Rules that support platforms can use toolchains and execution platforms. Rules that are part
+     * of configuring toolchains and platforms should set this to {@code false}.
+     */
     public Builder supportsPlatforms(boolean flag) {
       this.supportsPlatforms = flag;
       return this;
     }
 
     /**
+     * Specifies whether targets of this rule can add additional constraints on the execution
+     * platform selected. If this is {@link ExecutionPlatformConstraintsAllowed#PER_TARGET}, there
+     * will be an attribute named {@code exec_compatible_with} that can be used to add these
+     * constraints.
+     *
+     * <p>Please note that this value is not inherited by child rules, and must be re-set on them if
+     * the same behavior is required.
+     */
+    public Builder executionPlatformConstraintsAllowed(ExecutionPlatformConstraintsAllowed value) {
+      this.executionPlatformConstraintsAllowed = value;
+      return this;
+    }
+
+    /**
+     * Adds additional execution platform constraints that apply for all targets from this rule.
+     *
+     * <p>Please note that this value is inherited by child rules.
+     */
+    public Builder addExecutionPlatformConstraints(Label... constraints) {
+      return this.addExecutionPlatformConstraints(Lists.newArrayList(constraints));
+    }
+
+    /**
+     * Adds additional execution platform constraints that apply for all targets from this rule.
+     *
+     * <p>Please note that this value is inherited by child rules.
+     */
+    public Builder addExecutionPlatformConstraints(Iterable<Label> constraints) {
+      Iterables.addAll(this.executionPlatformConstraints, constraints);
+      return this;
+    }
+
+    /**
      * Returns an Attribute.Builder object which contains a replica of the
      * same attribute in the parent rule if exists.
      *
@@ -1294,6 +1382,8 @@ public class RuleClass {
 
   private final ImmutableSet<Label> requiredToolchains;
   private final boolean supportsPlatforms;
+  private final ExecutionPlatformConstraintsAllowed executionPlatformConstraintsAllowed;
+  private final ImmutableSet<Label> executionPlatformConstraints;
 
   /**
    * Constructs an instance of RuleClass whose name is 'name', attributes are 'attributes'. The
@@ -1343,6 +1433,8 @@ public class RuleClass {
       boolean supportsConstraintChecking,
       Set<Label> requiredToolchains,
       boolean supportsPlatforms,
+      ExecutionPlatformConstraintsAllowed executionPlatformConstraintsAllowed,
+      Set<Label> executionPlatformConstraints,
       OutputFile.Kind  outputFileKind,
       Collection<Attribute> attributes) {
     this.name = name;
@@ -1375,6 +1467,8 @@ public class RuleClass {
     this.supportsConstraintChecking = supportsConstraintChecking;
     this.requiredToolchains = ImmutableSet.copyOf(requiredToolchains);
     this.supportsPlatforms = supportsPlatforms;
+    this.executionPlatformConstraintsAllowed = executionPlatformConstraintsAllowed;
+    this.executionPlatformConstraints = ImmutableSet.copyOf(executionPlatformConstraints);
 
     // Create the index and collect non-configurable attributes.
     int index = 0;
@@ -2191,6 +2285,14 @@ public class RuleClass {
     return supportsPlatforms;
   }
 
+  public ExecutionPlatformConstraintsAllowed executionPlatformConstraintsAllowed() {
+    return executionPlatformConstraintsAllowed;
+  }
+
+  public ImmutableSet<Label> getExecutionPlatformConstraints() {
+    return executionPlatformConstraints;
+  }
+
   @Nullable
   public OutputFile.Kind  getOutputFileKind() {
     return outputFileKind;
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java
index 1e39c7693c..79498d734b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/AspectFunction.java
@@ -419,6 +419,7 @@ public final class AspectFunction implements SkyFunction {
                     aspect.getDescriptor().getDescription(),
                     associatedConfiguredTargetAndData.getTarget().toString()),
                 requiredToolchains,
+                /* execConstraintLabels= */ ImmutableSet.of(),
                 key.getAspectConfigurationKey());
       } catch (ToolchainContextException e) {
         // TODO(katre): better error handling
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java
index 8363ff7144..f1df1166bc 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ConfiguredTargetFunction.java
@@ -31,6 +31,7 @@ import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.ConfiguredTargetFactory;
 import com.google.devtools.build.lib.analysis.Dependency;
 import com.google.devtools.build.lib.analysis.DependencyResolver.InconsistentAspectOrderException;
+import com.google.devtools.build.lib.analysis.PlatformSemantics;
 import com.google.devtools.build.lib.analysis.TargetAndConfiguration;
 import com.google.devtools.build.lib.analysis.ToolchainContext;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
@@ -57,9 +58,11 @@ import com.google.devtools.build.lib.packages.Attribute;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.NoSuchTargetException;
 import com.google.devtools.build.lib.packages.NoSuchThingException;
+import com.google.devtools.build.lib.packages.NonconfigurableAttributeMapper;
 import com.google.devtools.build.lib.packages.Package;
 import com.google.devtools.build.lib.packages.RawAttributeMapper;
 import com.google.devtools.build.lib.packages.Rule;
+import com.google.devtools.build.lib.packages.RuleClass.ExecutionPlatformConstraintsAllowed;
 import com.google.devtools.build.lib.packages.RuleClassProvider;
 import com.google.devtools.build.lib.packages.RuleTransitionFactory;
 import com.google.devtools.build.lib.packages.Target;
@@ -272,11 +275,15 @@ public final class ConfiguredTargetFunction implements SkyFunction {
         if (rule.getRuleClassObject().supportsPlatforms()) {
           ImmutableSet<Label> requiredToolchains =
               rule.getRuleClassObject().getRequiredToolchains();
+
+          // Collect local (target, rule) constraints for filtering out execution platforms.
+          ImmutableSet<Label> execConstraintLabels = getExecutionPlatformConstraints(rule);
           toolchainContext =
               ToolchainUtil.createToolchainContext(
                   env,
                   rule.toString(),
                   requiredToolchains,
+                  execConstraintLabels,
                   configuredTargetKey.getConfigurationKey());
           if (env.valuesMissing()) {
             return null;
@@ -384,6 +391,25 @@ public final class ConfiguredTargetFunction implements SkyFunction {
   }
 
   /**
+   * Returns the target-specific execution platform constraints, based on the rule definition and
+   * any constraints added by the target.
+   */
+  private static ImmutableSet<Label> getExecutionPlatformConstraints(Rule rule) {
+    NonconfigurableAttributeMapper mapper = NonconfigurableAttributeMapper.of(rule);
+    ImmutableSet.Builder<Label> execConstraintLabels = new ImmutableSet.Builder<>();
+
+    execConstraintLabels.addAll(rule.getRuleClassObject().getExecutionPlatformConstraints());
+
+    if (rule.getRuleClassObject().executionPlatformConstraintsAllowed()
+        == ExecutionPlatformConstraintsAllowed.PER_TARGET) {
+      execConstraintLabels.addAll(
+          mapper.get(PlatformSemantics.EXEC_COMPATIBLE_WITH_ATTR, BuildType.LABEL_LIST));
+    }
+
+    return execConstraintLabels.build();
+  }
+
+  /**
    * Computes the direct dependencies of a node in the configured target graph (a configured target
    * or an aspects).
    *
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 906924fa9f..6cf2bd5d77 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -968,7 +968,11 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     SkyFunctionEnvironmentForTesting env =
         new SkyFunctionEnvironmentForTesting(buildDriver, eventHandler, this);
     return ToolchainUtil.createToolchainContext(
-        env, "", requiredToolchains, config == null ? null : BuildConfigurationValue.key(config));
+        env,
+        "",
+        requiredToolchains,
+        /* execConstraintLabels= */ ImmutableSet.of(),
+        config == null ? null : BuildConfigurationValue.key(config));
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java
index 40d6701c33..39149f54e4 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ToolchainUtil.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.skyframe;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.stream.Collectors.joining;
 
 import com.google.auto.value.AutoValue;
@@ -28,6 +29,7 @@ import com.google.devtools.build.lib.analysis.PlatformConfiguration;
 import com.google.devtools.build.lib.analysis.PlatformOptions;
 import com.google.devtools.build.lib.analysis.ToolchainContext;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
+import com.google.devtools.build.lib.analysis.platform.ConstraintValueInfo;
 import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
 import com.google.devtools.build.lib.analysis.platform.PlatformProviderUtils;
 import com.google.devtools.build.lib.cmdline.Label;
@@ -52,17 +54,26 @@ import javax.annotation.Nullable;
 /**
  * Common code to create a {@link ToolchainContext} given a set of required toolchain type labels.
  */
+// TODO(katre): Refactor this and ToolchainContext into something nicer to work with and with
+// fewer static methods everywhere.
 public class ToolchainUtil {
 
   /**
    * Returns a new {@link ToolchainContext}, with the correct toolchain labels based on the results
    * of the {@link ToolchainResolutionFunction}.
+   *
+   * @param env the Skyframe environment to use to acquire dependencies
+   * @param targetDescription a description of the target use, for error and debug message context
+   * @param requiredToolchains the required toolchain types that must be resolved
+   * @param execConstraintLabels extra constraints on the execution platform to select
+   * @param configurationKey the build configuration to use for resolving other targets
    */
   @Nullable
   static ToolchainContext createToolchainContext(
       Environment env,
       String targetDescription,
       Set<Label> requiredToolchains,
+      Set<Label> execConstraintLabels,
       @Nullable BuildConfigurationValue.Key configurationKey)
       throws ToolchainContextException, InterruptedException {
 
@@ -91,6 +102,11 @@ public class ToolchainUtil {
     ConfiguredTargetKey hostPlatformKey = ConfiguredTargetKey.of(hostPlatformLabel, configuration);
     ConfiguredTargetKey targetPlatformKey =
         ConfiguredTargetKey.of(targetPlatformLabel, configuration);
+    ImmutableList<ConfiguredTargetKey> execConstraintKeys =
+        execConstraintLabels
+            .stream()
+            .map(label -> ConfiguredTargetKey.of(label, configuration))
+            .collect(toImmutableList());
 
     // Load the host and target platforms early, to check for errors.
     getPlatformInfo(ImmutableList.of(hostPlatformKey, targetPlatformKey), env);
@@ -108,6 +124,14 @@ public class ToolchainUtil {
             .addAll(registeredExecutionPlatforms.registeredExecutionPlatformKeys())
             .add(hostPlatformKey)
             .build();
+
+    // Filter out execution platforms that don't satisfy the extra constraints.
+    availableExecutionPlatformKeys =
+        filterPlatforms(availableExecutionPlatformKeys, execConstraintKeys, env);
+    if (availableExecutionPlatformKeys == null) {
+      return null;
+    }
+
     Optional<ResolvedToolchains> resolvedToolchains =
         resolveToolchainLabels(
             env,
@@ -421,6 +445,94 @@ public class ToolchainUtil {
     return labels.build();
   }
 
+  @Nullable
+  private static ImmutableList<ConfiguredTargetKey> filterPlatforms(
+      ImmutableList<ConfiguredTargetKey> platformKeys,
+      ImmutableList<ConfiguredTargetKey> constraintKeys,
+      Environment env)
+      throws ToolchainContextException, InterruptedException {
+
+    // Short circuit if not needed.
+    if (constraintKeys.isEmpty()) {
+      return platformKeys;
+    }
+
+    Map<ConfiguredTargetKey, PlatformInfo> platformInfoMap = getPlatformInfo(platformKeys, env);
+    if (platformInfoMap == null) {
+      return null;
+    }
+    List<ConstraintValueInfo> constraints = getConstraintValueInfo(constraintKeys, env);
+    if (constraints == null) {
+      return null;
+    }
+
+    return platformKeys
+        .stream()
+        .filter(key -> filterPlatform(platformInfoMap.get(key), constraints))
+        .collect(toImmutableList());
+  }
+
+  @Nullable
+  private static List<ConstraintValueInfo> getConstraintValueInfo(
+      ImmutableList<ConfiguredTargetKey> constraintKeys, Environment env)
+      throws InterruptedException, ToolchainContextException {
+
+    Map<SkyKey, ValueOrException<ConfiguredValueCreationException>> values =
+        env.getValuesOrThrow(constraintKeys, ConfiguredValueCreationException.class);
+    boolean valuesMissing = env.valuesMissing();
+    List<ConstraintValueInfo> constraintValues = valuesMissing ? null : new ArrayList<>();
+    try {
+      for (ConfiguredTargetKey key : constraintKeys) {
+        ConstraintValueInfo constraintValueInfo = findConstraintValueInfo(values.get(key));
+        if (!valuesMissing && constraintValueInfo != null) {
+          constraintValues.add(constraintValueInfo);
+        }
+      }
+    } catch (ConfiguredValueCreationException e) {
+      throw new ToolchainContextException(e);
+    }
+    if (valuesMissing) {
+      return null;
+    }
+    return constraintValues;
+  }
+
+  @Nullable
+  private static ConstraintValueInfo findConstraintValueInfo(
+      ValueOrException<ConfiguredValueCreationException> valueOrException)
+      throws ConfiguredValueCreationException, ToolchainContextException {
+
+    ConfiguredTargetValue configuredTargetValue = (ConfiguredTargetValue) valueOrException.get();
+    if (configuredTargetValue == null) {
+      return null;
+    }
+
+    ConfiguredTarget configuredTarget = configuredTargetValue.getConfiguredTarget();
+    ConstraintValueInfo constraintValueInfo =
+        PlatformProviderUtils.constraintValue(configuredTarget);
+    if (constraintValueInfo == null) {
+      throw new ToolchainContextException(
+          new InvalidConstraintValueException(configuredTarget.getLabel()));
+    }
+
+    return constraintValueInfo;
+  }
+
+  private static boolean filterPlatform(
+      PlatformInfo platformInfo, List<ConstraintValueInfo> constraints) {
+    for (ConstraintValueInfo filterConstraint : constraints) {
+      ConstraintValueInfo platformInfoConstraint =
+          platformInfo.getConstraint(filterConstraint.constraint());
+      if (platformInfoConstraint == null || !platformInfoConstraint.equals(filterConstraint)) {
+        // The value for this setting is not present in the platform, or doesn't match the expected
+        // value.
+        return false;
+      }
+    }
+
+    return true;
+  }
+
   /**
    * Exception used when an error occurs in {@link #expandTargetPatterns(Environment, List,
    * FilteringPolicy)}.
@@ -460,6 +572,24 @@ public class ToolchainUtil {
     }
   }
 
+  /** Exception used when a constraint value label is not a valid constraint value. */
+  static final class InvalidConstraintValueException extends Exception {
+    InvalidConstraintValueException(Label label) {
+      super(formatError(label));
+    }
+
+    InvalidConstraintValueException(Label label, ConfiguredValueCreationException e) {
+      super(formatError(label), e);
+    }
+
+    private static String formatError(Label label) {
+      return String.format(
+          "Target %s was referenced as a constraint_value,"
+              + " but does not provide ConstraintValueInfo",
+          label);
+    }
+  }
+
   /** Exception used when a toolchain type is required but no matching toolchain is found. */
   public static final class UnresolvedToolchainsException extends Exception {
     private final ImmutableList<Label> missingToolchainTypes;
@@ -483,6 +613,10 @@ public class ToolchainUtil {
       super(e);
     }
 
+    public ToolchainContextException(InvalidConstraintValueException e) {
+      super(e);
+    }
+
     public ToolchainContextException(UnresolvedToolchainsException e) {
       super(e);
     }