Remove special handling of SpawnActionContext in Executor/ActionExecContext

Instead, internally look up the correct context by mnemonic. This simplifies
all the callers. We still need a little bit of special casing when constructing
the action context map.

PiperOrigin-RevId: 197572357

12 files changed, 46 insertions, 49 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionContext.java b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionContext.java
index 411475e8b1..eae8b9c58a 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionContext.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionExecutionContext.java
@@ -232,11 +232,6 @@ public class ActionExecutionContext implements Closeable {
     return executor.getContext(type);
   }
 
-  /** Returns the action context implementation for a given spawn action. */
-  public SpawnActionContext getSpawnActionContext(Spawn spawn) {
-    return executor.getSpawnActionContext(spawn);
-  }
-
   /**
    * Whether this Executor reports subcommands. If not, reportSubcommand has no effect.
    * This is provided so the caller of reportSubcommand can avoid wastefully constructing the
diff --git a/src/main/java/com/google/devtools/build/lib/actions/Executor.java b/src/main/java/com/google/devtools/build/lib/actions/Executor.java
index 5511052220..a2de3d6a78 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/Executor.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/Executor.java
@@ -86,7 +86,4 @@ public interface Executor {
    * Looks up and returns an action context implementation of the given interface type.
    */
   <T extends ActionContext> T getContext(Class<? extends T> type);
-
-  /** Returns the action context implementation for the given spawn. */
-  SpawnActionContext getSpawnActionContext(Spawn spawn);
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
index 6191b272a8..e2cbcc89bd 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java
@@ -283,7 +283,8 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
   protected List<SpawnResult> internalExecute(ActionExecutionContext actionExecutionContext)
       throws ExecException, InterruptedException, CommandLineExpansionException {
     Spawn spawn = getSpawn(actionExecutionContext);
-    return getContext(actionExecutionContext, spawn).exec(spawn, actionExecutionContext);
+    return actionExecutionContext.getContext(SpawnActionContext.class)
+        .exec(spawn, actionExecutionContext);
   }
 
   @Override
@@ -495,11 +496,6 @@ public class SpawnAction extends AbstractAction implements ExecutionInfoSpecifie
     return executionInfo;
   }
 
-  protected SpawnActionContext getContext(
-      ActionExecutionContext actionExecutionContext, Spawn spawn) {
-    return actionExecutionContext.getSpawnActionContext(spawn);
-  }
-
   /** A spawn instance that is tied to a specific SpawnAction. */
   private class ActionSpawn extends BaseSpawn {
 
diff --git a/src/main/java/com/google/devtools/build/lib/exec/BlazeExecutor.java b/src/main/java/com/google/devtools/build/lib/exec/BlazeExecutor.java
index 7d41a156ba..afd04483b0 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/BlazeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/BlazeExecutor.java
@@ -18,8 +18,6 @@ import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.Executor;
 import com.google.devtools.build.lib.actions.ExecutorInitException;
-import com.google.devtools.build.lib.actions.Spawn;
-import com.google.devtools.build.lib.actions.SpawnActionContext;
 import com.google.devtools.build.lib.clock.Clock;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
 import com.google.devtools.build.lib.events.Event;
@@ -58,7 +56,6 @@ public final class BlazeExecutor implements Executor {
   private AtomicBoolean inExecutionPhase;
 
   private final Map<Class<? extends ActionContext>, ActionContext> contextMap;
-  private final SpawnActionContextMaps spawnActionContextMaps;
 
   /**
    * Constructs an Executor, bound to a specified output base path, and which will use the specified
@@ -90,7 +87,6 @@ public final class BlazeExecutor implements Executor {
     this.eventBus = eventBus;
     this.clock = clock;
     this.options = options;
-    this.spawnActionContextMaps = spawnActionContextMaps;
     this.inExecutionPhase = new AtomicBoolean(false);
     this.contextMap = spawnActionContextMaps.contextMap();
 
@@ -180,20 +176,9 @@ public final class BlazeExecutor implements Executor {
 
   @Override
   public <T extends ActionContext> T getContext(Class<? extends T> type) {
-    Preconditions.checkArgument(type != SpawnActionContext.class,
-        "should use getSpawnActionContext instead");
     return type.cast(contextMap.get(type));
   }
 
-  /**
-   * Returns the {@link SpawnActionContext} to use for the given mnemonic. If no execution mode is
-   * set, then it returns the default strategy for spawn actions.
-   */
-  @Override
-  public SpawnActionContext getSpawnActionContext(Spawn spawn) {
-    return spawnActionContextMaps.getSpawnActionContext(spawn, reporter);
-  }
-
   /** Returns true iff the --verbose_failures option was enabled. */
   @Override
   public boolean getVerboseFailures() {
diff --git a/src/main/java/com/google/devtools/build/lib/exec/SpawnActionContextMaps.java b/src/main/java/com/google/devtools/build/lib/exec/SpawnActionContextMaps.java
index 20adb6040f..4fc32aa37d 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/SpawnActionContextMaps.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/SpawnActionContextMaps.java
@@ -29,12 +29,16 @@ import com.google.common.collect.Ordering;
 import com.google.common.collect.Table;
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.ActionContextMarker;
+import com.google.devtools.build.lib.actions.ActionExecutionContext;
+import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionStrategy;
 import com.google.devtools.build.lib.actions.ExecutorInitException;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnActionContext;
+import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.analysis.test.TestActionContext;
 import com.google.devtools.build.lib.events.Event;
+import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.Reporter;
 import com.google.devtools.build.lib.util.ExitCode;
 import com.google.devtools.build.lib.util.RegexFilter;
@@ -86,7 +90,7 @@ public final class SpawnActionContextMaps {
    * <p>If the reason for selecting the context is worth mentioning to the user, logs a message
    * using the given {@link Reporter}.
    */
-  public SpawnActionContext getSpawnActionContext(Spawn spawn, Reporter reporter) {
+  public SpawnActionContext getSpawnActionContext(Spawn spawn, EventHandler reporter) {
     Preconditions.checkNotNull(spawn);
     if (!spawnStrategyRegexList.isEmpty() && spawn.getResourceOwner() != null) {
       String description = spawn.getResourceOwner().getProgressMessage();
@@ -117,6 +121,7 @@ public final class SpawnActionContextMaps {
       }
       contextMap.put(context.getClass(), context);
     }
+    contextMap.put(SpawnActionContext.class, new ProxySpawnActionContext());
     return ImmutableMap.copyOf(contextMap);
   }
 
@@ -346,4 +351,20 @@ public final class SpawnActionContextMaps {
       return marker != null ? marker.name() : context.getSimpleName();
     }
   }
+
+  /** Proxy that looks up the right SpawnActionContext for a spawn during exec. */
+  @VisibleForTesting
+  public final class ProxySpawnActionContext implements SpawnActionContext {
+    @Override
+    public List<SpawnResult> exec(Spawn spawn, ActionExecutionContext actionExecutionContext)
+        throws ExecException, InterruptedException {
+      return resolve(spawn, actionExecutionContext.getEventHandler())
+          .exec(spawn, actionExecutionContext);
+    }
+
+    @VisibleForTesting
+    public SpawnActionContext resolve(Spawn spawn, EventHandler eventHandler) {
+      return getSpawnActionContext(spawn, eventHandler);
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java b/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
index 7ecef0cfef..21bf683f07 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/StandaloneTestStrategy.java
@@ -338,7 +338,8 @@ public class StandaloneTestStrategy extends TestStrategy {
     TestResultData.Builder builder = TestResultData.newBuilder();
 
     long startTime = actionExecutionContext.getClock().currentTimeMillis();
-    SpawnActionContext spawnActionContext = actionExecutionContext.getSpawnActionContext(spawn);
+    SpawnActionContext spawnActionContext =
+        actionExecutionContext.getContext(SpawnActionContext.class);
     List<SpawnResult> spawnResults = ImmutableList.of();
     BuildEventStreamProtos.TestResult.ExecutionInfo.Builder executionInfo =
         BuildEventStreamProtos.TestResult.ExecutionInfo.newBuilder();
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
index 4547dbe35e..53801f3e69 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppLinkAction.java
@@ -40,6 +40,7 @@ import com.google.devtools.build.lib.actions.ExecutionRequirements;
 import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.actions.SimpleSpawn;
 import com.google.devtools.build.lib.actions.Spawn;
+import com.google.devtools.build.lib.actions.SpawnActionContext;
 import com.google.devtools.build.lib.actions.extra.CppLinkInfo;
 import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.analysis.RuleContext;
@@ -318,7 +319,7 @@ public final class CppLinkAction extends AbstractAction
                 estimateResourceConsumptionLocal());
         return ActionResult.create(
             actionExecutionContext
-                .getSpawnActionContext(spawn)
+                .getContext(SpawnActionContext.class)
                 .exec(spawn, actionExecutionContext));
       } catch (ExecException e) {
         throw e.toActionExecutionException(
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/SpawnGccStrategy.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/SpawnGccStrategy.java
index 2ce7ec48ab..0ba465a6b8 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/SpawnGccStrategy.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/SpawnGccStrategy.java
@@ -62,7 +62,8 @@ public class SpawnGccStrategy implements CppCompileActionContext {
             action.estimateResourceConsumptionLocal());
 
     List<SpawnResult> spawnResults =
-        actionExecutionContext.getSpawnActionContext(spawn).exec(spawn, actionExecutionContext);
+        actionExecutionContext.getContext(SpawnActionContext.class)
+            .exec(spawn, actionExecutionContext);
     return CppCompileActionResult.builder().setSpawnResults(spawnResults).build();
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java
index 0c29b59382..7f6436f8c9 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java
@@ -154,7 +154,7 @@ public class JavaHeaderCompileAction extends SpawnAction {
   protected List<SpawnResult> internalExecute(ActionExecutionContext actionExecutionContext)
       throws ExecException, InterruptedException {
     Spawn spawn = getDirectSpawn();
-    SpawnActionContext context = getContext(actionExecutionContext, spawn);
+    SpawnActionContext context = actionExecutionContext.getContext(SpawnActionContext.class);
     try {
       return context.exec(spawn, actionExecutionContext);
     } catch (ExecException e) {
diff --git a/src/test/java/com/google/devtools/build/lib/actions/util/DummyExecutor.java b/src/test/java/com/google/devtools/build/lib/actions/util/DummyExecutor.java
index 1d0cd9be99..7b09b6b3f5 100644
--- a/src/test/java/com/google/devtools/build/lib/actions/util/DummyExecutor.java
+++ b/src/test/java/com/google/devtools/build/lib/actions/util/DummyExecutor.java
@@ -16,8 +16,6 @@ package com.google.devtools.build.lib.actions.util;
 import com.google.common.eventbus.EventBus;
 import com.google.devtools.build.lib.actions.ActionContext;
 import com.google.devtools.build.lib.actions.Executor;
-import com.google.devtools.build.lib.actions.Spawn;
-import com.google.devtools.build.lib.actions.SpawnActionContext;
 import com.google.devtools.build.lib.clock.BlazeClock;
 import com.google.devtools.build.lib.clock.Clock;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
@@ -84,11 +82,6 @@ public final class DummyExecutor implements Executor {
   }
 
   @Override
-  public SpawnActionContext getSpawnActionContext(Spawn spawn) {
-    throw new UnsupportedOperationException();
-  }
-
-  @Override
   public OptionsClassProvider getOptions() {
     throw new UnsupportedOperationException();
   }
diff --git a/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java b/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
index 0f0a02363d..7c89da42ef 100644
--- a/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
+++ b/src/test/java/com/google/devtools/build/lib/exec/StandaloneTestStrategyTest.java
@@ -19,6 +19,7 @@ import static com.google.devtools.build.lib.testutil.TestConstants.WORKSPACE_NAM
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.same;
 import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableList;
@@ -135,7 +136,8 @@ public final class StandaloneTestStrategyTest extends BuildViewTestCase {
             .build();
     when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
 
-    when(actionExecutionContext.getSpawnActionContext(any())).thenReturn(spawnActionContext);
+    when(actionExecutionContext.getContext(same(SpawnActionContext.class)))
+        .thenReturn(spawnActionContext);
 
     // actual StandaloneTestStrategy execution
     List<SpawnResult> spawnResults =
@@ -221,7 +223,8 @@ public final class StandaloneTestStrategyTest extends BuildViewTestCase {
         .thenThrow(new SpawnExecException("test failed", failSpawnResult, false))
         .thenReturn(ImmutableList.of(passSpawnResult));
 
-    when(actionExecutionContext.getSpawnActionContext(any())).thenReturn(spawnActionContext);
+    when(actionExecutionContext.getContext(same(SpawnActionContext.class)))
+        .thenReturn(spawnActionContext);
 
     // actual StandaloneTestStrategy execution
     List<SpawnResult> spawnResults =
@@ -306,7 +309,8 @@ public final class StandaloneTestStrategyTest extends BuildViewTestCase {
             .build();
     when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
 
-    when(actionExecutionContext.getSpawnActionContext(any())).thenReturn(spawnActionContext);
+    when(actionExecutionContext.getContext(same(SpawnActionContext.class)))
+        .thenReturn(spawnActionContext);
 
     // actual StandaloneTestStrategy execution
     List<SpawnResult> spawnResults =
@@ -383,7 +387,8 @@ public final class StandaloneTestStrategyTest extends BuildViewTestCase {
             .build();
     when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
 
-    when(actionExecutionContext.getSpawnActionContext(any())).thenReturn(spawnActionContext);
+    when(actionExecutionContext.getContext(same(SpawnActionContext.class)))
+        .thenReturn(spawnActionContext);
 
     // actual StandaloneTestStrategy execution
     List<SpawnResult> spawnResults =
@@ -485,7 +490,8 @@ public final class StandaloneTestStrategyTest extends BuildViewTestCase {
                 expectedSpawnResult,
                 /*forciblyRunRemotely=*/ false,
                 /*catastrophe=*/ false));
-    when(actionExecutionContext.getSpawnActionContext(any())).thenReturn(spawnActionContext);
+    when(actionExecutionContext.getContext(same(SpawnActionContext.class)))
+        .thenReturn(spawnActionContext);
 
     // actual StandaloneTestStrategy execution
     List<SpawnResult> spawnResults =
@@ -563,7 +569,8 @@ public final class StandaloneTestStrategyTest extends BuildViewTestCase {
     SpawnResult expectedSpawnResult =
         new SpawnResult.Builder().setStatus(Status.SUCCESS).setRunnerName("test").build();
     when(spawnActionContext.exec(any(), any())).thenReturn(ImmutableList.of(expectedSpawnResult));
-    when(actionExecutionContext.getSpawnActionContext(any())).thenReturn(spawnActionContext);
+    when(actionExecutionContext.getContext(same(SpawnActionContext.class)))
+        .thenReturn(spawnActionContext);
 
     // actual StandaloneTestStrategy execution
     List<SpawnResult> spawnResults =
diff --git a/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java b/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
index d4ecd6ceb9..da31f120d2 100644
--- a/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
+++ b/src/test/java/com/google/devtools/build/lib/standalone/StandaloneSpawnStrategyTest.java
@@ -172,14 +172,14 @@ public class StandaloneSpawnStrategyTest {
   @Test
   public void testBinTrueExecutesFine() throws Exception {
     Spawn spawn = createSpawn(getTrueCommand());
-    executor.getSpawnActionContext(spawn).exec(spawn, createContext());
+    executor.getContext(SpawnActionContext.class).exec(spawn, createContext());
 
     assertThat(out()).isEmpty();
     assertThat(err()).isEmpty();
   }
 
   private List<SpawnResult> run(Spawn spawn) throws Exception {
-    return executor.getSpawnActionContext(spawn).exec(spawn, createContext());
+    return executor.getContext(SpawnActionContext.class).exec(spawn, createContext());
   }
 
   private ActionExecutionContext createContext() {