Actively kill off still-active workers when stopping work on interrupt.

larsrc-google · copybara-github · commit 308bce36cba4 · 2021-01-18T12:13:52.000-08:00
Due to limits of GenericKeyedObjectPool, we can't selectively kill the active workers, but must instead clear the entire pool.

RELNOTES: None.
PiperOrigin-RevId: 352441326
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
@@ -203,12 +203,20 @@ public void buildComplete(BuildCompleteEvent event) {
     }
   }
 
-  // Kill workers on Ctrl-C to quickly end the interrupted build.
-  // TODO(philwo) - make sure that this actually *kills* the workers and not just politely waits
-  // for them to finish.
+  /**
+   * Stops any workers that are still executing.
+   *
+   * <p>This currently kills off some amount of workers, losing the warmed-up state.
+   * TODO(b/119701157): Cancel running workers instead (requires some way to reach each worker).
+   */
   @Subscribe
   public void buildInterrupted(BuildInterruptedEvent event) {
-    shutdownPool("Build interrupted, shutting down worker pool...");
+    if (workerPool != null) {
+      if ((options != null && options.workerVerbose)) {
+        env.getReporter().handle(Event.info("Build interrupted, stopping active workers..."));
+      }
+      workerPool.stopWork();
+    }
   }
 
   /** Shuts down the worker pool and sets {#code workerPool} to null. */
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java
@@ -193,4 +193,24 @@ public void close() {
     workerPools.values().forEach(GenericKeyedObjectPool::close);
     multiplexPools.values().forEach(GenericKeyedObjectPool::close);
   }
+
+  /** Stops any ongoing work in the worker pools. This may entail killing the worker processes. */
+  public void stopWork() {
+    workerPools
+        .values()
+        .forEach(
+            pool -> {
+              if (pool.getNumActive() > 0) {
+                pool.clear();
+              }
+            });
+    multiplexPools
+        .values()
+        .forEach(
+            pool -> {
+              if (pool.getNumActive() > 0) {
+                pool.clear();
+              }
+            });
+  }
 }
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
@@ -238,4 +238,50 @@ public void testBorrow_twoHiPrioBlocks() throws Exception {
     verify(factoryMock, times(2)).makeObject(workerKey1);
     verify(factoryMock, times(1)).makeObject(workerKey2);
   }
+
+  @Test
+  public void testStopWork_activePoolsStopped() throws Exception {
+    WorkerPool pool =
+        new WorkerPool(
+            factoryMock,
+            // Have to declare the mnemonics, or they all fall into the default SimpleWorkerPool.
+            ImmutableMap.of("mnem1", 2, "mnem2", 2),
+            ImmutableMap.of("mnem2", 2, "mnem3", 2),
+            Lists.newArrayList());
+    WorkerKey singleKey1 = createWorkerKey(fileSystem, "mnem1", false);
+    // These workers get borrowed, then both get destroyed in stopWork because they share mnemonic
+    WorkerKey singleKey1a = createWorkerKey(fileSystem, "mnem1", false, "anArg");
+    pool.borrowObject(singleKey1);
+    Worker worker1a = pool.borrowObject(singleKey1a);
+    pool.returnObject(singleKey1a, worker1a);
+    WorkerKey singleKey2 = createWorkerKey(fileSystem, "mnem2", false);
+    // This worker gets borrowed, then returned, doesn't get destroyed in stopWork
+    Worker worker1 = pool.borrowObject(singleKey2);
+    pool.returnObject(singleKey2, worker1);
+    WorkerKey multiKey1 = createWorkerKey(fileSystem, "mnem2", true);
+    // This worker gets borrowed, then destroyed in stopWork, but separately from the singleplex
+    // worker2 even though they share a mnemonic.
+    pool.borrowObject(multiKey1);
+    WorkerKey multiKey2 = createWorkerKey(fileSystem, "mnem3", true);
+    // This worker gets borrowed, then returned, doesn't get destroyed during stopWork.
+    Worker worker2 = pool.borrowObject(multiKey2);
+    pool.returnObject(multiKey2, worker2);
+    verify(factoryMock, times(1)).makeObject(singleKey1);
+    verify(factoryMock, times(1)).makeObject(singleKey1a);
+    verify(factoryMock, times(1)).makeObject(singleKey2);
+    verify(factoryMock, times(1)).makeObject(multiKey1);
+    verify(factoryMock, times(1)).makeObject(multiKey2);
+    pool.stopWork();
+    pool.borrowObject(singleKey1);
+    pool.borrowObject(singleKey1a);
+    pool.borrowObject(singleKey2);
+    pool.borrowObject(multiKey1);
+    pool.borrowObject(multiKey2);
+    // After stopWork, we had to create new workers for the keys that got their pools destroyed.
+    verify(factoryMock, times(2)).makeObject(singleKey1);
+    verify(factoryMock, times(2)).makeObject(singleKey1a);
+    verify(factoryMock, times(1)).makeObject(singleKey2);
+    verify(factoryMock, times(2)).makeObject(multiKey1);
+    verify(factoryMock, times(1)).makeObject(multiKey2);
+  }
 }
