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

3 files changed, 24 insertions, 17 deletions

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();