commit 55dfe56cc10cef528c36f058553f4d77e5d90d10
author Googler <tanzhengwei@google.com> Fri Apr 19 06:54:26 2024 -0700
committer Copybara-Service <copybara-worker@google.com> Fri Apr 19 06:56:26 2024 -0700
tree 87d0cfbc76129f6023b609233c57d1701c085ac8
parent 6f953cb037558b17e99018d8f5d190c6421c7c21

Implement missing WorkerKeyPool.close() that kills workers when shutting down the pool.

PiperOrigin-RevId: 626352634
Change-Id: I3f8a816b4f21619fe1e159a228e374ef1765a8a0

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

diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java
index e79ef12..a08ef53 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerPoolImpl.java
@@ -357,7 +357,18 @@
           key.getMnemonic(), key.hashCode());
     }
 
-    private void close() {}
+    // Destroys all workers created in this pool.
+    private synchronized void close() {
+      for (Worker worker : idleQueue) {
+        factory.destroyWorker(worker.getWorkerKey(), worker);
+      }
+      for (Worker worker : activeSet) {
+        logger.atInfo().log(
+            "Interrupting and shutting down active worker %s (id %d) due to pool shutdown",
+            key.getMnemonic(), worker.getWorkerId());
+        factory.destroyWorker(worker.getWorkerKey(), worker);
+      }
+    }
   }
 
   /**

src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java

diff --git a/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java b/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java
index 73940af..3ac4e63 100644
--- a/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java
+++ b/src/test/java/com/google/devtools/build/lib/worker/WorkerPoolTest.java
@@ -401,4 +401,16 @@
     workerPool.reset();
     assertThat(workerPool.getMaxTotalPerKey(workerKey)).isEqualTo(2);
   }
+
+  @Test
+  public void testClose_destroysWorkers() throws Exception {
+    WorkerKey workerKey = createWorkerKey(fileSystem, "mnem", false);
+    Worker worker1 = workerPool.borrowObject(workerKey);
+    Worker worker2 = workerPool.borrowObject(workerKey);
+    workerPool.returnObject(workerKey, worker1);
+    workerPool.returnObject(workerKey, worker2);
+    workerPool.close();
+    verify(factoryMock).destroyWorker(workerKey, worker1);
+    verify(factoryMock).destroyWorker(workerKey, worker2);
+  }
 }