Add asyncClose to the ResourceHandle.

It helps to release the resources from  worker async finishing.
Also added check of workers availability in case of unused resources.

PiperOrigin-RevId: 465038659
Change-Id: Ib1e6e1fc5b4159ff2235c47d306850fd4e879658

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

@@ -309,23 +309,22 @@ public boolean threadHasResources() {
     return threadLocked.get();
   }
 
-  void releaseResources(ActionExecutionMetadata owner, ResourceSet resources)
-      throws IOException, InterruptedException {
-    releaseResources(owner, resources, /* worker= */ null);
-    return;
-  }
-
   /**
-   * Releases previously requested resource =.
+   * Releases previously requested resource.
    *
    * <p>NB! This method must be thread-safe!
+   *
+   * @param owner action metadata, which resources should ve released
+   * @param resources resources should be released
+   * @param worker the worker, which used during execution
+   * @throws java.io.IOException if could not return worker to the workerPool
    */
-  @VisibleForTesting
   void releaseResources(
       ActionExecutionMetadata owner, ResourceSet resources, @Nullable Worker worker)
       throws IOException, InterruptedException {
     Preconditions.checkNotNull(
         resources, "releaseResources called with resources == NULL during %s", owner);
+
     Preconditions.checkState(
         threadHasResources(), "releaseResources without resource lock during %s", owner);
 
@@ -343,6 +342,15 @@ void releaseResources(
     }
   }
 
+  // TODO (b/241066751) find better way to change resource ownership
+  public void releaseResourceOwnership() {
+    threadLocked.set(false);
+  }
+
+  public void acquireResourceOwnership() {
+    threadLocked.set(true);
+  }
+
   /**
    * Returns the pair of worker and latch. Worker should be null if there is no workerKey in
    * resources. The latch isn't null if we could not acquire the resources right now and need to
@@ -436,7 +444,17 @@ 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) {
+
+    WorkerKey workerKey = resources.getWorkerKey();
+    int availableWorkers = 0;
+    int activeWorkers = 0;
+    if (workerKey != null) {
+      availableWorkers = this.workerPool.getMaxTotalPerKey(workerKey);
+      activeWorkers = this.workerPool.getNumActive(workerKey);
+    }
+    boolean workerIsAvailable = workerKey == null || activeWorkers < availableWorkers;
+
+    if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0 && workerIsAvailable) {
       return true;
     }
     // Use only MIN_NECESSARY_???_RATIO of the resource value to check for
@@ -454,14 +472,6 @@ private boolean areResourcesAvailable(ResourceSet resources) {
 
     double remainingRam = availableRam - usedRam;
 
-    WorkerKey workerKey = resources.getWorkerKey();
-    int availableWorkers = 0;
-    int activeWorkers = 0;
-    if (workerKey != null) {
-      availableWorkers = this.workerPool.getMaxTotalPerKey(workerKey);
-      activeWorkers = this.workerPool.getNumActive(workerKey);
-    }
-
     // Resources are considered available if any one of the conditions below is true:
     // 1) If resource is not requested at all, it is available.
     // 2) If resource is not used at the moment, it is considered to be
@@ -471,9 +481,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 localTestCountIsAvailable = localTestCount == 0 || usedLocalTestCount == 0
-        || usedLocalTestCount + localTestCount <= availableLocalTestCount;
-    boolean workerIsAvailable = workerKey == null || activeWorkers < availableWorkers;
+    boolean localTestCountIsAvailable =
+        localTestCount == 0
+            || usedLocalTestCount == 0
+            || usedLocalTestCount + localTestCount <= availableLocalTestCount;
     return cpuIsAvailable && ramIsAvailable && localTestCountIsAvailable && workerIsAvailable;
   }
 

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

@@ -602,6 +602,9 @@ private WorkResponse executeRequest(
             handle,
             workerAsResource);
         workerOwner.setWorker(null);
+        if (workerAsResource) {
+          resourceManager.releaseResourceOwnership();
+        }
       } else if (!context.speculating()) {
         // Non-sandboxed workers interrupted outside of dynamic execution can only mean that
         // the user interrupted the build, and we don't want to delay finishing. Instead we
@@ -662,6 +665,10 @@ private void finishWorkAsync(
     Thread reaper =
         new Thread(
             () -> {
+              if (workerAsResource) {
+                resourceManager.acquireResourceOwnership();
+              }
+
               Worker w = worker;
               try {
                 if (canCancel) {
@@ -692,10 +699,18 @@ private void finishWorkAsync(
                 }
               } finally {
                 if (w != null) {
-                  try {
-                    workers.returnObject(key, w);
-                  } catch (IllegalStateException e3) {
-                    // The worker already not part of the pool
+                  if (workerAsResource) {
+                    try {
+                      resourceHandle.close();
+                    } catch (IOException | InterruptedException | IllegalStateException e) {
+                      // Error while returning worker to the pool. Could not do anythinng.
+                    }
+                  } else {
+                    try {
+                      workers.returnObject(key, w);
+                    } catch (IllegalStateException e3) {
+                      // The worker already not part of the pool
+                    }
                   }
                 }
               }

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

@@ -115,7 +115,7 @@ private ResourceHandle acquire(double ram, double cpu, int tests, String mnemoni
   }
 
   private void release(double ram, double cpu, int tests) throws IOException, InterruptedException {
-    rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests));
+    rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker=*/ null);
   }
 
   private void validate(int count) {