Shrink worker pool if memory pressure too high.

`WorkerLifecycleManager` catches workers in idle state too rarely, therefore I implemented strategy to shrink worker pools if the memory pressure is high even after evicting of idle workers.

This strategy sets workers as `doomed` and on the next return to worker pool invalidates them and shrinks worker pool.

To enable this strategy set flag `experimental_shrink_worker_pool`.

This strategy is tested and could decrease worker memory usage twice without almost any wall time degradation.

PiperOrigin-RevId: 522302353
Change-Id: I0c1d09cb3d70c711862dc87ca486bd5c978f0813

Author: Googler

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

@@ -157,7 +157,6 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/util:ps_info_collector",
         "//third_party:auto_value",
         "//third_party:guava",
-        "//third_party:jsr305",
     ],
 )
 
@@ -202,6 +201,7 @@ java_library(
         ":worker",
         ":worker_key",
         "//third_party:apache_commons_pool2",
+        "//third_party:guava",
     ],
 )
 

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

@@ -15,6 +15,8 @@
 
 import com.google.common.base.Throwables;
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
@@ -29,6 +31,14 @@
 @ThreadSafe
 final class SimpleWorkerPool extends GenericKeyedObjectPool<WorkerKey, Worker> {
 
+  /**
+   * The subtrahend for maximal toal number of object per key. Unfortunately
+   * GenericKeyedObjectPoolConfig doesn't support different number of objects per key, so we need to
+   * use this adhoc variable to manage that. We need this variable to automatically adjusting pool
+   * size per worker key.
+   */
+  private Map<WorkerKey, Integer> shrunkBy = new HashMap<>();
+
   public SimpleWorkerPool(WorkerFactory factory, int max) {
     super(factory, makeConfig(max));
   }
@@ -77,12 +87,38 @@ public Worker borrowObject(WorkerKey key) throws IOException, InterruptedExcepti
   public void invalidateObject(WorkerKey key, Worker obj) throws InterruptedException {
     try {
       super.invalidateObject(key, obj);
+      if (obj.isDoomed()) {
+        updateShrunkBy(key);
+      }
     } catch (Throwable t) {
       Throwables.propagateIfPossible(t, InterruptedException.class);
       throw new RuntimeException("unexpected", t);
     }
   }
 
+  @Override
+  public void returnObject(WorkerKey key, Worker obj) {
+    super.returnObject(key, obj);
+    if (obj.isDoomed()) {
+      updateShrunkBy(key);
+    }
+  }
+
+  public int getMaxTotalPerKey(WorkerKey key) {
+    return getMaxTotalPerKey() - shrunkBy.getOrDefault(key, 0);
+  }
+
+  private synchronized void updateShrunkBy(WorkerKey workerKey) {
+    int currentValue = shrunkBy.getOrDefault(workerKey, 0);
+    if (getMaxTotalPerKey() - currentValue > 1) {
+      shrunkBy.put(workerKey, currentValue + 1);
+    }
+  }
+
+  void clearShrunkBy() {
+    shrunkBy = new HashMap<>();
+  }
+
   /**
    * Our own configuration class for the {@code SimpleWorkerPool} that correctly implements {@code
    * equals()} and {@code hashCode()}.

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

@@ -39,18 +39,24 @@ public abstract class Worker {
   protected final int workerId;
   /** The path of the log file for this worker. */
   protected final Path logFile;
+  /**
+   * Indicated that worker should be destroyed after usage. If worker doomed, then after its
+   * desctruction we automacically shrink the pool size for its worker key.
+   */
+  protected boolean doomed;
 
   public Worker(WorkerKey workerKey, int workerId, Path logFile) {
     this.workerKey = workerKey;
     this.workerId = workerId;
     this.logFile = logFile;
+    this.doomed = false;
   }
 
   /**
    * Returns a unique id for this worker. This is used to distinguish different worker processes in
    * logs and messages.
    */
-  int getWorkerId() {
+  public int getWorkerId() {
     return this.workerId;
   }
 
@@ -64,6 +70,14 @@ public WorkerKey getWorkerKey() {
     return workerKey;
   }
 
+  public boolean isDoomed() {
+    return doomed;
+  }
+
+  void setDoomed(boolean doomed) {
+    this.doomed = doomed;
+  }
+
   HashCode getWorkerFilesCombinedHash() {
     return workerKey.getWorkerFilesCombinedHash();
   }

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

@@ -144,6 +144,9 @@ public void destroyObject(WorkerKey key, PooledObject<Worker> p) {
   @Override
   public boolean validateObject(WorkerKey key, PooledObject<Worker> p) {
     Worker worker = p.getObject();
+    if (worker.isDoomed()) {
+      return false;
+    }
     Optional<Integer> exitValue = worker.getExitValue();
     if (exitValue.isPresent()) {
       if (worker.diedUnexpectedly()) {

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

@@ -17,6 +17,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.GoogleLogger;
+import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.Reporter;
 import java.time.Duration;
 import java.util.Comparator;
 import java.util.HashSet;
@@ -41,12 +43,17 @@ final class WorkerLifecycleManager extends Thread {
   private boolean isWorking = false;
   private final WorkerPool workerPool;
   private final WorkerOptions options;
+  private Reporter reporter;
 
   public WorkerLifecycleManager(WorkerPool workerPool, WorkerOptions options) {
     this.workerPool = workerPool;
     this.options = options;
   }
 
+  public void setReporter(Reporter reporter) {
+    this.reporter = reporter;
+  }
+
   @Override
   public void run() {
     if (options.totalWorkerMemoryLimitMb == 0) {
@@ -89,7 +96,6 @@ void evictWorkers(ImmutableList<WorkerMetric> workerMetrics) throws InterruptedE
 
     int workerMemeoryUsage =
         workerMetrics.stream()
-            .filter(metric -> metric.getWorkerStat() != null)
             .mapToInt(metric -> metric.getWorkerStat().getUsedMemoryInKB() / 1000)
             .sum();
 
@@ -101,18 +107,67 @@ void evictWorkers(ImmutableList<WorkerMetric> workerMetrics) throws InterruptedE
         collectEvictionCandidates(
             workerMetrics, options.totalWorkerMemoryLimitMb, workerMemeoryUsage);
 
-    if (candidates.isEmpty()) {
-      return;
+    String msg =
+        String.format("Going to evict %d workers with ids: %s", candidates.size(), candidates);
+    logger.atInfo().log("%s", msg);
+    if (reporter != null) {
+      reporter.handle(Event.info(msg));
     }
 
-    logger.atInfo().log("Going to evict %d workers with ids: %s", candidates.size(), candidates);
+    ImmutableSet<Integer> evictedWorkers = evictCandidates(workerPool, candidates);
 
-    evictCandidates(workerPool, candidates);
+    msg =
+        String.format(
+            "Total evicted idle workers %d. With ids: %s", evictedWorkers.size(), evictedWorkers);
+    logger.atInfo().log("%s", msg);
+    if (reporter != null) {
+      reporter.handle(Event.info(msg));
+    }
+
+    if (options.shrinkWorkerPool) {
+      List<WorkerMetric> notEvictedWorkerMetrics =
+          workerMetrics.stream()
+              .filter(
+                  metric ->
+                      !evictedWorkers.containsAll(metric.getWorkerProperties().getWorkerIds()))
+              .collect(Collectors.toList());
+
+      int notEvictedWorkerMemeoryUsage =
+          notEvictedWorkerMetrics.stream()
+              .mapToInt(metric -> metric.getWorkerStat().getUsedMemoryInKB() / 1000)
+              .sum();
+
+      if (notEvictedWorkerMemeoryUsage <= options.totalWorkerMemoryLimitMb) {
+        return;
+      }
+
+      postponeInvalidation(notEvictedWorkerMetrics, notEvictedWorkerMemeoryUsage);
+    }
   }
 
-  private void evictCandidates(WorkerPool pool, ImmutableSet<Integer> candidates)
-      throws InterruptedException {
-    pool.evictWithPolicy(new CandidateEvictionPolicy(candidates));
+  private void postponeInvalidation(
+      List<WorkerMetric> workerMetrics, int notEvictedWorkerMemeoryUsage) {
+    ImmutableSet<Integer> potentialCandidates =
+        getCandidates(
+            workerMetrics, options.totalWorkerMemoryLimitMb, notEvictedWorkerMemeoryUsage);
+
+    String msg = String.format("New doomed workers candidates %s", potentialCandidates);
+    logger.atInfo().log("%s", msg);
+    if (reporter != null) {
+      reporter.handle(Event.info(msg));
+    }
+    workerPool.setDoomedWorkers(potentialCandidates);
+  }
+
+  /**
+   * Applies eviction police for candidates. Returns the worker ids of evicted workers. We don't
+   * garantee that every candidate is going to be evicted. Returns worker ids of evicted workers.
+   */
+  private static ImmutableSet<Integer> evictCandidates(
+      WorkerPool pool, ImmutableSet<Integer> candidates) throws InterruptedException {
+    CandidateEvictionPolicy policy = new CandidateEvictionPolicy(candidates);
+    pool.evictWithPolicy(policy);
+    return policy.getEvictedWorkers();
   }
 
   /** Collects worker candidates to evict. Choses workers with the largest memory consumption. */
@@ -124,27 +179,27 @@ ImmutableSet<Integer> collectEvictionCandidates(
 
     List<WorkerMetric> idleWorkerMetrics =
         workerMetrics.stream()
-            .filter(
-                metric ->
-                    metric.getWorkerStat() != null
-                        && idleWorkers.containsAll(metric.getWorkerProperties().getWorkerIds()))
+            .filter(metric -> idleWorkers.containsAll(metric.getWorkerProperties().getWorkerIds()))
             .collect(Collectors.toList());
 
-    if (idleWorkerMetrics.size() != idleWorkers.size()) {
-      logger.atInfo().log(
-          "Difference between idle workers and idle worker metrics is %d",
-          idleWorkers.size() - idleWorkerMetrics.size());
-    }
+    return getCandidates(idleWorkerMetrics, memoryLimitMb, workerMemeoryUsageMb);
+  }
 
-    idleWorkerMetrics.sort(new MemoryComparator());
+  /**
+   * Chooses the worker ids of workers with the most usage of memory. Selects workers until total
+   * memory usage is less than memoryLimitMb.
+   */
+  private static ImmutableSet<Integer> getCandidates(
+      List<WorkerMetric> workerMetrics, int memoryLimitMb, int usedMemoryMb) {
 
+    workerMetrics.sort(new MemoryComparator());
     ImmutableSet.Builder<Integer> candidates = ImmutableSet.builder();
     int freeMemoryMb = 0;
-    for (WorkerMetric metric : idleWorkerMetrics) {
+    for (WorkerMetric metric : workerMetrics) {
       candidates.addAll(metric.getWorkerProperties().getWorkerIds());
       freeMemoryMb += metric.getWorkerStat().getUsedMemoryInKB() / 1000;
 
-      if (workerMemeoryUsageMb - freeMemoryMb <= memoryLimitMb) {
+      if (usedMemoryMb - freeMemoryMb <= memoryLimitMb) {
         break;
       }
     }
@@ -183,27 +238,33 @@ public Set<Integer> getWorkerIds() {
 
   /** Eviction policy for WorkerPool. Evict all idle workers, which were passed in constructor. */
   private static class CandidateEvictionPolicy implements EvictionPolicy<Worker> {
-    private final ImmutableSet<Integer> workerIds;
+    private final ImmutableSet<Integer> workerCandidates;
+    private final Set<Integer> evictedWorkers;
 
-    public CandidateEvictionPolicy(ImmutableSet<Integer> workerIds) {
-      this.workerIds = workerIds;
+    public CandidateEvictionPolicy(ImmutableSet<Integer> workerCandidates) {
+      this.workerCandidates = workerCandidates;
+      this.evictedWorkers = new HashSet<>();
     }
 
     @Override
     public boolean evict(EvictionConfig config, PooledObject<Worker> underTest, int idleCount) {
       int workerId = underTest.getObject().getWorkerId();
-      if (workerIds.contains(workerId)) {
+      if (workerCandidates.contains(workerId)) {
+        evictedWorkers.add(workerId);
         logger.atInfo().log(
             "Evicting worker %d with mnemonic %s",
             workerId, underTest.getObject().getWorkerKey().getMnemonic());
         return true;
       }
-
       return false;
     }
+
+    public ImmutableSet<Integer> getEvictedWorkers() {
+      return ImmutableSet.copyOf(evictedWorkers);
+    }
   }
 
-  /** Compare workers memory in descending order. */
+  /** Compare worker metrics by memory consupmtion in descending order. */
   private static class MemoryComparator implements Comparator<WorkerMetric> {
     @Override
     public int compare(WorkerMetric m1, WorkerMetric m2) {

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

@@ -18,7 +18,6 @@
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.BuildMetrics.WorkerMetrics;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.BuildMetrics.WorkerMetrics.WorkerStats;
 import java.time.Instant;
-import javax.annotation.Nullable;
 
 /**
  * Contains data about worker statistics during execution. This class contains data for {@link
@@ -29,13 +28,12 @@ public abstract class WorkerMetric {
 
   public abstract WorkerProperties getWorkerProperties();
 
-  @Nullable
   public abstract WorkerStat getWorkerStat();
 
   public abstract boolean isMeasurable();
 
   public static WorkerMetric create(
-      WorkerProperties workerProperties, @Nullable WorkerStat workerStat, boolean isMeasurable) {
+      WorkerProperties workerProperties, WorkerStat workerStat, boolean isMeasurable) {
     return new AutoValue_WorkerMetric(workerProperties, workerStat, isMeasurable);
   }
 

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

@@ -162,12 +162,18 @@ public void buildStarting(BuildStartingEvent event) {
       workerPool = new WorkerPoolImpl(newConfig);
       // If workerPool is restarted then we should recreate metrics.
       WorkerMetricsCollector.instance().clear();
+
+      // Start collecting after a pool is defined
+      workerLifecycleManager = new WorkerLifecycleManager(workerPool, options);
+      if (options.workerVerbose) {
+        workerLifecycleManager.setReporter(env.getReporter());
+      }
+      workerLifecycleManager.setDaemon(true);
+      workerLifecycleManager.start();
     }
 
-    // Start collecting after a pool is defined
-    workerLifecycleManager = new WorkerLifecycleManager(workerPool, options);
-    workerLifecycleManager.setDaemon(true);
-    workerLifecycleManager.start();
+    // Clean doomed workers on the beginning of a build.
+    workerPool.clearDoomedWorkers();
   }
 
   @Override

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

@@ -221,4 +221,14 @@ public String getTypeDescription() {
       effectTags = {OptionEffectTag.EXECUTION},
       help = "If enabled, workers are run in a hardened sandbox, if the implementation allows it.")
   public boolean sandboxHardening;
+
+  @Option(
+      name = "experimental_shrink_worker_pool",
+      defaultValue = "false",
+      documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
+      effectTags = {OptionEffectTag.EXECUTION, OptionEffectTag.HOST_MACHINE_RESOURCE_OPTIMIZATIONS},
+      help =
+          "If enabled, could shrink worker pool if worker memory pressure is high. This flag works"
+              + " only when flag experimental_total_worker_memory_limit_mb is enabled.")
+  public boolean shrinkWorkerPool;
 }

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 package com.google.devtools.build.lib.worker;
 
+import com.google.common.collect.ImmutableSet;
 import java.io.IOException;
 import org.apache.commons.pool2.impl.EvictionPolicy;
 
@@ -43,5 +44,9 @@ public interface WorkerPool {
 
   void invalidateObject(WorkerKey key, Worker obj) throws InterruptedException;
 
+  void setDoomedWorkers(ImmutableSet<Integer> workerIds);
+
+  void clearDoomedWorkers();
+
   void close();
 }