Filter spawn strategies by execution platform

src/main/java/com/google/devtools/build/lib/bazel/rules/BazelStrategyModule.java

@@ -18,6 +18,7 @@
 import com.google.devtools.build.lib.analysis.actions.FileWriteActionContext;
 import com.google.devtools.build.lib.analysis.actions.TemplateExpansionContext;
 import com.google.devtools.build.lib.buildtool.BuildRequest;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.exec.ExecutionOptions;
 import com.google.devtools.build.lib.exec.ModuleActionContextRegistry;
 import com.google.devtools.build.lib.exec.SpawnCache;
@@ -90,5 +91,9 @@ public void registerSpawnStrategies(
     for (Map.Entry<RegexFilter, List<String>> entry : options.strategyByRegexp) {
       registryBuilder.addDescriptionFilter(entry.getKey(), entry.getValue());
     }
+
+    for (Map.Entry<String, List<String>> strategy : options.allowedStrategiesByExecPlatform) {
+      registryBuilder.addExecPlatformFilter(Label.parseCanonicalUnchecked(strategy.getKey()), strategy.getValue());
+    }
   }
 }

src/main/java/com/google/devtools/build/lib/exec/BUILD

@@ -107,6 +107,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/util:resource_converter",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/common/options",
+        "//src/main/java/com/google/devtools/build/lib/analysis:config/fragment_options",
         "//third_party:guava",
     ],
 )
@@ -363,6 +364,7 @@ java_library(
         ":remote_local_fallback_registry",
         ":spawn_strategy_policy",
         "//src/main/java/com/google/devtools/build/lib/actions",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",

src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java

@@ -122,6 +122,23 @@ public class ExecutionOptions extends OptionsBase {
               + "the 'local' strategy, but reversing the order would run it with 'sandboxed'. ")
   public List<Map.Entry<RegexFilter, List<String>>> strategyByRegexp;
 
+  @Option(
+      name = "allowed_strategies_by_exec_platform",
+      allowMultiple = true,
+      converter = Converters.StringToStringListConverter.class,
+      defaultValue = "null",
+      documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
+      effectTags = {OptionEffectTag.EXECUTION},
+      help =
+          """
+          Filters spawn strategies by the execution platform.
+          Example: `--allowed_strategies_by_exec_platform=@platforms//host:host=local,sandboxed,worker`
+          to prevent actions configured for the host platform from being spawned remotely.
+          Example: `--allowed_strategies_by_exec_platform=//:linux_amd64=remote` to prevent actions
+          configured for a platform `//:linux_amd64` from being spawned locally.
+          """)
+  public List<Map.Entry<String, List<String>>> allowedStrategiesByExecPlatform;
+
   @Option(
       name = "materialize_param_files",
       defaultValue = "false",

src/main/java/com/google/devtools/build/lib/exec/RemoteLocalFallbackRegistry.java

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.exec;
 
 import com.google.devtools.build.lib.actions.ActionContext;
+import com.google.devtools.build.lib.actions.Spawn;
 import javax.annotation.Nullable;
 
 /**
@@ -29,5 +30,5 @@ public interface RemoteLocalFallbackRegistry extends ActionContext {
    * @return remote fallback strategy or {@code null} if none was registered
    */
   @Nullable
-  AbstractSpawnStrategy getRemoteLocalFallbackStrategy();
+  AbstractSpawnStrategy getRemoteLocalFallbackStrategy(Spawn spawn);
 }

src/main/java/com/google/devtools/build/lib/exec/SpawnStrategyRegistry.java

@@ -22,6 +22,7 @@
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.LinkedListMultimap;
 import com.google.common.collect.ListMultimap;
@@ -34,6 +35,7 @@
 import com.google.devtools.build.lib.actions.SandboxedSpawnStrategy;
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnStrategy;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.Reporter;
@@ -69,6 +71,7 @@ public final class SpawnStrategyRegistry
 
   private final ImmutableListMultimap<String, SpawnStrategy> mnemonicToStrategies;
   private final StrategyRegexFilter strategyRegexFilter;
+  private final StrategyPlatformFilter strategyPlatformFilter;
   private final ImmutableList<? extends SpawnStrategy> defaultStrategies;
   private final ImmutableMultimap<String, SandboxedSpawnStrategy> mnemonicToRemoteDynamicStrategies;
   private final ImmutableMultimap<String, SandboxedSpawnStrategy> mnemonicToLocalDynamicStrategies;
@@ -77,12 +80,14 @@ public final class SpawnStrategyRegistry
   private SpawnStrategyRegistry(
       ImmutableListMultimap<String, SpawnStrategy> mnemonicToStrategies,
       StrategyRegexFilter strategyRegexFilter,
+      StrategyPlatformFilter strategyPlatformFilter,
       ImmutableList<? extends SpawnStrategy> defaultStrategies,
       ImmutableMultimap<String, SandboxedSpawnStrategy> mnemonicToRemoteDynamicStrategies,
       ImmutableMultimap<String, SandboxedSpawnStrategy> mnemonicToLocalDynamicStrategies,
       @Nullable AbstractSpawnStrategy remoteLocalFallbackStrategy) {
     this.mnemonicToStrategies = mnemonicToStrategies;
     this.strategyRegexFilter = strategyRegexFilter;
+    this.strategyPlatformFilter = strategyPlatformFilter;
     this.defaultStrategies = defaultStrategies;
     this.mnemonicToRemoteDynamicStrategies = mnemonicToRemoteDynamicStrategies;
     this.mnemonicToLocalDynamicStrategies = mnemonicToLocalDynamicStrategies;
@@ -105,7 +110,8 @@ private SpawnStrategyRegistry(
    * using the given {@link Reporter}.
    */
   public List<? extends SpawnStrategy> getStrategies(Spawn spawn, @Nullable EventHandler reporter) {
-    return getStrategies(spawn.getResourceOwner(), spawn.getMnemonic(), reporter);
+    return strategyPlatformFilter.getStrategies(
+        spawn, getStrategies(spawn.getResourceOwner(), spawn.getMnemonic(), reporter));
   }
 
   /**
@@ -116,6 +122,8 @@ public List<? extends SpawnStrategy> getStrategies(Spawn spawn, @Nullable EventH
    *
    * <p>If the reason for selecting the context is worth mentioning to the user, logs a message
    * using the given {@link Reporter}.
+   *
+   * NOTE: This method is public for Blaze, getStrategies(Spawn, EventHandler) must be used in Bazel.
    */
   public List<? extends SpawnStrategy> getStrategies(
       ActionExecutionMetadata resourceOwner, String mnemonic, @Nullable EventHandler reporter) {
@@ -154,7 +162,7 @@ public ImmutableCollection<SandboxedSpawnStrategy> getDynamicSpawnActionContexts
             ? mnemonicToRemoteDynamicStrategies
             : mnemonicToLocalDynamicStrategies;
     if (mnemonicToDynamicStrategies.containsKey(spawn.getMnemonic())) {
-      return mnemonicToDynamicStrategies.get(spawn.getMnemonic());
+      return strategyPlatformFilter.getStrategies(spawn, mnemonicToDynamicStrategies.get(spawn.getMnemonic()));
     }
     if (mnemonicToDynamicStrategies.containsKey("")) {
       return mnemonicToDynamicStrategies.get("");
@@ -164,8 +172,9 @@ public ImmutableCollection<SandboxedSpawnStrategy> getDynamicSpawnActionContexts
 
   @Nullable
   @Override
-  public AbstractSpawnStrategy getRemoteLocalFallbackStrategy() {
-    return remoteLocalFallbackStrategy;
+  public AbstractSpawnStrategy getRemoteLocalFallbackStrategy(Spawn spawn) {
+    return strategyPlatformFilter.getStrategies(spawn, Lists.newArrayList(remoteLocalFallbackStrategy))
+        .getFirst();
   }
 
   /**
@@ -203,9 +212,16 @@ void logSpawnStrategies() {
         strategyRegexFilter.getFilterToStrategies().asMap().entrySet()) {
       Collection<SpawnStrategy> value = entry.getValue();
       logger.atInfo().log(
-          "FilterToStrategyImplementations: \"%s\" = [%s]",
+          "FilterDescriptionToStrategyImplementations: \"%s\" = [%s]",
           entry.getKey(), toImplementationNames(value));
     }
+    for (Map.Entry<Label, ImmutableList<SpawnStrategy>> entry :
+        strategyPlatformFilter.getFilterToStrategies().entrySet()) {
+      Collection<SpawnStrategy> value = entry.getValue();
+      logger.atInfo().log(
+          "FilterPlatformToStrategyImplementations: \"%s\" = [%s]",
+          entry.getKey().getCanonicalName(), toImplementationNames(value));
+    }
 
     logger.atInfo().log(
         "DefaultStrategyImplementations: [%s]", toImplementationNames(defaultStrategies));
@@ -287,6 +303,7 @@ public static final class Builder {
     private final HashMap<String, List<String>> mnemonicToRemoteDynamicIdentifiers =
         new HashMap<>();
     private final HashMap<String, List<String>> mnemonicToLocalDynamicIdentifiers = new HashMap<>();
+    private final HashMap<Label, List<String>> execPlatformFilters = new HashMap<>();
 
     @Nullable private String remoteLocalFallbackStrategyIdentifier;
 
@@ -316,6 +333,12 @@ public Builder addDescriptionFilter(RegexFilter filter, List<String> identifiers
       return this;
     }
 
+    @CanIgnoreReturnValue
+    public Builder addExecPlatformFilter(Label execPlatform, List<String> identifiers) {
+      this.execPlatformFilters.put(execPlatform, identifiers);
+      return this;
+    }
+
     /**
      * Adds a filter limiting any spawn whose {@linkplain Spawn#getMnemonic() mnemonic}
      * (case-sensitively) matches the given mnemonic to only use strategies with the given
@@ -446,6 +469,14 @@ public SpawnStrategyRegistry build() throws AbruptExitException {
         }
       }
 
+      ImmutableMap.Builder<Label, ImmutableList<SpawnStrategy>> platformToStrategies = ImmutableMap.builder();
+      for (Map.Entry<Label, List<String>> entry : execPlatformFilters.entrySet()) {
+        Label platform = entry.getKey();
+        platformToStrategies.put(
+          platform,
+          strategyMapper.toStrategies(entry.getValue(), "platform " + platform.getCanonicalName()));
+      }
+
       ImmutableListMultimap.Builder<String, SpawnStrategy> mnemonicToStrategies =
           new ImmutableListMultimap.Builder<>();
       for (Map.Entry<String, List<String>> entry : mnemonicToIdentifiers.entrySet()) {
@@ -516,6 +547,7 @@ public SpawnStrategyRegistry build() throws AbruptExitException {
           mnemonicToStrategies.build(),
           new StrategyRegexFilter(
               strategyMapper, strategyPolicy, filterToIdentifiers, filterToStrategies),
+          new StrategyPlatformFilter(strategyMapper, platformToStrategies.build()),
           defaultStrategies,
           mnemonicToRemoteStrategies.build(),
           mnemonicToLocalStrategies.build(),
@@ -597,6 +629,46 @@ public String toString() {
     }
   }
 
+  private static class StrategyPlatformFilter {
+    private final StrategyMapper strategyMapper;
+    private final ImmutableMap<Label, ImmutableList<SpawnStrategy>> platformToStrategies;
+
+    private StrategyPlatformFilter(
+      StrategyMapper strategyMapper,
+      ImmutableMap<Label, ImmutableList<SpawnStrategy>> platformToStrategies) {
+      this.strategyMapper = strategyMapper;
+      this.platformToStrategies = platformToStrategies;
+    }
+
+    public <T extends SpawnStrategy> List<T> getStrategies(
+        Spawn spawn, List<T> candidateStrategies) {
+      var platformLabel = spawn.getExecutionPlatformLabel();
+      Preconditions.checkNotNull(platformLabel, "Attempting to spawn action without an execution platform.");
+
+      var allowedStrategies = platformToStrategies.get(platformLabel);
+      if (allowedStrategies != null) {
+        List<T> filteredStrategies = new ArrayList<>();
+        for (var strategy : candidateStrategies) {
+          if (allowedStrategies.contains(strategy)) {
+            filteredStrategies.add(strategy);
+          }
+        }
+        return filteredStrategies;
+      }
+
+      return candidateStrategies;
+    }
+
+    public <T extends SpawnStrategy> ImmutableCollection<T> getStrategies(
+        Spawn spawn, ImmutableCollection<T> candidateStrategies) {
+      return ImmutableList.copyOf(getStrategies(spawn, Lists.newCopyOnWriteArrayList(candidateStrategies)));
+    }
+
+    ImmutableMap<Label, ImmutableList<SpawnStrategy>> getFilterToStrategies() {
+      return platformToStrategies;
+    }
+  }
+
   /* Maps the strategy identifier (e.g. "local", "worker"..) to the real strategy. */
   private static class StrategyMapper {
 

src/main/java/com/google/devtools/build/lib/exec/SpawnStrategyResolver.java

@@ -86,10 +86,12 @@ public List<? extends SpawnStrategy> resolve(
       if (fallbackStrategies.isEmpty()) {
         String message =
             String.format(
-                "%s spawn cannot be executed with any of the available strategies: %s. Your"
-                    + " --spawn_strategy, --genrule_strategy and/or --strategy flags are probably"
-                    + " too strict. Visit https://github.com/bazelbuild/bazel/issues/7480 for"
-                    + " advice",
+                """
+                %s spawn cannot be executed with any of the available strategies: %s. Your \
+                --spawn_strategy, --genrule_strategy, strategy and/or \
+                --allowed_strategies_by_exec_platform flags are probably too strict. \
+                Visit https://github.com/bazelbuild/bazel/issues/7480 for advice.
+                """,
                 spawn.getMnemonic(), strategies);
         throw new UserExecException(
             FailureDetail.newBuilder()

src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnRunner.java

@@ -555,7 +555,7 @@ private SpawnResult execLocally(Spawn spawn, SpawnExecutionContext context)
         context.getContext(RemoteLocalFallbackRegistry.class);
     checkNotNull(localFallbackRegistry, "Expected a RemoteLocalFallbackRegistry to be registered");
     AbstractSpawnStrategy remoteLocalFallbackStrategy =
-        localFallbackRegistry.getRemoteLocalFallbackStrategy();
+        localFallbackRegistry.getRemoteLocalFallbackStrategy(spawn);
     checkNotNull(
         remoteLocalFallbackStrategy,
         "A remote local fallback strategy must be set if using remote fallback.");

src/test/java/com/google/devtools/build/lib/exec/SpawnStrategyRegistryTest.java

@@ -28,6 +28,7 @@
 import com.google.devtools.build.lib.actions.Spawn;
 import com.google.devtools.build.lib.actions.SpawnResult;
 import com.google.devtools.build.lib.actions.SpawnStrategy;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
 import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.events.Event;
@@ -300,6 +301,24 @@ public void testDuplicatedDescriptionFilter() throws Exception {
         .containsExactly(strategy2);
   }
 
+  @Test
+  public void testPlatformFilter() throws Exception {
+    NoopStrategy strategy1 = new NoopStrategy("1");
+    NoopStrategy strategy2 = new NoopStrategy("2");
+    SpawnStrategyRegistry strategyRegistry =
+        SpawnStrategyRegistry.builder()
+            .registerStrategy(strategy1, "foo")
+            .registerStrategy(strategy2, "bar")
+            .addExecPlatformFilter(Label.parseCanonical("//:dummy_platform"), ImmutableList.of("foo"))
+            .build();
+
+    assertThat(
+            strategyRegistry.getStrategies(
+                createSpawnWithMnemonicAndDescription("", ""),
+                SpawnStrategyRegistryTest::noopEventHandler))
+        .containsExactly(strategy1);
+  }
+
   @Test
   public void testMultipleDefaultStrategies() throws Exception {
     NoopStrategy strategy1 = new NoopStrategy("1");
@@ -537,7 +556,7 @@ public void testRemoteLocalFallback() throws Exception {
             .setRemoteLocalFallbackStrategyIdentifier("bar")
             .build();
 
-    assertThat(strategyRegistry.getRemoteLocalFallbackStrategy()).isEqualTo(strategy2);
+    assertThat(strategyRegistry.getRemoteLocalFallbackStrategy(createSpawnWithMnemonicAndDescription("", ""))).isEqualTo(strategy2);
   }
 
   @Test
@@ -561,7 +580,7 @@ public void testRemoteLocalFallbackNotRegistered() throws Exception {
     SpawnStrategyRegistry strategyRegistry =
         SpawnStrategyRegistry.builder().registerStrategy(strategy1, "foo").build();
 
-    assertThat(strategyRegistry.getRemoteLocalFallbackStrategy()).isNull();
+    assertThat(strategyRegistry.getRemoteLocalFallbackStrategy(createSpawnWithMnemonicAndDescription("", ""))).isNull();
   }
 
   @Test

src/test/java/com/google/devtools/build/lib/exec/util/FakeOwner.java

@@ -81,7 +81,19 @@ private FakeOwner(
   }
 
   public FakeOwner(String mnemonic, String progressMessage, String ownerLabel) {
-    this(mnemonic, progressMessage, checkNotNull(ownerLabel), null);
+    this(
+        mnemonic,
+        progressMessage,
+        checkNotNull(ownerLabel),
+        createPlatformInfo());
+  }
+
+  private static PlatformInfo createPlatformInfo() {
+    try {
+      return PlatformInfo.builder().setLabel(Label.parseCanonical("//:dummy_platform")).build();
+    } catch (Exception ex) {
+      throw new RuntimeException("Fake PlatformInfo construction failed.", ex);
+    }
   }
 
   @Override

src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerTest.java

@@ -265,7 +265,7 @@ private FakeSpawnExecutionContext getSpawnContext(Spawn spawn) {
     AbstractSpawnStrategy fakeLocalStrategy =
         new AbstractSpawnStrategy(execRoot, localRunner, new ExecutionOptions()) {};
     ClassToInstanceMap<ActionContext> actionContextRegistry =
-        ImmutableClassToInstanceMap.of(RemoteLocalFallbackRegistry.class, () -> fakeLocalStrategy);
+        ImmutableClassToInstanceMap.of(RemoteLocalFallbackRegistry.class, (_spawn) -> fakeLocalStrategy);
 
     var actionInputFetcher =
         new RemoteActionInputFetcher(