fixup Implement extra resources for test actions

Addressed comments on #13996
Fixed issues in tests and built and tested with
lowRISC/opentitan#16436

Signed-off-by: Drew Macrae <drewmacrae@google.com>

src/main/java/com/google/devtools/build/lib/actions/BUILD

@@ -128,7 +128,6 @@ java_library(
  "//third_party:flogger",
  "//third_party:guava",
  "//third_party:jsr305",
- "//third_party:rxjava3",
  "//third_party/protobuf:protobuf_java",
  ],
 )
@@ -312,7 +311,6 @@ java_library(
  "//third_party:flogger",
  "//third_party:guava",
  "//third_party:jsr305",
- "//third_party:rxjava3",
  ],
 )
 

src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java

@@ -157,27 +157,22 @@ public String parseIfMatches(String tag) throws ValidationException {
  /** How many extra resources an action requires for execution. */
  public static final ParseableRequirement RESOURCES =
  ParseableRequirement.create(
- "resources:<str>:<int>",
+ "resources:<str>:<float>",
  Pattern.compile("resources:(.+:.+)"),
  s -> {
  Preconditions.checkNotNull(s);
 
  int splitIndex = s.indexOf(":");
  String resourceCount = s.substring(splitIndex+1);
- int value;
+ float value;
  try {
- value = Integer.parseInt(resourceCount);
+ value = Float.parseFloat(resourceCount);
  } catch (NumberFormatException e) {
- return "can't be parsed as an integer";
- }
-
- // De-and-reserialize & compare to only allow canonical integer formats.
- if (!Integer.toString(value).equals(resourceCount)) {
- return "must be in canonical format (e.g. '4' instead of '+04')";
+ return "can't be parsed as a float";
  }
 
- if (value < 1) {
- return "can't be zero or negative";
+ if (value < 0) {
+ return "can't be negative";
  }
 
  return null;

src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java

@@ -32,6 +32,7 @@
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import javax.annotation.Nullable;
@@ -185,10 +186,6 @@ public static ResourceManager instance() {
  /** If set, local-only actions are given priority over dynamically run actions. */
  private boolean prioritizeLocalActions;
 
- private ResourceManager() {
- usedExtraResources = new HashMap<>();
- }
-
  @VisibleForTesting
  public static ResourceManager instanceForTestingOnly() {
  return new ResourceManager();
@@ -380,7 +377,7 @@ public void acquireResourceOwnership() {
  * wait.
  */
  private synchronized LatchWithWorker acquire(ResourceSet resources, ResourcePriority priority)
- throws IOException, InterruptedException {
+ throws IOException, InterruptedException, NoSuchElementException {
  if (areResourcesAvailable(resources)) {
  Worker worker = incrementResources(resources);
  return new LatchWithWorker(/* latch= */ null, worker);
@@ -429,12 +426,6 @@ private synchronized void releaseResourcesOnly(ResourceSet resources) {
  usedCpu -= resources.getCpuUsage();
  usedRam -= resources.getMemoryMb();
 
- for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
- String key = (String)resource.getKey();
- float value = (float)usedExtraResources.get(key) - resource.getValue();
- usedExtraResources.put(key, value);
- }
-
  usedLocalTestCount -= resources.getLocalTestCount();
 
  // TODO(bazel-team): (2010) rounding error can accumulate and value below can end up being
@@ -448,14 +439,14 @@ private synchronized void releaseResourcesOnly(ResourceSet resources) {
  }
 
  Set<String> toRemove = new HashSet<>();
- usedExtraResources.entrySet().forEach(
- resource -> {
-  String key = (String)resource.getKey();
-  float value = (float)usedExtraResources.get(key);
-  if (value < epsilon) {
-  toRemove.add(key);
-  }
-  });
+ for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
+ String key = (String)resource.getKey();
+ float value = (float)usedExtraResources.get(key) - resource.getValue();
+ usedExtraResources.put(key, value);
+ if (value < epsilon) {
+ toRemove.add(key);
+ }
+ }
  for (String key : toRemove) {
  usedExtraResources.remove(key);
  }
@@ -497,15 +488,28 @@ private synchronized void processWaitingThreads(Deque<Pair<ResourceSet, LatchWit
  }
  }
 
+ /**
+ * Throws an exception if requested extra resource isn't being tracked
+ */
+ private void assertExtraResourcesTracked(ResourceSet resources)
+ throws NoSuchElementException {
+ for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
+ String key = (String)resource.getKey();
+ if (!availableResources.getExtraResourceUsage().containsKey(key)) {
+ throw new NoSuchElementException("Resource "+key+" is not tracked in this resource set.");
+ }
+ }
+ }
+
  /**
  * Return true iff all requested extra resources are considered to be available.
  */
- private boolean areExtraResourcesAvailable(ResourceSet resources) {
+ private boolean areExtraResourcesAvailable(ResourceSet resources) throws NoSuchElementException {
  for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
  String key = (String)resource.getKey();
  float used = (float)usedExtraResources.getOrDefault(key, 0f);
  float requested = resource.getValue();
- float available = (float)availableResources.getExtraResourceUsage().getOrDefault(key, 0f);
+ float available = availableResources.getExtraResourceUsage().get(key);
  float epsilon = 0.0001f; // Account for possible rounding errors.
  if (requested != 0.0 && used != 0.0 && requested + used > available + epsilon) {
  return false;
@@ -516,7 +520,7 @@ private boolean areExtraResourcesAvailable(ResourceSet resources) {
 
  // Method will return true if all requested resources are considered to be available.
  @VisibleForTesting
- boolean areResourcesAvailable(ResourceSet resources) {
+ boolean areResourcesAvailable(ResourceSet resources) throws NoSuchElementException {
  Preconditions.checkNotNull(availableResources);
  // Comparison below is robust, since any calculation errors will be fixed
  // by the release() method.
@@ -532,6 +536,10 @@ boolean areResourcesAvailable(ResourceSet resources) {
  workerKey == null
  || (activeWorkers < availableWorkers && workerPool.couldBeBorrowed(workerKey));
 
+ // We test for tracking of extra resources whenever acquired and throw an
+ // exception before acquiring any untracked resource.
+ assertExtraResourcesTracked(resources);
+
  if (usedCpu == 0.0 && usedRam == 0.0 && usedExtraResources.isEmpty() && usedLocalTestCount == 0 && workerIsAvailable) {
  return true;
  }

src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java

@@ -23,7 +23,6 @@
 import com.google.devtools.build.lib.worker.WorkerKey;
 import com.google.devtools.common.options.Converter;
 import com.google.devtools.common.options.OptionsParsingException;
-import io.reactivex.rxjava3.annotations.NonNull;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 import javax.annotation.Nullable;
@@ -58,11 +57,12 @@ public class ResourceSet implements ResourceSetOrBuilder {
  /** The workerKey of used worker. Null if no worker is used. */
  @Nullable private final WorkerKey workerKey;
 
- private ResourceSet(double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
+ private ResourceSet(
+ double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
  this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, workerKey);
  }
 
- private ResourceSet(double memoryMb, double cpuUsage, @NonNull ImmutableMap<String, Float> extraResourceUsage, int localTestCount, @Nullable WorkerKey workerKey) {
+ private ResourceSet(double memoryMb, double cpuUsage, @Nullable ImmutableMap<String, Float> extraResourceUsage, int localTestCount, @Nullable WorkerKey workerKey) {
  this.memoryMb = memoryMb;
  this.cpuUsage = cpuUsage;
  this.extraResourceUsage = extraResourceUsage;
@@ -102,7 +102,7 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) {
  * represent available resources.
  */
  public static ResourceSet create(double memoryMb, double cpuUsage, int localTestCount) {
- return ResourceSet.create(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* wolkerKey= */ null);
+ return ResourceSet.createWithWorkerKey(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* wolkerKey= */ null);
  }
 
  /**
@@ -112,7 +112,11 @@ public static ResourceSet create(double memoryMb, double cpuUsage, int localTest
  * ResourceSets that represent available resources.
  */
  public static ResourceSet create(double memoryMb, double cpuUsage, ImmutableMap<String, Float> extraResourceUsage, int localTestCount) {
- return createWithWorkerKey(memoryMb, cpuUsage, extraResourceUseage, localTestCount, /* workerKey= */ null);
+ return createWithWorkerKey(memoryMb, cpuUsage, extraResourceUsage, localTestCount, /* workerKey= */ null);
+ }
+
+ public static ResourceSet createWithWorkerKey(double memoryMb, double cpuUsage, int localTestCount, WorkerKey workerKey) {
+ return ResourceSet.createWithWorkerKey(memoryMb, cpuUsage, /* extraResourceUsage= */ImmutableMap.of(), localTestCount, workerKey);
  }
 
  public static ResourceSet createWithWorkerKey(

src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java

@@ -939,7 +939,7 @@ public static void configureResourceManager(ResourceManager resourceMgr, BuildRe
  ExecutionOptions options = request.getOptions(ExecutionOptions.class);
  resourceMgr.setPrioritizeLocalActions(options.prioritizeLocalActions);
  ImmutableMap<String, Float> extraResources = options.localExtraResources.stream().collect(
- ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (v1,v2) -> v1));
+ ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (v1,v2) -> v2));
 
  resourceMgr.setAvailableResources(
  ResourceSet.create(

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

@@ -334,11 +334,11 @@ public boolean shouldMaterializeParamFiles() {
  help =
  "Set the number of extra resources available to Bazel. "
  + "Takes in a string-float pair. Can be used multiple times to specify multiple "
- + "types of extra resources. Bazel will limit concurrently running test actions "
- + "based on the available extra resources and the extra resources required "
- + "by the test actions. Tests can declare the amount of extra resources they need "
+ + "types of extra resources. Bazel will limit concurrently running actions "
+ + "based on the available extra resources and the extra resources required. "
+ + "Tests can declare the amount of extra resources they need "
  + "by using a tag of the \"resources:<resoucename>:<amount>\" format. "
- + "Available CPU, RAM and test job resources cannot be set with this flag.",
+ + "Available CPU, RAM and resources cannot be set with this flag.",
  converter = Converters.StringToFloatAssignmentConverter.class)
  public List<Map.Entry<String, Float>> localExtraResources;
 

src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java

@@ -407,6 +407,7 @@ WorkResponse execInWorkerWorkerAsResource(
  ResourceSet.createWithWorkerKey(
  spawn.getLocalResources().getMemoryMb(),
  spawn.getLocalResources().getCpuUsage(),
+ spawn.getLocalResources().getExtraResourceUsage(),
  spawn.getLocalResources().getLocalTestCount(),
  key);
 

src/main/java/com/google/devtools/common/options/Converters.java

@@ -475,7 +475,7 @@ public String getTypeDescription() {
  /**
  * A converter for for assignments from a string value to a float value.
  */
- public static class StringToFloatAssignmentConverter implements Converter<Map.Entry<String, Float>> {
+ public static class StringToFloatAssignmentConverter extends Converter.Contextless<Map.Entry<String, Float>> {
  private static final AssignmentConverter baseConverter = new AssignmentConverter();
 
  @Override

src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java

@@ -46,6 +46,7 @@
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.NoSuchElementException;
 import javax.annotation.Nullable;
 import org.apache.commons.pool2.PooledObject;
 import org.junit.Before;
@@ -114,21 +115,42 @@ private ResourceHandle acquire(double ram, double cpu, int tests)
  return acquire(ram, cpu, tests, ResourcePriority.LOCAL);
  }
 
- private ResourceHandle acquire(double ram, double cpu, int tests, ImmutableMap<String, Float> extraResources, String mnemonic)
+ private ResourceHandle acquire(double ram, double cpu, int tests, String mnemonic)
  throws InterruptedException, IOException {
 
  return rm.acquireResources(
  resourceOwner,
- ResourceSet.createWithWorkerKey(ram, cpu, tests, extraResources, createWorkerKey(mnemonic)),
+ ResourceSet.createWithWorkerKey(ram, cpu, tests, createWorkerKey(mnemonic)),
+ ResourcePriority.LOCAL);
+ }
+
+ private ResourceHandle acquire(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests, ResourcePriority priority)
+ throws InterruptedException, IOException, NoSuchElementException {
+ return rm.acquireResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests), priority);
+ }
+
+ private ResourceHandle acquire(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests)
+ throws InterruptedException, IOException, NoSuchElementException {
+ return acquire(ram, cpu, extraResources, tests, ResourcePriority.LOCAL);
+ }
+
+ private ResourceHandle acquire(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests, String mnemonic)
+ throws InterruptedException, IOException, NoSuchElementException {
+
+ return rm.acquireResources(
+ resourceOwner,
+ ResourceSet.createWithWorkerKey(ram, cpu, extraResources, tests, createWorkerKey(mnemonic)),
  ResourcePriority.LOCAL);
  }
 
  private void release(double ram, double cpu, int tests) throws IOException, InterruptedException {
- rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker=*/ null);
+ rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker= */ null);
  }
 
- private void release(double ram, double cpu, int tests, ImmutableMap<String, Float> extraResources) {
- rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests));
+ private void release(double ram, double cpu, ImmutableMap<String, Float> extraResources, int tests)
+ throws InterruptedException, IOException {
+ rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests),
+ /* worker= */ null);
  }
 
  private void validate(int count) {
@@ -179,11 +201,11 @@ public void testOverBudgetRequests() throws Exception {
  release(0, 0, bigTests);
  assertThat(rm.inUse()).isFalse();
 
- // Ditto, for extra resources (even if they don't exist in the available resource set):
- ImmutableMap<String, Float> bigExtraResources = ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f, "nonexisting", 10.0f);
- acquire(0, 0, 0, bigExtraResources);
+ // Ditto, for extra resources:
+ ImmutableMap<String, Float> bigExtraResources = ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f);
+ acquire(0, 0, bigExtraResources, 0);
  assertThat(rm.inUse()).isTrue();
- release(0, 0, 0, bigExtraResources);
+ release(0, 0, bigExtraResources, 0);
  assertThat(rm.inUse()).isFalse();
  }
 
@@ -271,20 +293,21 @@ public void testThatExtraResourcesCannotBeOverallocated() throws Exception {
  assertThat(rm.inUse()).isFalse();
 
  // Given a partially acquired extra resources:
- acquire(0, 0, 1, ImmutableMap.of("gpu", 1.0f));
+ acquire(0, 0, ImmutableMap.of("gpu", 1.0f), 1);
 
  // When a request for extra resources is made that would overallocate,
  // Then the request fails:
- TestThread thread1 = new TestThread(() -> assertThat(acquireNonblocking(0, 0, 0, ImmutableMap.of("gpu", 1.1f))).isNull());
+ TestThread thread1 = new TestThread(() -> acquire(0, 0, ImmutableMap.of("gpu", 1.1f), 0));
  thread1.start();
- thread1.joinAndAssertState(10000);
+ AssertionError e = assertThrows(AssertionError.class, () -> thread1.joinAndAssertState(1000));
+ assertThat(e).hasCauseThat().hasMessageThat().contains("is still alive");
  }
 
  @Test
  public void testHasResources() throws Exception {
  assertThat(rm.inUse()).isFalse();
  assertThat(rm.threadHasResources()).isFalse();
- acquire(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f));
+ acquire(1, 0.1, ImmutableMap.of("gpu", 1.0f), 1);
  assertThat(rm.threadHasResources()).isTrue();
 
  // We have resources in this thread - make sure other threads
@@ -305,15 +328,15 @@ public void testHasResources() throws Exception {
  assertThat(rm.threadHasResources()).isTrue();
  release(0, 0, 1);
  assertThat(rm.threadHasResources()).isFalse();
- acquire(0, 0, 0, ImmutableMap.of("gpu", 1.0f));
+ acquire(0, 0, ImmutableMap.of("gpu", 1.0f), 0);
  assertThat(rm.threadHasResources()).isTrue();
- release(0, 0, 0, ImmutableMap.of("gpu", 1.0f));
+ release(0, 0, ImmutableMap.of("gpu", 1.0f), 0);
  assertThat(rm.threadHasResources()).isFalse();
  });
  thread1.start();
  thread1.joinAndAssertState(10000);
 
- release(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f));
+ release(1, 0.1, ImmutableMap.of("gpu", 1.0f), 1);
  assertThat(rm.threadHasResources()).isFalse();
  assertThat(rm.inUse()).isFalse();
  }
@@ -704,6 +727,18 @@ private CyclicBarrier startAcquireReleaseThread(ResourcePriority priority) {
  return sync;
  }
 
+ @Test
+ public void testNonexistingResource() throws Exception {
+ // If we try to use nonexisting resource we should return an error
+ TestThread thread1 =
+ new TestThread(
+ () -> {
+ assertThrows(NoSuchElementException.class, () -> acquire(0, 0, ImmutableMap.of("nonexisting", 1.0f), 0));
+ });
+ thread1.start();
+ thread1.joinAndAssertState(1000);
+ }
+
  @Test
  public void testAcquireWithWorker_acquireAndRelease() throws Exception {
  int memory = 100;