Simplified ActionContextConsumer by having it operate on a new class which holds a variety of strategy/context maps.

PiperOrigin-RevId: 190491357

1 files changed, 19 insertions, 129 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
index 3b238d96c7..3323f58f26 100644
--- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
+++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java
@@ -16,24 +16,15 @@ package com.google.devtools.build.lib.buildtool;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
-import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.base.Stopwatch;
-import com.google.common.base.Strings;
-import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
-import com.google.common.collect.Multimap;
-import com.google.common.collect.Ordering;
-import com.google.common.collect.Table;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionCacheChecker;
-import com.google.devtools.build.lib.actions.ActionContext;
-import com.google.devtools.build.lib.actions.ActionContextMarker;
 import com.google.devtools.build.lib.actions.ActionGraph;
 import com.google.devtools.build.lib.actions.ActionInputFileCache;
 import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
@@ -41,14 +32,12 @@ import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactFactory;
 import com.google.devtools.build.lib.actions.BuildFailedException;
 import com.google.devtools.build.lib.actions.ExecException;
-import com.google.devtools.build.lib.actions.ExecutionStrategy;
 import com.google.devtools.build.lib.actions.Executor;
 import com.google.devtools.build.lib.actions.ExecutorInitException;
 import com.google.devtools.build.lib.actions.LocalHostCapacity;
 import com.google.devtools.build.lib.actions.PackageRoots;
 import com.google.devtools.build.lib.actions.ResourceManager;
 import com.google.devtools.build.lib.actions.ResourceSet;
-import com.google.devtools.build.lib.actions.SpawnActionContext;
 import com.google.devtools.build.lib.actions.TestExecException;
 import com.google.devtools.build.lib.actions.cache.ActionCache;
 import com.google.devtools.build.lib.actions.cache.Protos.ActionCacheStatistics;
@@ -60,7 +49,6 @@ import com.google.devtools.build.lib.analysis.WorkspaceStatusAction;
 import com.google.devtools.build.lib.analysis.actions.SymlinkTreeActionContext;
 import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.config.BuildConfigurationCollection;
-import com.google.devtools.build.lib.analysis.test.TestActionContext;
 import com.google.devtools.build.lib.buildtool.buildevent.ExecutionPhaseCompleteEvent;
 import com.google.devtools.build.lib.buildtool.buildevent.ExecutionStartingEvent;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
@@ -76,6 +64,7 @@ import com.google.devtools.build.lib.exec.ExecutionOptions;
 import com.google.devtools.build.lib.exec.ExecutorBuilder;
 import com.google.devtools.build.lib.exec.FilesetActionContextImpl;
 import com.google.devtools.build.lib.exec.SingleBuildFileCache;
+import com.google.devtools.build.lib.exec.SpawnActionContextMaps;
 import com.google.devtools.build.lib.exec.SymlinkTreeStrategy;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.profiler.ProfilePhase;
@@ -90,7 +79,6 @@ import com.google.devtools.build.lib.skyframe.Builder;
 import com.google.devtools.build.lib.skyframe.OutputService;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import com.google.devtools.build.lib.util.AbruptExitException;
-import com.google.devtools.build.lib.util.ExitCode;
 import com.google.devtools.build.lib.util.LoggingUtil;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.ModifiedFileSet;
@@ -102,16 +90,11 @@ import java.io.OutputStream;
 import java.io.OutputStreamWriter;
 import java.io.PrintWriter;
 import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.TreeMap;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
@@ -129,50 +112,6 @@ import java.util.logging.Logger;
  * @see com.google.devtools.build.lib.analysis.BuildView
  */
 public class ExecutionTool {
-  private static class StrategyConverter {
-    private Table<Class<? extends ActionContext>, String, ActionContext> classMap =
-        HashBasedTable.create();
-    private Map<Class<? extends ActionContext>, ActionContext> defaultClassMap =
-        new HashMap<>();
-
-    /**
-     * Aggregates all {@link ActionContext}s that are in {@code contextProviders}.
-     */
-    @SuppressWarnings("unchecked")
-    private StrategyConverter(Iterable<ActionContextProvider> contextProviders) {
-      for (ActionContextProvider provider : contextProviders) {
-        for (ActionContext strategy : provider.getActionContexts()) {
-          ExecutionStrategy annotation =
-              strategy.getClass().getAnnotation(ExecutionStrategy.class);
-          // TODO(ulfjack): Don't silently ignore action contexts without annotation.
-          if (annotation != null) {
-            defaultClassMap.put(annotation.contextType(), strategy);
-
-            for (String name : annotation.name()) {
-              classMap.put(annotation.contextType(), name, strategy);
-            }
-          }
-        }
-      }
-    }
-
-    @SuppressWarnings("unchecked")
-    private <T extends ActionContext> T getStrategy(Class<T> clazz, String name) {
-      return (T) (name.isEmpty() ? defaultClassMap.get(clazz) : classMap.get(clazz, name));
-    }
-
-    private String getValidValues(Class<? extends ActionContext> context) {
-      return Joiner.on(", ").join(Ordering.natural().sortedCopy(classMap.row(context).keySet()));
-    }
-
-    private String getUserFriendlyName(Class<? extends ActionContext> context) {
-      ActionContextMarker marker = context.getAnnotation(ActionContextMarker.class);
-      return marker != null
-          ? marker.name()
-          : context.getSimpleName();
-    }
-  }
-
   static final Logger logger = Logger.getLogger(ExecutionTool.class.getName());
 
   private final CommandEnvironment env;
@@ -182,10 +121,7 @@ public class ExecutionTool {
   private final ActionInputFileCache fileCache;
   private final ActionInputPrefetcher prefetcher;
   private final ImmutableList<ActionContextProvider> actionContextProviders;
-
-  private Map<String, SpawnActionContext> spawnStrategyMap =
-      new TreeMap<>(String.CASE_INSENSITIVE_ORDER);
-  private List<ActionContext> strategies = new ArrayList<>();
+  private SpawnActionContextMaps spawnActionContextMaps;
 
   ExecutionTool(CommandEnvironment env, BuildRequest request) throws ExecutorInitException {
     this.env = env;
@@ -208,20 +144,11 @@ public class ExecutionTool {
     // these dependencies should be added to the ActionContextConsumer of the module that actually
     // depends on them.
     builder.addActionContextConsumer(
-        new ActionContextConsumer() {
-          @Override
-          public ImmutableMap<String, String> getSpawnActionContexts() {
-            return ImmutableMap.of();
-          }
-
-          @Override
-          public Multimap<Class<? extends ActionContext>, String> getActionContexts() {
-            return ImmutableMultimap.<Class<? extends ActionContext>, String>builder()
-                .put(FilesetActionContext.class, "")
-                .put(WorkspaceStatusAction.Context.class, "")
-                .put(SymlinkTreeActionContext.class, "")
-                .build();
-          }
+        b -> {
+          b.strategyByContextMap()
+              .put(FilesetActionContext.class, "")
+              .put(WorkspaceStatusAction.Context.class, "")
+              .put(SymlinkTreeActionContext.class, "");
         });
 
     // Unfortunately, the exec root cache is not shared with caches in the remote execution client.
@@ -229,60 +156,24 @@ public class ExecutionTool {
         new SingleBuildFileCache(
             env.getExecRoot().getPathString(), env.getRuntime().getFileSystem());
     this.prefetcher = builder.getActionInputPrefetcher();
-        
+
     this.actionContextProviders = builder.getActionContextProviders();
     for (ActionContextProvider provider : actionContextProviders) {
       provider.init(fileCache);
     }
 
-    StrategyConverter strategyConverter = new StrategyConverter(actionContextProviders);
-
+    // There are many different SpawnActions, and we want to control the action context they use
+    // independently from each other, for example, to run genrules locally and Java compile action
+    // in prod. Thus, for SpawnActions, we decide the action context to use not only based on the
+    // context class, but also the mnemonic of the action.
+    SpawnActionContextMaps.Builder spawnActionContextMapsBuilder =
+        new SpawnActionContextMaps.Builder();
     for (ActionContextConsumer consumer : builder.getActionContextConsumers()) {
-      // There are many different SpawnActions, and we want to control the action context they use
-      // independently from each other, for example, to run genrules locally and Java compile action
-      // in prod. Thus, for SpawnActions, we decide the action context to use not only based on the
-      // context class, but also the mnemonic of the action.
-      for (Map.Entry<String, String> entry : consumer.getSpawnActionContexts().entrySet()) {
-        SpawnActionContext context =
-            strategyConverter.getStrategy(SpawnActionContext.class, entry.getValue());
-        if (context == null) {
-          String strategy = Strings.emptyToNull(entry.getKey());
-          throw makeExceptionForInvalidStrategyValue(
-              entry.getValue(),
-              Joiner.on(' ').skipNulls().join(strategy, "spawn"),
-              strategyConverter.getValidValues(SpawnActionContext.class));
-        }
-        spawnStrategyMap.put(entry.getKey(), context);
-      }
-
-      for (Map.Entry<Class<? extends ActionContext>, String> entry :
-          consumer.getActionContexts().entries()) {
-        ActionContext context = strategyConverter.getStrategy(entry.getKey(), entry.getValue());
-        if (context == null) {
-          throw makeExceptionForInvalidStrategyValue(
-              entry.getValue(),
-              strategyConverter.getUserFriendlyName(entry.getKey()),
-              strategyConverter.getValidValues(entry.getKey()));
-        }
-        strategies.add(context);
-      }
-    }
-
-    String testStrategyValue = request.getOptions(ExecutionOptions.class).testStrategy;
-    ActionContext context = strategyConverter.getStrategy(TestActionContext.class,
-        testStrategyValue);
-    if (context == null) {
-      throw makeExceptionForInvalidStrategyValue(testStrategyValue, "test",
-          strategyConverter.getValidValues(TestActionContext.class));
+      consumer.populate(spawnActionContextMapsBuilder);
     }
-    strategies.add(context);
-  }
-
-  private static ExecutorInitException makeExceptionForInvalidStrategyValue(String value,
-      String strategy, String validValues) {
-    return new ExecutorInitException(String.format(
-        "'%s' is an invalid value for %s strategy. Valid values are: %s", value, strategy,
-        validValues), ExitCode.COMMAND_LINE_ERROR);
+    spawnActionContextMaps =
+        spawnActionContextMapsBuilder.build(
+            actionContextProviders, request.getOptions(ExecutionOptions.class).testStrategy);
   }
 
   Executor getExecutor() throws ExecutorInitException {
@@ -304,8 +195,7 @@ public class ExecutionTool {
         env.getEventBus(),
         runtime.getClock(),
         request,
-        strategies,
-        spawnStrategyMap,
+        spawnActionContextMaps,
         actionContextProviders);
   }