Allow Skylark rules to specify whether targets can add execution platform constraints.

Closes #5341.

Change-Id: Ib74e59fec48102469a5039e045e3f3d0e0d86d8c
PiperOrigin-RevId: 200526448

4 files changed, 285 insertions, 1 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleClassFunctions.java
index 9237cea109..b1eef7dc85 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleClassFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/skylark/SkylarkRuleClassFunctions.java
@@ -318,6 +318,8 @@ public class SkylarkRuleClassFunctions implements SkylarkRuleFunctionsApi<Artifa
       SkylarkList<String> toolchains,
       String doc,
       SkylarkList<?> providesArg,
+      Boolean executionPlatformConstraintsAllowed,
+      SkylarkList<?> execCompatibleWith,
       FuncallExpression ast,
       Environment funcallEnv)
       throws EvalException, ConversionException {
@@ -399,6 +401,16 @@ public class SkylarkRuleClassFunctions implements SkylarkRuleFunctionsApi<Artifa
       builder.advertiseSkylarkProvider(skylarkProvider);
     }
 
+    if (!execCompatibleWith.isEmpty()) {
+      builder.addExecutionPlatformConstraints(
+          collectConstraintLabels(
+              execCompatibleWith.getContents(String.class, "exec_compatile_with"),
+              ast.getLocation()));
+    }
+    if (executionPlatformConstraintsAllowed) {
+      builder.executionPlatformConstraintsAllowed(ExecutionPlatformConstraintsAllowed.PER_TARGET);
+    }
+
     return new SkylarkRuleFunction(builder, type, attributes, ast.getLocation());
   }
 
@@ -445,6 +457,22 @@ public class SkylarkRuleClassFunctions implements SkylarkRuleFunctionsApi<Artifa
     return requiredToolchains.build();
   }
 
+  private static ImmutableList<Label> collectConstraintLabels(
+      Iterable<String> rawLabels, Location loc) throws EvalException {
+    ImmutableList.Builder<Label> constraintLabels = new ImmutableList.Builder<>();
+    for (String rawLabel : rawLabels) {
+      try {
+        Label constraintLabel = Label.parseAbsolute(rawLabel);
+        constraintLabels.add(constraintLabel);
+      } catch (LabelSyntaxException e) {
+        throw new EvalException(
+            loc, String.format("Unable to parse constraint %s: %s", rawLabel, e.getMessage()), e);
+      }
+    }
+
+    return constraintLabels.build();
+  }
+
   @Override
   public SkylarkAspect aspect(
       BaseFunction implementation,
@@ -657,7 +685,11 @@ public class SkylarkRuleClassFunctions implements SkylarkRuleFunctionsApi<Artifa
         addAttribute(definitionLocation, builder,
             descriptor.build(attribute.getFirst()));
       }
-      this.ruleClass = builder.build(ruleClassName, skylarkLabel + "%" + ruleClassName);
+      try {
+        this.ruleClass = builder.build(ruleClassName, skylarkLabel + "%" + ruleClassName);
+      } catch (IllegalArgumentException | IllegalStateException ex) {
+        throw new EvalException(location, ex);
+      }
 
       this.builder = null;
       this.attributes = null;
diff --git a/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkRuleFunctionsApi.java b/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkRuleFunctionsApi.java
index 53d25e8073..645db6cd72 100644
--- a/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkRuleFunctionsApi.java
+++ b/src/main/java/com/google/devtools/build/lib/skylarkbuildapi/SkylarkRuleFunctionsApi.java
@@ -266,6 +266,30 @@ public interface SkylarkRuleFunctionsApi<FileApiT extends FileApi> {
                 + "It is an error if a provider is listed here and the rule "
                 + "implementation function does not return it."
       ),
+      @Param(
+        name = "execution_platform_constraints_allowed",
+        type = Boolean.class,
+        named = true,
+        positional = false,
+        defaultValue = "False",
+        doc =
+            "If true, a special attribute named <code>exec_compatible_with</code> of "
+                + "label-list type is added, which must not already exist in "
+                + "<code>attrs</code>. Targets may use this attribute to specify additional "
+                + "constraints on the execution platform beyond those given in the "
+                + "<code>exec_compatible_with</code> argument to <code>rule()</code>."
+      ),
+      @Param(
+        name = "exec_compatible_with",
+        type = SkylarkList.class,
+        generic1 = String.class,
+        named = true,
+        positional = false,
+        defaultValue = "[]",
+        doc =
+            "A list of constraints on the execution platform that apply to all targets of "
+                + "this rule type."
+      )
     },
     useAst = true,
     useEnvironment = true
@@ -283,6 +307,8 @@ public interface SkylarkRuleFunctionsApi<FileApiT extends FileApi> {
       SkylarkList<String> toolchains,
       String doc,
       SkylarkList<?> providesArg,
+      Boolean executionPlatformConstraintsAllowed,
+      SkylarkList<?> execCompatibleWith,
       FuncallExpression ast,
       Environment funcallEnv)
       throws EvalException;
diff --git a/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleClassFunctionsTest.java b/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleClassFunctionsTest.java
index e1029ecf0a..dc1a170f40 100644
--- a/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleClassFunctionsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skylark/SkylarkRuleClassFunctionsTest.java
@@ -39,6 +39,7 @@ import com.google.devtools.build.lib.packages.PredicateWithMessage;
 import com.google.devtools.build.lib.packages.RequiredProviders;
 import com.google.devtools.build.lib.packages.RuleClass;
 import com.google.devtools.build.lib.packages.RuleClass.Builder.RuleClassType;
+import com.google.devtools.build.lib.packages.RuleClass.ExecutionPlatformConstraintsAllowed;
 import com.google.devtools.build.lib.packages.SkylarkAspectClass;
 import com.google.devtools.build.lib.packages.SkylarkDefinedAspect;
 import com.google.devtools.build.lib.packages.SkylarkInfo;
@@ -1625,6 +1626,51 @@ public class SkylarkRuleClassFunctionsTest extends SkylarkTestCase {
   }
 
   @Test
+  public void testRuleAddExecutionConstraints() throws Exception {
+    registerDummyUserDefinedFunction();
+    scratch.file("test/BUILD", "toolchain_type(name = 'my_toolchain_type')");
+    evalAndExport(
+        "r1 = rule(",
+        "  implementation = impl,",
+        "  toolchains=['//test:my_toolchain_type'],",
+        "  exec_compatible_with=['//constraint:cv1', '//constraint:cv2'],",
+        ")");
+    RuleClass c = ((SkylarkRuleFunction) lookup("r1")).getRuleClass();
+    assertThat(c.getExecutionPlatformConstraints())
+        .containsExactly(makeLabel("//constraint:cv1"), makeLabel("//constraint:cv2"));
+  }
+
+  @Test
+  public void testTargetsCanAddExecutionPlatformConstraints() throws Exception {
+    registerDummyUserDefinedFunction();
+    scratch.file("test/BUILD", "toolchain_type(name = 'my_toolchain_type')");
+    evalAndExport(
+        "r1 = rule(impl, ",
+        "  toolchains=['//test:my_toolchain_type'],",
+        "  execution_platform_constraints_allowed=True,",
+        ")");
+    RuleClass c = ((SkylarkRuleFunction) lookup("r1")).getRuleClass();
+    assertThat(c.executionPlatformConstraintsAllowed())
+        .isEqualTo(ExecutionPlatformConstraintsAllowed.PER_TARGET);
+  }
+
+  @Test
+  public void testTargetsCanAddExecutionPlatformConstraints_attrAlreadyDefined() throws Exception {
+    registerDummyUserDefinedFunction();
+    scratch.file("test/BUILD", "toolchain_type(name = 'my_toolchain_type')");
+    ev.setFailFast(false);
+    evalAndExport(
+        "r1 = rule(impl, ",
+        "  attrs = {",
+        "    'exec_compatible_with': attr.label_list(),",
+        "  },",
+        "  toolchains=['//test:my_toolchain_type'],",
+        "  execution_platform_constraints_allowed=True,",
+        ")");
+    ev.assertContainsError("Rule should not already define the attribute \"exec_compatible_with\"");
+  }
+
+  @Test
   public void testRuleFunctionReturnsNone() throws Exception {
     scratch.file("test/rule.bzl",
         "def _impl(ctx):",
diff --git a/src/test/shell/bazel/toolchain_test.sh b/src/test/shell/bazel/toolchain_test.sh
index 748182c0bc..097e874641 100755
--- a/src/test/shell/bazel/toolchain_test.sh
+++ b/src/test/shell/bazel/toolchain_test.sh
@@ -734,4 +734,184 @@ EOF
 }
 
 
+function test_rule_with_default_execution_constraints() {
+  write_test_toolchain
+  write_register_toolchain
+
+  # Add test platforms.
+  mkdir -p platforms
+  cat >> platforms/BUILD <<EOF
+constraint_setting(name = 'setting')
+constraint_value(name = 'value1', constraint_setting = ':setting')
+constraint_value(name = 'value2', constraint_setting = ':setting')
+
+platform(
+    name = 'platform1',
+    constraint_values = [':value1'],
+    visibility = ['//visibility:public'])
+platform(
+    name = 'platform2',
+    constraint_values = [':value2'],
+    visibility = ['//visibility:public'])
+EOF
+
+  # Add a rule with default execution constraints.
+  mkdir -p demo
+  cat >> demo/rule.bzl <<EOF
+def _impl(ctx):
+  return []
+
+sample_rule = rule(
+  implementation = _impl,
+  attrs = {},
+  exec_compatible_with = [
+    '//platforms:value2',
+  ],
+  toolchains = ['//toolchain:test_toolchain'],
+)
+EOF
+
+  # Use the new rule.
+  cat >> demo/BUILD <<EOF
+load(':rule.bzl', 'sample_rule')
+
+sample_rule(name = 'use')
+EOF
+
+  # Build the target, using debug messages to verify the correct platform was selected.
+  bazel build \
+    --extra_execution_platforms=//platforms:all \
+    --toolchain_resolution_debug \
+    //demo:use &> $TEST_log || fail "Build failed"
+    expect_log "Selected execution platform //platforms:platform2"
+}
+
+
+function test_target_with_execution_constraints() {
+  write_test_toolchain
+  write_register_toolchain
+
+  # Add test platforms.
+  mkdir -p platforms
+  cat >> platforms/BUILD <<EOF
+package(default_visibility = ['//visibility:public'])
+constraint_setting(name = 'setting')
+constraint_value(name = 'value1', constraint_setting = ':setting')
+constraint_value(name = 'value2', constraint_setting = ':setting')
+
+platform(
+    name = 'platform1',
+    constraint_values = [':value1'],
+    visibility = ['//visibility:public'])
+platform(
+    name = 'platform2',
+    constraint_values = [':value2'],
+    visibility = ['//visibility:public'])
+EOF
+
+  # Add a rule with default execution constraints.
+  mkdir -p demo
+  cat >> demo/rule.bzl <<EOF
+def _impl(ctx):
+  return []
+
+sample_rule = rule(
+  implementation = _impl,
+  attrs = {},
+  toolchains = ['//toolchain:test_toolchain'],
+  execution_platform_constraints_allowed = True,
+)
+EOF
+
+  # Use the new rule.
+  cat >> demo/BUILD <<EOF
+load(':rule.bzl', 'sample_rule')
+
+sample_rule(
+  name = 'use',
+  exec_compatible_with = [
+    '//platforms:value2',
+  ],
+)
+EOF
+
+  # Build the target, using debug messages to verify the correct platform was selected.
+  bazel build \
+    --extra_execution_platforms=//platforms:all \
+    --toolchain_resolution_debug \
+    //demo:use &> $TEST_log || fail "Build failed"
+    expect_log "Selected execution platform //platforms:platform2"
+}
+
+function test_rule_and_target_with_execution_constraints() {
+  write_test_toolchain
+  write_register_toolchain
+
+  # Add test platforms.
+  mkdir -p platforms
+  cat >> platforms/BUILD <<EOF
+package(default_visibility = ['//visibility:public'])
+constraint_setting(name = 'setting1')
+constraint_value(name = 'value1', constraint_setting = ':setting1')
+constraint_value(name = 'value2', constraint_setting = ':setting1')
+
+constraint_setting(name = 'setting2')
+constraint_value(name = 'value3', constraint_setting = ':setting2')
+constraint_value(name = 'value4', constraint_setting = ':setting2')
+
+platform(
+    name = 'platform1_3',
+    constraint_values = [':value1', ':value3'],
+    visibility = ['//visibility:public'])
+platform(
+    name = 'platform1_4',
+    constraint_values = [':value1', ':value4'],
+    visibility = ['//visibility:public'])
+platform(
+    name = 'platform2_3',
+    constraint_values = [':value2', ':value3'],
+    visibility = ['//visibility:public'])
+platform(
+    name = 'platform2_4',
+    constraint_values = [':value2', ':value4'],
+    visibility = ['//visibility:public'])
+EOF
+
+  # Add a rule with default execution constraints.
+  mkdir -p demo
+  cat >> demo/rule.bzl <<EOF
+def _impl(ctx):
+  return []
+
+sample_rule = rule(
+  implementation = _impl,
+  attrs = {},
+  exec_compatible_with = [
+    '//platforms:value2',
+  ],
+  toolchains = ['//toolchain:test_toolchain'],
+  execution_platform_constraints_allowed = True,
+)
+EOF
+
+  # Use the new rule.
+  cat >> demo/BUILD <<EOF
+load(':rule.bzl', 'sample_rule')
+
+sample_rule(
+  name = 'use',
+  exec_compatible_with = [
+    '//platforms:value4',
+  ],
+)
+EOF
+
+  # Build the target, using debug messages to verify the correct platform was selected.
+  bazel build \
+    --extra_execution_platforms=//platforms:all \
+    --toolchain_resolution_debug \
+    //demo:use &> $TEST_log || fail "Build failed"
+    expect_log "Selected execution platform //platforms:platform2_4"
+}
+
 run_suite "toolchain tests"