Allow --worker_max_instances to take MnemonicName=value to specify max for each named worker.

RELNOTES: Allow --worker_max_instances to take MnemonicName=value to specify max for each worker.
PiperOrigin-RevId: 195244295

4 files changed, 142 insertions, 58 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/worker/SimpleWorkerPool.java b/src/main/java/com/google/devtools/build/lib/worker/SimpleWorkerPool.java
new file mode 100644
index 0000000000..68e512565e
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/worker/SimpleWorkerPool.java
@@ -0,0 +1,54 @@
+// Copyright 2018 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.worker;
+
+import com.google.common.base.Throwables;
+import java.io.IOException;
+import javax.annotation.concurrent.ThreadSafe;
+import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
+import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
+
+/**
+ * A worker pool that spawns multiple workers and delegates work to them.
+ *
+ * <p>This is useful when the worker cannot handle multiple parallel requests on its own and we need
+ * to pre-fork a couple of them instead.
+ */
+@ThreadSafe
+final class SimpleWorkerPool extends GenericKeyedObjectPool<WorkerKey, Worker> {
+
+  public SimpleWorkerPool(WorkerFactory factory, GenericKeyedObjectPoolConfig config) {
+    super(factory, config);
+  }
+
+  @Override
+  public Worker borrowObject(WorkerKey key) throws IOException, InterruptedException {
+    try {
+      return super.borrowObject(key);
+    } catch (Throwable t) {
+      Throwables.propagateIfPossible(t, IOException.class, InterruptedException.class);
+      throw new RuntimeException("unexpected", t);
+    }
+  }
+
+  @Override
+  public void invalidateObject(WorkerKey key, Worker obj) throws IOException, InterruptedException {
+    try {
+      super.invalidateObject(key, obj);
+    } catch (Throwable t) {
+      Throwables.propagateIfPossible(t, IOException.class, InterruptedException.class);
+      throw new RuntimeException("unexpected", t);
+    }
+  }
+}
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
index 716be4dcbf..c5160d9a0d 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerModule.java
@@ -15,6 +15,7 @@ package com.google.devtools.build.lib.worker;
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.eventbus.Subscribe;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
 import com.google.devtools.build.lib.buildtool.buildevent.BuildCompleteEvent;
@@ -29,16 +30,16 @@ import com.google.devtools.build.lib.runtime.commands.CleanCommand.CleanStarting
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.common.options.OptionsBase;
 import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
-/**
- * A module that adds the WorkerActionContextProvider to the available action context providers.
- */
+/** A module that adds the WorkerActionContextProvider to the available action context providers. */
 public class WorkerModule extends BlazeModule {
   private CommandEnvironment env;
 
   private WorkerFactory workerFactory;
   private WorkerPool workerPool;
-  private WorkerPoolConfig workerPoolConfig;
+  private ImmutableMap<String, Integer> workerPoolConfig;
   private WorkerOptions options;
 
   @Override
@@ -96,7 +97,18 @@ public class WorkerModule extends BlazeModule {
     workerFactory.setReporter(env.getReporter());
     workerFactory.setOptions(options);
 
-    WorkerPoolConfig newConfig = createWorkerPoolConfig(options);
+    // Use a LinkedHashMap instead of an ImmutableMap.Builder to allow duplicates; the last value
+    // passed wins.
+    LinkedHashMap<String, Integer> newConfigBuilder = new LinkedHashMap<>();
+    for (Map.Entry<String, Integer> entry : options.workerMaxInstances) {
+      newConfigBuilder.put(entry.getKey(), entry.getValue());
+    }
+    if (!newConfigBuilder.containsKey("")) {
+      // Empty string gives the number of workers for any type of worker not explicitly specified.
+      // If no value is given, use the default, 4.
+      newConfigBuilder.put("", 4);
+    }
+    ImmutableMap<String, Integer> newConfig = ImmutableMap.copyOf(newConfigBuilder);
 
     // If the config changed compared to the last run, we have to create a new pool.
     if (workerPoolConfig != null && !workerPoolConfig.equals(newConfig)) {
@@ -109,37 +121,6 @@ public class WorkerModule extends BlazeModule {
     }
   }
 
-  private WorkerPoolConfig createWorkerPoolConfig(WorkerOptions options) {
-    WorkerPoolConfig config = new WorkerPoolConfig();
-
-    // It's better to re-use a worker as often as possible and keep it hot, in order to profit
-    // from JIT optimizations as much as possible.
-    config.setLifo(true);
-
-    // Keep a fixed number of workers running per key.
-    config.setMaxIdlePerKey(options.workerMaxInstances);
-    config.setMaxTotalPerKey(options.workerMaxInstances);
-    config.setMinIdlePerKey(options.workerMaxInstances);
-
-    // Don't limit the total number of worker processes, as otherwise the pool might be full of
-    // e.g. Java workers and could never accommodate another request for a different kind of
-    // worker.
-    config.setMaxTotal(-1);
-
-    // Wait for a worker to become ready when a thread needs one.
-    config.setBlockWhenExhausted(true);
-
-    // Always test the liveliness of worker processes.
-    config.setTestOnBorrow(true);
-    config.setTestOnCreate(true);
-    config.setTestOnReturn(true);
-
-    // No eviction of idle workers.
-    config.setTimeBetweenEvictionRunsMillis(-1);
-
-    return config;
-  }
-
   @Override
   public void executorInit(CommandEnvironment env, BuildRequest request, ExecutorBuilder builder) {
     Preconditions.checkNotNull(workerPool);
@@ -149,8 +130,7 @@ public class WorkerModule extends BlazeModule {
 
   @Subscribe
   public void buildComplete(BuildCompleteEvent event) {
-    if (options != null
-        && options.workerQuitAfterBuild) {
+    if (options != null && options.workerQuitAfterBuild) {
       shutdownPool("Build completed, shutting down worker pool...");
     }
   }
@@ -163,9 +143,7 @@ public class WorkerModule extends BlazeModule {
     shutdownPool("Build interrupted, shutting down worker pool...");
   }
 
-  /**
-   * Shuts down the worker pool and sets {#code workerPool} to null.
-   */
+  /** Shuts down the worker pool and sets {#code workerPool} to null. */
   private void shutdownPool(String reason) {
     Preconditions.checkArgument(!reason.isEmpty());
 
diff --git a/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
index 6055818680..c55a139ffe 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerOptions.java
@@ -46,14 +46,17 @@ public class WorkerOptions extends OptionsBase {
 
   @Option(
     name = "worker_max_instances",
-    defaultValue = "4",
+    converter = Converters.NamedIntegersConverter.class,
+    defaultValue = "",
     documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
     effectTags = {OptionEffectTag.UNKNOWN},
     help =
         "How many instances of a worker process (like the persistent Java compiler) may be "
-            + "launched if you use the 'worker' strategy."
+            + "launched if you use the 'worker' strategy. May be specified as [name=value] to "
+            + "give a different value per worker mnemonic.",
+    allowMultiple = true
   )
-  public int workerMaxInstances;
+  public List<Map.Entry<String, Integer>> workerMaxInstances;
 
   @Option(
       name = "high_priority_workers",
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
index 9b3afbbdda..6fe749e74f 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerPool.java
@@ -14,12 +14,13 @@
 package com.google.devtools.build.lib.worker;
 
 import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import java.io.IOException;
+import java.util.HashSet;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.concurrent.ThreadSafe;
-import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
-import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
 
 /**
  * A worker pool that spawns multiple workers and delegates work to them.
@@ -28,21 +29,66 @@ import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
  * to pre-fork a couple of them instead.
  */
 @ThreadSafe
-final class WorkerPool extends GenericKeyedObjectPool<WorkerKey, Worker> {
+final class WorkerPool {
   private final AtomicInteger highPriorityWorkersInUse = new AtomicInteger(0);
   private final ImmutableSet<String> highPriorityWorkerMnemonics;
+  private final ImmutableMap<String, Integer> config;
+  private final ImmutableMap<Integer, SimpleWorkerPool> pools;
 
   /**
    * @param factory worker factory
-   * @param config pool configuration
+   * @param config pool configuration; max number of workers per worker mnemonic; the empty string
+   *     key specifies the default maximum
    * @param highPriorityWorkers mnemonics of high priority workers
    */
   public WorkerPool(
-      WorkerFactory factory,
-      GenericKeyedObjectPoolConfig config,
-      Iterable<String> highPriorityWorkers) {
-    super(factory, config);
+      WorkerFactory factory, Map<String, Integer> config, Iterable<String> highPriorityWorkers) {
     highPriorityWorkerMnemonics = ImmutableSet.copyOf(highPriorityWorkers);
+    this.config = ImmutableMap.copyOf(config);
+    ImmutableMap.Builder<Integer, SimpleWorkerPool> poolsBuilder = ImmutableMap.builder();
+    for (Integer max : new HashSet<>(config.values())) {
+      poolsBuilder.put(max, new SimpleWorkerPool(factory, makeConfig(max)));
+    }
+    pools = poolsBuilder.build();
+  }
+
+  private WorkerPoolConfig makeConfig(int max) {
+    WorkerPoolConfig config = new WorkerPoolConfig();
+
+    // It's better to re-use a worker as often as possible and keep it hot, in order to profit
+    // from JIT optimizations as much as possible.
+    config.setLifo(true);
+
+    // Keep a fixed number of workers running per key.
+    config.setMaxIdlePerKey(max);
+    config.setMaxTotalPerKey(max);
+    config.setMinIdlePerKey(max);
+
+    // Don't limit the total number of worker processes, as otherwise the pool might be full of
+    // e.g. Java workers and could never accommodate another request for a different kind of
+    // worker.
+    config.setMaxTotal(-1);
+
+    // Wait for a worker to become ready when a thread needs one.
+    config.setBlockWhenExhausted(true);
+
+    // Always test the liveliness of worker processes.
+    config.setTestOnBorrow(true);
+    config.setTestOnCreate(true);
+    config.setTestOnReturn(true);
+
+    // No eviction of idle workers.
+    config.setTimeBetweenEvictionRunsMillis(-1);
+
+    return config;
+  }
+
+  private SimpleWorkerPool getPool(WorkerKey key) {
+    Integer max = config.get(key.getMnemonic());
+    if (max == null) {
+      max = config.get("");
+    }
+    return pools.get(max);
   }
 
   /**
@@ -51,11 +97,10 @@ final class WorkerPool extends GenericKeyedObjectPool<WorkerKey, Worker> {
    * @param key worker key
    * @return a worker
    */
-  @Override
   public Worker borrowObject(WorkerKey key) throws IOException, InterruptedException {
     Worker result;
     try {
-      result = super.borrowObject(key);
+      result = getPool(key).borrowObject(key);
     } catch (Throwable t) {
       Throwables.propagateIfPossible(t, IOException.class, InterruptedException.class);
       throw new RuntimeException("unexpected", t);
@@ -75,21 +120,19 @@ final class WorkerPool extends GenericKeyedObjectPool<WorkerKey, Worker> {
     return result;
   }
 
-  @Override
   public void returnObject(WorkerKey key, Worker obj) {
     if (highPriorityWorkerMnemonics.contains(key.getMnemonic())) {
       decrementHighPriorityWorkerCount();
     }
-    super.returnObject(key, obj);
+    getPool(key).returnObject(key, obj);
   }
 
-  @Override
   public void invalidateObject(WorkerKey key, Worker obj) throws IOException, InterruptedException {
     if (highPriorityWorkerMnemonics.contains(key.getMnemonic())) {
       decrementHighPriorityWorkerCount();
     }
     try {
-      super.invalidateObject(key, obj);
+      getPool(key).invalidateObject(key, obj);
     } catch (Throwable t) {
       Throwables.propagateIfPossible(t, IOException.class, InterruptedException.class);
       throw new RuntimeException("unexpected", t);
@@ -118,4 +161,10 @@ final class WorkerPool extends GenericKeyedObjectPool<WorkerKey, Worker> {
       }
     }
   }
+
+  public void close() {
+    for (SimpleWorkerPool pool : pools.values()) {
+      pool.close();
+    }
+  }
 }