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

5 files changed, 211 insertions, 118 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();
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/common/options/Converters.java b/src/main/java/com/google/devtools/common/options/Converters.java
index 35f2da4f93..fb3bbfaef2 100644
--- a/src/main/java/com/google/devtools/common/options/Converters.java
+++ b/src/main/java/com/google/devtools/common/options/Converters.java
@@ -26,10 +26,7 @@ import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.regex.PatternSyntaxException;
 
-/**
- * Some convenient converters used by blaze. Note: These are specific to
- * blaze.
- */
+/** Some convenient converters used by blaze. Note: These are specific to blaze. */
 public final class Converters {
 
   /** Standard converter for booleans. Accepts common shorthands/synonyms. */
@@ -180,9 +177,7 @@ public final class Converters {
     }
   }
 
-  /**
-   * Standard converter for the {@link java.time.Duration} type.
-   */
+  /** Standard converter for the {@link java.time.Duration} type. */
   public static class DurationConverter implements Converter<Duration> {
     private final Pattern durationRegex = Pattern.compile("^([0-9]+)(d|h|m|s|ms)$");
 
@@ -198,7 +193,7 @@ public final class Converters {
       }
       long duration = Long.parseLong(m.group(1));
       String unit = m.group(2);
-      switch(unit) {
+      switch (unit) {
         case "d":
           return Duration.ofDays(duration);
         case "h":
@@ -210,8 +205,8 @@ public final class Converters {
         case "ms":
           return Duration.ofMillis(duration);
         default:
-          throw new IllegalStateException("This must not happen. Did you update the regex without "
-              + "the switch case?");
+          throw new IllegalStateException(
+              "This must not happen. Did you update the regex without the switch case?");
       }
     }
 
@@ -240,14 +235,8 @@ public final class Converters {
           .build();
 
   /**
-   * Join a list of words as in English.  Examples:
-   * "nothing"
-   * "one"
-   * "one or two"
-   * "one and two"
-   * "one, two or three".
-   * "one, two and three".
-   * The toString method of each element is used.
+   * Join a list of words as in English. Examples: "nothing" "one" "one or two" "one and two" "one,
+   * two or three". "one, two and three". The toString method of each element is used.
    */
   static String joinEnglishList(Iterable<?> choices) {
     StringBuilder buf = new StringBuilder();
@@ -261,14 +250,12 @@ public final class Converters {
     return buf.length() == 0 ? "nothing" : buf.toString();
   }
 
-  public static class SeparatedOptionListConverter
-      implements Converter<List<String>> {
+  public static class SeparatedOptionListConverter implements Converter<List<String>> {
 
     private final String separatorDescription;
     private final Splitter splitter;
 
-    protected SeparatedOptionListConverter(char separator,
-                                           String separatorDescription) {
+    protected SeparatedOptionListConverter(char separator, String separatorDescription) {
       this.separatorDescription = separatorDescription;
       this.splitter = Splitter.on(separator);
     }
@@ -284,8 +271,7 @@ public final class Converters {
     }
   }
 
-  public static class CommaSeparatedOptionListConverter
-      extends SeparatedOptionListConverter {
+  public static class CommaSeparatedOptionListConverter extends SeparatedOptionListConverter {
     public CommaSeparatedOptionListConverter() {
       super(',', "comma");
     }
@@ -299,10 +285,10 @@ public final class Converters {
 
   public static class LogLevelConverter implements Converter<Level> {
 
-    public static final Level[] LEVELS = new Level[] {
-      Level.OFF, Level.SEVERE, Level.WARNING, Level.INFO, Level.FINE,
-      Level.FINER, Level.FINEST
-    };
+    public static final Level[] LEVELS =
+        new Level[] {
+          Level.OFF, Level.SEVERE, Level.WARNING, Level.INFO, Level.FINE, Level.FINER, Level.FINEST
+        };
 
     @Override
     public Level convert(String input) throws OptionsParsingException {
@@ -318,12 +304,9 @@ public final class Converters {
     public String getTypeDescription() {
       return "0 <= an integer <= " + (LEVELS.length - 1);
     }
-
   }
 
-  /**
-   * Checks whether a string is part of a set of strings.
-   */
+  /** Checks whether a string is part of a set of strings. */
   public static class StringSetConverter implements Converter<String> {
 
     // TODO(bazel-team): if this class never actually contains duplicates, we could s/List/Set/
@@ -349,9 +332,7 @@ public final class Converters {
     }
   }
 
-  /**
-   * Checks whether a string is a valid regex pattern and compiles it.
-   */
+  /** Checks whether a string is a valid regex pattern and compiles it. */
   public static class RegexPatternConverter implements Converter<Pattern> {
 
     @Override
@@ -369,9 +350,7 @@ public final class Converters {
     }
   }
 
-  /**
-   * Limits the length of a string argument.
-   */
+  /** Limits the length of a string argument. */
   public static class LengthLimitingConverter implements Converter<String> {
     private final int maxSize;
 
@@ -393,9 +372,7 @@ public final class Converters {
     }
   }
 
-  /**
-   * Checks whether an integer is in the given range.
-   */
+  /** Checks whether an integer is in the given range. */
   public static class RangeConverter implements Converter<Integer> {
     final int minValue;
     final int maxValue;
@@ -432,25 +409,27 @@ public final class Converters {
         return "an integer, >= " + minValue;
       } else {
         return "an integer in "
-            + (minValue < 0 ? "(" + minValue + ")" : minValue) + "-" + maxValue + " range";
+            + (minValue < 0 ? "(" + minValue + ")" : minValue)
+            + "-"
+            + maxValue
+            + " range";
       }
     }
   }
 
   /**
-   * A converter for variable assignments from the parameter list of a blaze
-   * command invocation. Assignments are expected to have the form "name=value",
-   * where names and values are defined to be as permissive as possible.
+   * A converter for variable assignments from the parameter list of a blaze command invocation.
+   * Assignments are expected to have the form "name=value", where names and values are defined to
+   * be as permissive as possible.
    */
   public static class AssignmentConverter implements Converter<Map.Entry<String, String>> {
 
     @Override
-    public Map.Entry<String, String> convert(String input)
-        throws OptionsParsingException {
+    public Map.Entry<String, String> convert(String input) throws OptionsParsingException {
       int pos = input.indexOf("=");
       if (pos <= 0) {
-        throw new OptionsParsingException("Variable definitions must be in the form of a "
-            + "'name=value' assignment");
+        throw new OptionsParsingException(
+            "Variable definitions must be in the form of a 'name=value' assignment");
       }
       String name = input.substring(0, pos);
       String value = input.substring(pos + 1);
@@ -461,24 +440,22 @@ public final class Converters {
     public String getTypeDescription() {
       return "a 'name=value' assignment";
     }
-
   }
 
   /**
-   * A converter for variable assignments from the parameter list of a blaze
-   * command invocation. Assignments are expected to have the form "name[=value]",
-   * where names and values are defined to be as permissive as possible and value
-   * part can be optional (in which case it is considered to be null).
+   * A converter for variable assignments from the parameter list of a blaze command invocation.
+   * Assignments are expected to have the form "name[=value]", where names and values are defined to
+   * be as permissive as possible and value part can be optional (in which case it is considered to
+   * be null).
    */
   public static class OptionalAssignmentConverter implements Converter<Map.Entry<String, String>> {
 
     @Override
-    public Map.Entry<String, String> convert(String input)
-        throws OptionsParsingException {
-      int pos = input.indexOf("=");
+    public Map.Entry<String, String> convert(String input) throws OptionsParsingException {
+      int pos = input.indexOf('=');
       if (pos == 0 || input.length() == 0) {
-        throw new OptionsParsingException("Variable definitions must be in the form of a "
-            + "'name=value' or 'name' assignment");
+        throw new OptionsParsingException(
+            "Variable definitions must be in the form of a 'name=value' or 'name' assignment");
       } else if (pos < 0) {
         return Maps.immutableEntry(input, null);
       }
@@ -491,7 +468,40 @@ public final class Converters {
     public String getTypeDescription() {
       return "a 'name=value' assignment with an optional value part";
     }
+  }
+
+  /**
+   * A converter for named integers of the form "[name=]value". When no name is specified, an empty
+   * string is used for the key.
+   */
+  public static class NamedIntegersConverter implements Converter<Map.Entry<String, Integer>> {
 
+    @Override
+    public Map.Entry<String, Integer> convert(String input) throws OptionsParsingException {
+      int pos = input.indexOf('=');
+      if (pos == 0 || input.length() == 0) {
+        throw new OptionsParsingException(
+            "Specify either 'value' or 'name=value', where 'value' is an integer");
+      } else if (pos < 0) {
+        try {
+          return Maps.immutableEntry("", Integer.parseInt(input));
+        } catch (NumberFormatException e) {
+          throw new OptionsParsingException("'" + input + "' is not an int");
+        }
+      }
+      String name = input.substring(0, pos);
+      String value = input.substring(pos + 1);
+      try {
+        return Maps.immutableEntry(name, Integer.parseInt(value));
+      } catch (NumberFormatException e) {
+        throw new OptionsParsingException("'" + value + "' is not an int");
+      }
+    }
+
+    @Override
+    public String getTypeDescription() {
+      return "an integer or a named integer, 'name=value'";
+    }
   }
 
   public static class HelpVerbosityConverter extends EnumConverter<OptionsParser.HelpVerbosity> {
@@ -508,5 +518,4 @@ public final class Converters {
       super(0, 100);
     }
   }
-
 }