diff --git a/src/main/java/com/google/devtools/build/lib/actions/BUILD b/src/main/java/com/google/devtools/build/lib/actions/BUILD index b859a4ac7beaab..425b90d4bbdf40 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/BUILD +++ b/src/main/java/com/google/devtools/build/lib/actions/BUILD @@ -124,6 +124,7 @@ java_library( "//third_party:flogger", "//third_party:guava", "//third_party:jsr305", + "//third_party:rxjava3", "//third_party/protobuf:protobuf_java", ], ) @@ -305,6 +306,7 @@ java_library( "//src/main/java/com/google/devtools/common/options", "//third_party:flogger", "//third_party:guava", + "//third_party:rxjava3", ], ) diff --git a/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java b/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java index dcdd4358b497c4..e029c5ac28122f 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java @@ -152,6 +152,35 @@ public String parseIfMatches(String tag) throws ValidationException { return null; }); + /** How many extra resources an action requires for execution. */ + public static final ParseableRequirement RESOURCES = + ParseableRequirement.create( + "resources::", + Pattern.compile("resources:(.+:.+)"), + s -> { + Preconditions.checkNotNull(s); + + int splitIndex = s.indexOf(":"); + String resourceCount = s.substring(splitIndex+1); + int value; + try { + value = Integer.parseInt(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')"; + } + + if (value < 1) { + return "can't be zero or negative"; + } + + return null; + }); + /** If an action supports running in persistent worker mode. */ public static final String SUPPORTS_WORKERS = "supports-workers"; diff --git a/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java b/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java index e8a57f6266fcce..fcb3f1b8415309 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java @@ -25,9 +25,13 @@ import com.google.devtools.build.lib.util.OS; import com.google.devtools.build.lib.util.Pair; import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.LinkedList; import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.concurrent.CountDownLatch; /** @@ -137,6 +141,10 @@ public static ResourceManager instance() { // definition in the ResourceSet class. private double usedRam; + // Used amount of extra resources. Corresponds to the extra resource + // definition in the ResourceSet class. + private Map usedExtraResources; + // Used local test count. Corresponds to the local test count definition in the ResourceSet class. private int usedLocalTestCount; @@ -145,6 +153,7 @@ public static ResourceManager instance() { private ResourceManager() { requestList = new LinkedList<>(); + usedExtraResources = new HashMap<>(); } @VisibleForTesting public static ResourceManager instanceForTestingOnly() { @@ -158,6 +167,7 @@ private ResourceManager() { public synchronized void resetResourceUsage() { usedCpu = 0; usedRam = 0; + usedExtraResources = new HashMap<>(); usedLocalTestCount = 0; for (Pair request : requestList) { // CountDownLatch can be set only to 0 or 1. @@ -177,6 +187,7 @@ public synchronized void setAvailableResources(ResourceSet resources) { ResourceSet.create( staticResources.getMemoryMb(), staticResources.getCpuUsage(), + staticResources.getExtraResourceUsage(), staticResources.getLocalTestCount()); processWaitingThreads(); } @@ -265,6 +276,17 @@ ResourceHandle tryAcquire(ActionExecutionMetadata owner, ResourceSet resources) private void incrementResources(ResourceSet resources) { usedCpu += resources.getCpuUsage(); usedRam += resources.getMemoryMb(); + + resources.getExtraResourceUsage().entrySet().forEach( + resource -> { + String key = (String)resource.getKey(); + float value = resource.getValue(); + if (usedExtraResources.containsKey(key)) { + value += (float)usedExtraResources.get(key); + } + usedExtraResources.put(key, value); + }); + usedLocalTestCount += resources.getLocalTestCount(); } @@ -272,7 +294,7 @@ private void incrementResources(ResourceSet resources) { * Return true if any resources have been claimed through this manager. */ public synchronized boolean inUse() { - return usedCpu != 0.0 || usedRam != 0.0 || usedLocalTestCount != 0 || !requestList.isEmpty(); + return usedCpu != 0.0 || usedRam != 0.0 || !usedExtraResources.isEmpty() || usedLocalTestCount != 0 || !requestList.isEmpty(); } @@ -323,6 +345,13 @@ private synchronized CountDownLatch acquire(ResourceSet resources) { private synchronized boolean release(ResourceSet resources) { usedCpu -= resources.getCpuUsage(); usedRam -= resources.getMemoryMb(); + + for (Map.Entry 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 @@ -334,6 +363,20 @@ private synchronized boolean release(ResourceSet resources) { if (usedRam < epsilon) { usedRam = 0; } + + Set 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 (String key : toRemove) { + usedExtraResources.remove(key); + } + if (!requestList.isEmpty()) { processWaitingThreads(); return true; @@ -361,12 +404,29 @@ private synchronized void processWaitingThreads() { } } + /** + * Return true iff all requested extra resources are considered to be available. + */ + private boolean areExtraResourcesAvailable(ResourceSet resources) { + for (Map.Entry 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 epsilon = 0.0001f; // Account for possible rounding errors. + if (requested != 0.0 && used != 0.0 && requested + used > available + epsilon) { + return false; + } + } + return true; + } + // Method will return true if all requested resources are considered to be available. private boolean areResourcesAvailable(ResourceSet resources) { Preconditions.checkNotNull(availableResources); // Comparison below is robust, since any calculation errors will be fixed // by the release() method. - if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0) { + if (usedCpu == 0.0 && usedRam == 0.0 && usedExtraResources.isEmpty() && usedLocalTestCount == 0) { return true; } // Use only MIN_NECESSARY_???_RATIO of the resource value to check for @@ -410,9 +470,10 @@ private boolean areResourcesAvailable(ResourceSet resources) { // 3) If used resource amount is less than total available resource amount. boolean cpuIsAvailable = cpu == 0.0 || usedCpu == 0.0 || usedCpu + cpu <= availableCpu; boolean ramIsAvailable = ram == 0.0 || usedRam == 0.0 || ram <= remainingRam; + boolean extraResourcesIsAvailable = areExtraResourcesAvailable(resources); boolean localTestCountIsAvailable = localTestCount == 0 || usedLocalTestCount == 0 || usedLocalTestCount + localTestCount <= availableLocalTestCount; - return cpuIsAvailable && ramIsAvailable && localTestCountIsAvailable; + return cpuIsAvailable && ramIsAvailable && extraResourcesIsAvailable && localTestCountIsAvailable; } @VisibleForTesting diff --git a/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java b/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java index b62f74d416907a..583aafd87e7914 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java @@ -14,13 +14,16 @@ package com.google.devtools.build.lib.actions; +import com.google.common.base.Joiner; import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableMap; import com.google.common.primitives.Doubles; import com.google.devtools.build.lib.collect.nestedset.NestedSet; import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable; import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec; 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; @@ -43,12 +46,23 @@ public class ResourceSet implements ResourceSetOrBuilder { /** The number of CPUs, or fractions thereof. */ private final double cpuUsage; + /** + * Map of extra resources (for example: GPUs, embedded boards, ...) mapping + * name of the resource to a value. + */ + private final ImmutableMap extraResourceUsage; + /** The number of local tests. */ private final int localTestCount; private ResourceSet(double memoryMb, double cpuUsage, int localTestCount) { + this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount); + } + + private ResourceSet(double memoryMb, double cpuUsage, @NonNull ImmutableMap extraResourceUsage, int localTestCount) { this.memoryMb = memoryMb; this.cpuUsage = cpuUsage; + this.extraResourceUsage = extraResourceUsage; this.localTestCount = localTestCount; } @@ -74,18 +88,28 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) { } /** - * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, ioUsage, and + * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, and + * localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or + * {@link #createWithLocalTestCount(int)}. Use this method primarily when constructing + * ResourceSets that represent available resources. + */ + public static ResourceSet create(double memoryMb, double cpuUsage, int localTestCount) { + return ResourceSet.create(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount); + } + + /** + * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, extraResources, and * localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or * {@link #createWithLocalTestCount(int)}. Use this method primarily when constructing * ResourceSets that represent available resources. */ @AutoCodec.Instantiator public static ResourceSet create( - double memoryMb, double cpuUsage, int localTestCount) { - if (memoryMb == 0 && cpuUsage == 0 && localTestCount == 0) { + double memoryMb, double cpuUsage, ImmutableMap extraResourceUsage, int localTestCount) { + if (memoryMb == 0 && cpuUsage == 0 && extraResourceUsage.size() == 0 && localTestCount == 0) { return ZERO; } - return new ResourceSet(memoryMb, cpuUsage, localTestCount); + return new ResourceSet(memoryMb, cpuUsage, extraResourceUsage, localTestCount); } /** Returns the amount of real memory (resident set size) used in MB. */ @@ -105,6 +129,10 @@ public double getCpuUsage() { return cpuUsage; } + public ImmutableMap getExtraResourceUsage() { + return extraResourceUsage; + } + /** Returns the local test count used. */ public int getLocalTestCount() { return localTestCount; @@ -115,6 +143,7 @@ public String toString() { return "Resources: \n" + "Memory: " + memoryMb + "M\n" + "CPU: " + cpuUsage + "\n" + + Joiner.on("\n").withKeyValueSeparator(": ").join(extraResourceUsage.entrySet()) + "Local tests: " + localTestCount + "\n"; } diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java index 0c97167cfc8390..5f9bc5023d6f3f 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java +++ b/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java @@ -31,6 +31,7 @@ import com.google.devtools.build.lib.server.FailureDetails.FailureDetail; import com.google.devtools.build.lib.server.FailureDetails.TestAction; import com.google.devtools.build.lib.server.FailureDetails.TestAction.Code; +import java.util.HashMap; import java.util.List; import java.util.Map; @@ -162,32 +163,22 @@ public boolean isPersistentTestRunner() { return isPersistentTestRunner; } - public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs) - throws UserExecException { - if (usingLocalTestJobs) { - return LOCAL_TEST_JOBS_BASED_RESOURCES; - } - + private double getLocalCpuResourceUsage(Label label) throws UserExecException { ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size); - // Tests can override their CPU reservation with a "cpu:" tag. - ResourceSet testResourcesFromTag = null; + double cpuCount = -1.0; for (String tag : executionInfo.keySet()) { try { String cpus = ExecutionRequirements.CPU.parseIfMatches(tag); if (cpus != null) { - if (testResourcesFromTag != null) { + if (cpuCount != -1.0) { String message = String.format( "%s has more than one '%s' tag, but duplicate tags aren't allowed", label, ExecutionRequirements.CPU.userFriendlyName()); throw new UserExecException(createFailureDetail(message, Code.DUPLICATE_CPU_TAGS)); } - testResourcesFromTag = - ResourceSet.create( - testResourcesFromSize.getMemoryMb(), - Float.parseFloat(cpus), - testResourcesFromSize.getLocalTestCount()); + cpuCount = Float.parseFloat(cpus); } } catch (ValidationException e) { String message = @@ -200,8 +191,54 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs throw new UserExecException(createFailureDetail(message, Code.INVALID_CPU_TAG)); } } + return cpuCount != -1.0 ? cpuCount : testResourcesFromSize.getCpuUsage(); + } - return testResourcesFromTag != null ? testResourcesFromTag : testResourcesFromSize; + private ImmutableMap getLocalExtraResourceUsage(Label label) throws UserExecException { + // Tests can specify requirements for extra resources using "resources::" tag. + Map extraResources = new HashMap<>(); + for (String tag : executionInfo.keySet()) { + try { + String extras = ExecutionRequirements.RESOURCES.parseIfMatches(tag); + if (extras != null) { + int splitIndex = extras.indexOf(":"); + String resourceName = extras.substring(0, splitIndex); + String resourceCount = extras.substring(splitIndex+1); + if (extraResources.get(resourceName) != null) { + String message = + String.format( + "%s has more than one '%s' tag, but duplicate tags aren't allowed", + label, ExecutionRequirements.RESOURCES.userFriendlyName()); + throw new UserExecException(createFailureDetail(message, Code.DUPLICATE_CPU_TAGS)); + } + extraResources.put(resourceName, Float.parseFloat(resourceCount)); + } + } catch (ValidationException e) { + String message = + String.format( + "%s has a '%s' tag, but its value '%s' didn't pass validation: %s", + label, + ExecutionRequirements.RESOURCES.userFriendlyName(), + e.getTagValue(), + e.getMessage()); + throw new UserExecException(createFailureDetail(message, Code.INVALID_CPU_TAG)); + } + } + return ImmutableMap.copyOf(extraResources); + } + + public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs) + throws UserExecException { + if (usingLocalTestJobs) { + return LOCAL_TEST_JOBS_BASED_RESOURCES; + } + + ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size); + return ResourceSet.create( + testResourcesFromSize.getMemoryMb(), + getLocalCpuResourceUsage(label), + getLocalExtraResourceUsage(label), + testResourcesFromSize.getLocalTestCount()); } private static FailureDetail createFailureDetail(String message, Code detailedCode) { diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java index a6b182555c07f6..17d10ebc1b15ae 100644 --- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java +++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java @@ -107,6 +107,7 @@ import java.util.HashSet; import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; @@ -883,10 +884,14 @@ public static void configureResourceManager(ResourceManager resourceMgr, BuildRe resources = ResourceSet.createWithRamCpu(options.localRamResources, options.localCpuResources); resourceMgr.setUseLocalMemoryEstimate(options.localMemoryEstimate); + ImmutableMap extraResources = options.localExtraResources.stream().collect( + ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (v1,v2) -> v1)); + resourceMgr.setAvailableResources( ResourceSet.create( resources.getMemoryMb(), resources.getCpuUsage(), + extraResources, request.getExecutionOptions().usingLocalTestJobs() ? request.getExecutionOptions().localTestJobs : Integer.MAX_VALUE)); diff --git a/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java b/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java index 387c04b9987e96..f135c00ba8ae98 100644 --- a/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java +++ b/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java @@ -302,8 +302,7 @@ public boolean shouldMaterializeParamFiles() { + "an integer, or \"HOST_CPUS\", optionally followed by [-|*] " + "(eg. HOST_CPUS*.5 to use half the available CPU cores)." + "By default, (\"HOST_CPUS\"), Bazel will query system configuration to estimate " - + "number of CPU cores available for the locally executed build actions. " - + "Note: This is a no-op if --local_resources is set.", + + "number of CPU cores available for the locally executed build actions.", converter = CpuResourceConverter.class) public float localCpuResources; @@ -318,11 +317,27 @@ public boolean shouldMaterializeParamFiles() { + "(eg. HOST_RAM*.5 to use half the available RAM)." + "By default, (\"HOST_RAM*.67\"), Bazel will query system configuration to estimate " + "amount of RAM available for the locally executed build actions and will use 67% " - + "of available RAM. " - + "Note: This is a no-op if --local_resources is set.", + + "of available RAM.", converter = RamResourceConverter.class) public float localRamResources; + @Option( + name = "local_extra_resources", + defaultValue = "null", + documentationCategory = OptionDocumentationCategory.UNCATEGORIZED, + effectTags = {OptionEffectTag.UNKNOWN}, + allowMultiple = true, + 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 " + + "by using a tag of the \"resources::\" format. " + + "Available CPU, RAM and test job resources cannot be set with this flag.", + converter = Converters.StringToFloatAssignmentConverter.class) + public List> localExtraResources; + @Option( name = "experimental_local_memory_estimate", defaultValue = "false", diff --git a/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java b/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java index 0947cf5cf5e8dc..b1960e878dc85c 100644 --- a/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java +++ b/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java @@ -56,7 +56,8 @@ private static boolean legalExecInfoKeys(String tag) { || tag.startsWith("disable-") || tag.startsWith("cpu:") || tag.equals(ExecutionRequirements.LOCAL) - || tag.equals(ExecutionRequirements.WORKER_KEY_MNEMONIC); + || tag.equals(ExecutionRequirements.WORKER_KEY_MNEMONIC) + || tag.startsWith("resources:"); } private TargetUtils() {} // Uninstantiable. diff --git a/src/main/java/com/google/devtools/common/options/Converters.java b/src/main/java/com/google/devtools/common/options/Converters.java index 8563e9cff816a6..76ee1169f2077e 100644 --- a/src/main/java/com/google/devtools/common/options/Converters.java +++ b/src/main/java/com/google/devtools/common/options/Converters.java @@ -469,6 +469,24 @@ public String getTypeDescription() { } } + /** + * A converter for for assignments from a string value to a float value. + */ + public static class StringToFloatAssignmentConverter implements Converter> { + private static final AssignmentConverter baseConverter = new AssignmentConverter(); + + @Override + public Map.Entry convert(String input) throws OptionsParsingException, NumberFormatException { + Map.Entry stringEntry = baseConverter.convert(input); + return Maps.immutableEntry(stringEntry.getKey(), Float.parseFloat(stringEntry.getValue())); + } + + @Override + public String getTypeDescription() { + return "a named float, 'name=value'"; + } + } + /** * A converter for command line flag aliases. It does additional validation on the name and value * of the assignment to ensure they conform to the naming limitations. diff --git a/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java b/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java index bac53ea4ede5c5..813fdc5dd5ad11 100644 --- a/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java +++ b/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java @@ -51,7 +51,14 @@ public class ResourceManagerTest { @Before public final void configureResourceManager() throws Exception { rm.setAvailableResources( - ResourceSet.create(/*memoryMb=*/ 1000, /*cpuUsage=*/ 1, /* localTestCount= */ 2)); + ResourceSet.create( + /*memoryMb=*/ 1000, + /*cpuUsage=*/ 1, + /*extraResourceUsage*/ ImmutableMap.of( + "gpu", 2.0f, + "fancyresource", 1.5f + ), + /* localTestCount= */ 2)); counter = new AtomicInteger(0); sync = new CyclicBarrier(2); sync2 = new CyclicBarrier(2); @@ -63,14 +70,27 @@ private ResourceHandle acquire(double ram, double cpu, int tests) return rm.acquireResources(resourceOwner, ResourceSet.create(ram, cpu, tests)); } + private ResourceHandle acquire(double ram, double cpu, int tests, ImmutableMap extraResources) + throws InterruptedException { + return rm.acquireResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests)); + } + private ResourceHandle acquireNonblocking(double ram, double cpu, int tests) { return rm.tryAcquire(resourceOwner, ResourceSet.create(ram, cpu, tests)); } + private ResourceHandle acquireNonblocking(double ram, double cpu, int tests, ImmutableMap extraResources) { + return rm.tryAcquire(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests)); + } + private void release(double ram, double cpu, int tests) { rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests)); } + private void release(double ram, double cpu, int tests, ImmutableMap extraResources) { + rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests)); + } + private void validate(int count) { assertThat(counter.incrementAndGet()).isEqualTo(count); } @@ -104,6 +124,13 @@ public void testOverBudgetRequests() throws Exception { assertThat(rm.inUse()).isTrue(); release(0, 0, bigTests); assertThat(rm.inUse()).isFalse(); + + // Ditto, for extra resources (even if they don't exist in the available resource set): + ImmutableMap bigExtraResources = ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f, "nonexisting", 10.0f); + acquire(0, 0, 0, bigExtraResources); + assertThat(rm.inUse()).isTrue(); + release(0, 0, 0, bigExtraResources); + assertThat(rm.inUse()).isFalse(); } @Test @@ -183,11 +210,25 @@ public void testThatTestsCannotBeOverallocated() throws Exception { thread1.joinAndAssertState(10000); } + @Test + public void testThatExtraResourcesCannotBeOverallocated() throws Exception { + assertThat(rm.inUse()).isFalse(); + + // Given a partially acquired extra resources: + acquire(0, 0, 1, ImmutableMap.of("gpu", 1.0f)); + + // 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()); + thread1.start(); + thread1.joinAndAssertState(10000); + } + @Test public void testHasResources() throws Exception { assertThat(rm.inUse()).isFalse(); assertThat(rm.threadHasResources()).isFalse(); - acquire(1, 0.1, 1); + acquire(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f)); assertThat(rm.threadHasResources()).isTrue(); // We have resources in this thread - make sure other threads @@ -208,11 +249,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)); + assertThat(rm.threadHasResources()).isTrue(); + release(0, 0, 0, ImmutableMap.of("gpu", 1.0f)); + assertThat(rm.threadHasResources()).isFalse(); }); thread1.start(); thread1.joinAndAssertState(10000); - release(1, 0.1, 1); + release(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f)); assertThat(rm.threadHasResources()).isFalse(); assertThat(rm.inUse()).isFalse(); }