Compose flaky attribute with --flaky_test_attempts

site/en/docs/user-manual.md

@@ -1521,10 +1521,15 @@ or total number of passed tests. Tests that fail all allowed attempts are
 considered to be failed.
 
 By default (when this option is not specified, or when it is set to
-default), only a single attempt is allowed for regular tests, and
-3 for test rules with the `flaky` attribute set. You can specify
-an integer value to override the maximum limit of test attempts. Bazel allows
-a maximum of 10 test attempts in order to prevent abuse of the system.
+default), non-flaky targets (`flaky = 0`) run once. Flaky targets run
+`1 + flaky` times. For example, `flaky = 1` runs twice and `flaky = 2`
+runs three times.
+
+When set to an integer, that value is the total number of attempts for
+all matching targets, including non-flaky ones, and overrides the flaky
+attribute. You can specify an integer value globally or per target regex.
+Bazel allows a maximum of 10 test attempts in order to prevent abuse of
+the system.
 
 #### `--runs_per_test={{ "<var>" }}[regex@]number{{ "</var>" }}` {:#runs-per-test}
 

src/main/java/com/google/devtools/build/docgen/templates/attributes/test/flaky.html

@@ -1,13 +1,17 @@
-<p>Boolean; <a href="#configurable-attributes">nonconfigurable</a>;
-  default is <code>False</code></p>
+<p>Integer; <a href="#configurable-attributes">nonconfigurable</a>;
+  default is <code>0</code>. Legacy boolean values are accepted (<code>True</code>
+  is treated as <code>1</code>, <code>False</code> as <code>0</code>).</p>
 
 <p>
-Marks test as flaky.
+Declares how many additional retry attempts this flaky test should receive
+when <code>--flaky_test_attempts</code> is <code>default</code>. A value of
+<code>0</code> means the test is not flaky and runs once unless a numeric
+<code>--flaky_test_attempts</code> value overrides all targets. For example,
+with the default flag, <code>flaky = 1</code> runs the test twice and
+<code>flaky = 2</code> runs it three times.
 </p>
 
 <p>
-If set, executes the test up to three times, marking it as failed only if it
-fails each time. By default, this attribute is set to False and the test is
-executed only once. Note, that use of this attribute is generally discouraged -
-tests should pass reliably when their assertions are upheld.
+Note that use of this attribute is generally discouraged - tests should pass
+reliably when their assertions are upheld.
 </p>

src/main/java/com/google/devtools/build/lib/analysis/BaseRuleClasses.java

@@ -230,8 +230,8 @@ public RuleClass build(RuleClass.Builder builder, RuleDefinitionEnvironment env)
                   .nonconfigurable("policy decision: should be consistent across configurations")
                   .value(TIMEOUT_DEFAULT))
           .add(
-              attr("flaky", BOOLEAN)
-                  .value(false)
+              attr("flaky", Type.FLAKY_TEST_RETRIES)
+                  .value(StarlarkInt.of(0))
                   .taggable()
                   .nonconfigurable("policy decision: should be consistent across configurations"))
           .add(attr("shard_count", INTEGER).value(StarlarkInt.of(-1)))

src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java

@@ -238,8 +238,8 @@ public static RuleClass getTestBaseRule(RuleDefinitionEnvironment env) {
                     .nonconfigurable("policy decision: should be consistent across configurations")
                     .value(TIMEOUT_DEFAULT))
             .add(
-                attr("flaky", BOOLEAN)
-                    .value(false)
+                attr("flaky", Type.FLAKY_TEST_RETRIES)
+                    .value(StarlarkInt.of(0))
                     .taggable()
                     .nonconfigurable("taggable - called in Rule.getRuleTags"))
             .add(attr("shard_count", INTEGER).value(StarlarkInt.of(-1)))

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

@@ -276,40 +276,44 @@ private static void addRunUnderArgs(TestRunnerAction testAction, List<String> ar
   }
 
   /**
-   * Returns the number of attempts specific test action can be retried.
+   * Returns the number of attempts for the given test action.
    *
-   * <p>For rules with "flaky = 1" attribute, this method will return 3 unless --flaky_test_attempts
-   * option is given and specifies another value.
+   * <p>When {@code --flaky_test_attempts} resolves to a numeric value for this target, that value
+   * is the total number of attempts for all tests (including {@code flaky = 0}). When the flag is
+   * {@code default}, non-flaky targets run once and flaky targets run {@code 1 + flaky} times.
    */
   @VisibleForTesting /* protected */
   public int getTestAttempts(TestRunnerAction action) {
-    return action.getTestProperties().isFlaky()
-        ? getTestAttemptsForFlakyTest(action)
-        : getTestAttempts(action, /* defaultTestAttempts= */ 1);
+    return computeTestAttempts(
+        action.getTestProperties().getFlakyRetries(),
+        executionOptions,
+        action.getOwner().getLabel());
   }
 
-  private int getTestAttempts(TestRunnerAction action, int defaultTestAttempts) {
-    Label testLabel = action.getOwner().getLabel();
-    return getTestAttemptsPerLabel(executionOptions, testLabel, defaultTestAttempts);
+  @VisibleForTesting
+  static int computeTestAttempts(int flakyRetries, ExecutionOptions options, Label label) {
+    String flagValue = resolveFlakyTestAttemptsFlagValue(options, label);
+    if (!"default".equals(flagValue)) {
+      return Math.min(MAX_TEST_ATTEMPTS, Integer.parseInt(flagValue));
+    }
+    if (flakyRetries == 0) {
+      return 1;
+    }
+    return Math.min(MAX_TEST_ATTEMPTS, 1 + flakyRetries);
   }
 
-  public int getTestAttemptsForFlakyTest(TestRunnerAction action) {
-    return getTestAttempts(action, /* defaultTestAttempts= */ 3);
-  }
+  private static final int MAX_TEST_ATTEMPTS = 10;
 
-  private static int getTestAttemptsPerLabel(
-      ExecutionOptions options, Label label, int defaultTestAttempts) {
-    // Check from the last provided, so that the last option provided takes precedence.
+  private static String resolveFlakyTestAttemptsFlagValue(ExecutionOptions options, Label label) {
+    if (options.testAttempts == null || options.testAttempts.isEmpty()) {
+      return "default";
+    }
     for (PerLabelOptions perLabelAttempts : Lists.reverse(options.testAttempts)) {
       if (perLabelAttempts.isIncluded(label)) {
-        String attempts = Iterables.getOnlyElement(perLabelAttempts.getOptions());
-        if ("default".equals(attempts)) {
-          return defaultTestAttempts;
-        }
-        return Integer.parseInt(attempts);
+        return Iterables.getOnlyElement(perLabelAttempts.getOptions());
       }
     }
-    return defaultTestAttempts;
+    return "default";
   }
 
   /**

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

@@ -68,7 +68,7 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
   private final TestTimeout timeout;
   private final List<String> tags;
   private final boolean isRemotable;
-  private final boolean isFlaky;
+  private final int flakyRetries;
   private final boolean isExternal;
   private final String language;
   private final ImmutableMap<String, String> executionInfo;
@@ -87,7 +87,8 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
     tags = ruleContext.attributes().get("tags", Types.STRING_LIST);
 
     // We need to use method on ruleConfiguredTarget to perform validation.
-    isFlaky = ruleContext.attributes().get("flaky", Type.BOOLEAN);
+    flakyRetries =
+        ruleContext.attributes().get("flaky", Type.FLAKY_TEST_RETRIES).toIntUnchecked();
     isExternal = TargetUtils.isExternalTestRule(rule);
 
     Map<String, String> executionInfo = Maps.newLinkedHashMap();
@@ -148,8 +149,12 @@ public boolean isRemotable() {
     return isRemotable;
   }
 
+  public int getFlakyRetries() {
+    return flakyRetries;
+  }
+
   public boolean isFlaky() {
-    return isFlaky;
+    return flakyRetries > 0;
   }
 
   public boolean isExternal() {

src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java

@@ -222,16 +222,16 @@ public boolean shouldMaterializeParamFiles() {
           "Each test will be retried up to the specified number of times in case of any test"
               + " failure. Tests that required more than one attempt to pass are marked as 'FLAKY'"
               + " in the test summary. Normally the value specified is just an integer or the"
-              + " string 'default'. If an integer, then all tests will be run up to N times. If"
-              + " 'default', then only a single test attempt will be made for regular tests and"
-              + " three for tests marked explicitly as flaky by their rule (flaky=1 attribute)."
-              + " Alternate syntax: regex_filter@flaky_test_attempts. Where flaky_test_attempts is"
-              + " as above and regex_filter stands for a list of include and exclude regular"
-              + " expression patterns (Also see --runs_per_test). Example:"
-              + " --flaky_test_attempts=//foo/.*,-//foo/bar/.*@3 deflakes all tests in //foo/"
-              + " except those under foo/bar three times. This option can be passed multiple"
-              + " times. The most recently passed argument that matches takes precedence. If"
-              + " nothing matches, behavior is as if 'default' above.")
+              + " string 'default'. If an integer, then that value is the total number of attempts"
+              + " for all matching targets, overriding the flaky attribute. If 'default', non-flaky"
+              + " targets (flaky = 0) run once and flaky targets run 1 + flaky times. Alternate"
+              + " syntax: regex_filter@flaky_test_attempts. Where flaky_test_attempts is as above"
+              + " and regex_filter stands for a list of include and exclude regular expression"
+              + " patterns (Also see --runs_per_test). Example:"
+              + " --flaky_test_attempts=//foo/.*,-//foo/bar/.*@3 runs all tests in //foo/ except"
+              + " foo/bar up to three times. This option can be passed multiple times. The most"
+              + " recently passed argument that matches takes precedence. If nothing matches,"
+              + " behavior is as if 'default' above.")
   public List<PerLabelOptions> testAttempts;
 
   @Option(

src/main/java/com/google/devtools/build/lib/packages/ProtoUtils.java

@@ -30,6 +30,7 @@
 import static com.google.devtools.build.lib.packages.BuildType.OUTPUT_LIST;
 import static com.google.devtools.build.lib.packages.BuildType.TRISTATE;
 import static com.google.devtools.build.lib.packages.Type.BOOLEAN;
+import static com.google.devtools.build.lib.packages.Type.FLAKY_TEST_RETRIES;
 import static com.google.devtools.build.lib.packages.Type.INTEGER;
 import static com.google.devtools.build.lib.packages.Type.STRING;
 import static com.google.devtools.build.lib.packages.Type.STRING_NO_INTERN;
@@ -48,6 +49,7 @@ public class ProtoUtils {
   private static final ImmutableMap<Type<?>, Discriminator> TYPE_MAP =
       new ImmutableMap.Builder<Type<?>, Discriminator>()
           .put(INTEGER, Discriminator.INTEGER)
+          .put(FLAKY_TEST_RETRIES, Discriminator.INTEGER)
           .put(DISTRIBUTIONS, Discriminator.DISTRIBUTION_SET)
           .put(LABEL, Discriminator.LABEL)
           // NODEP_LABEL attributes are not really strings. This is implemented

src/main/java/com/google/devtools/build/lib/packages/RuleClass.java

@@ -684,7 +684,7 @@ public String toString() {
             attr("tags", Types.STRING_LIST).build(),
             attr("size", Type.STRING).build(),
             attr("timeout", Type.STRING).build(),
-            attr("flaky", Type.BOOLEAN).build(),
+            attr("flaky", Type.FLAKY_TEST_RETRIES).build(),
             attr("shard_count", Type.INTEGER).build(),
             attr("local", Type.BOOLEAN).build());
 

src/main/java/com/google/devtools/build/lib/packages/Type.java

@@ -260,6 +260,13 @@ public Set<String> toTagSet(Object value, String name) {
   /** The type of a Starlark integer in the signed 32-bit range. */
   @SerializationConstant public static final Type<StarlarkInt> INTEGER = new IntegerType();
 
+  /**
+   * The type of the {@code flaky} test attribute: a non-negative int counting additional retry
+   * attempts, composed with {@code --flaky_test_attempts}. Also accepts legacy boolean values.
+   */
+  @SerializationConstant
+  public static final Type<StarlarkInt> FLAKY_TEST_RETRIES = new FlakyTestRetriesType();
+
   /** The type of a string which interns the instance with String#intern. */
   @SerializationConstant
   public static final Type<String> STRING = new StringType(/* internString= */ true);
@@ -386,6 +393,60 @@ public StarlarkInt concat(Iterable<StarlarkInt> elements) {
     }
   }
 
+  private static final class FlakyTestRetriesType extends Type<StarlarkInt> {
+    @Override
+    public StarlarkInt cast(Object value) {
+      return (StarlarkInt) value;
+    }
+
+    @Override
+    public StarlarkInt getDefaultValue() {
+      return StarlarkInt.of(0);
+    }
+
+    @Override
+    public void visitLabels(LabelVisitor visitor, StarlarkInt value, @Nullable Attribute context) {}
+
+    @Override
+    public String toString() {
+      return "int";
+    }
+
+    @Override
+    public StarlarkInt convert(Object x, Object what, LabelConverter labelConverter)
+        throws ConversionException {
+      if (x instanceof Boolean) {
+        return StarlarkInt.of((Boolean) x ? 1 : 0);
+      }
+      if (x instanceof StarlarkInt) {
+        StarlarkInt i = (StarlarkInt) x;
+        try {
+          if (i.toIntUnchecked() < 0) {
+            throw new ConversionException("non-negative int", x, what);
+          }
+        } catch (
+            @SuppressWarnings("UnusedException")
+            IllegalArgumentException ex) {
+          String prefix = what != null ? ("for " + what + ", ") : "";
+          throw new ConversionException(
+              String.format("%sgot %s, want value in signed 32-bit range", prefix, i));
+        }
+        return i;
+      }
+      throw new ConversionException(this, x, what);
+    }
+
+    @Override
+    public Set<String> toTagSet(Object value, String name) {
+      if (value == null) {
+        String msg = "Illegal tag conversion from null on Attribute " + name + ".";
+        throw new IllegalStateException(msg);
+      }
+      int retries = ((StarlarkInt) value).toIntUnchecked();
+      return retries == 0 ? ImmutableSet.of("noflaky") : ImmutableSet.of("flaky");
+    }
+  }
+
   private static final class BooleanType extends Type<Boolean> {
     @Override
     public Boolean cast(Object value) {

src/test/java/com/google/devtools/build/lib/analysis/BUILD

@@ -79,6 +79,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:build_view",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_configuration",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/build_options",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/per_label_options",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/core_options",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/execution_transition_factory",
         "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment",
@@ -126,6 +127,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/collect/nestedset",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/exec:bin_tools",
+        "//src/main/java/com/google/devtools/build/lib/exec:execution_options",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/packages:configured_attribute_mapper",
         "//src/main/java/com/google/devtools/build/lib/packages:exec_group",

src/test/java/com/google/devtools/build/lib/analysis/test/TestActionBuilderTest.java

@@ -267,28 +267,52 @@ public void testFlakyAttributeValidation() throws Exception {
             srcs = ["a.sh"],
             flaky = 1,
         )
+
+        foo_test(
+            name = "flaky_true_test",
+            srcs = ["a.sh"],
+            flaky = True,
+        )
+
+        foo_test(
+            name = "very_flaky_test",
+            srcs = ["a.sh"],
+            flaky = 2,
+        )
         """);
     Artifact testStatus = Iterables.getOnlyElement(getTestStatusArtifacts("//flaky:good_test"));
     assertThat(testStatus).isNotNull();
     TestRunnerAction action = (TestRunnerAction) getGeneratingAction(testStatus);
+    assertThat(action.getTestProperties().getFlakyRetries()).isEqualTo(0);
     assertThat(action.getTestProperties().isFlaky()).isFalse();
 
     testStatus = Iterables.getOnlyElement(getTestStatusArtifacts("//flaky:flaky_test"));
     assertThat(testStatus).isNotNull();
     action = (TestRunnerAction) getGeneratingAction(testStatus);
+    assertThat(action.getTestProperties().getFlakyRetries()).isEqualTo(1);
     assertThat(action.getTestProperties().isFlaky()).isTrue();
+
+    testStatus = Iterables.getOnlyElement(getTestStatusArtifacts("//flaky:flaky_true_test"));
+    assertThat(testStatus).isNotNull();
+    action = (TestRunnerAction) getGeneratingAction(testStatus);
+    assertThat(action.getTestProperties().getFlakyRetries()).isEqualTo(1);
+
+    testStatus = Iterables.getOnlyElement(getTestStatusArtifacts("//flaky:very_flaky_test"));
+    assertThat(testStatus).isNotNull();
+    action = (TestRunnerAction) getGeneratingAction(testStatus);
+    assertThat(action.getTestProperties().getFlakyRetries()).isEqualTo(2);
   }
 
   @Test
-  public void testIllegalBooleanFlakySetting() throws Exception {
+  public void testIllegalNegativeFlakySetting() throws Exception {
     checkError(
         "flaky",
         "bad_test",
-        "expected one of [False, True, 0, 1]",
+        "expected non-negative int",
         "load('//test_defs:foo_test.bzl', 'foo_test')",
         "foo_test(name = 'bad_test',",
         "        srcs = ['a.sh'],",
-        "        flaky = 2)");
+        "        flaky = -1)");
   }
 
   @Test