Filter spawn strategies by execution platform

Author: Silic0nS0ldier

src/main/java/com/google/devtools/build/lib/analysis/platform/BUILD

@@ -85,6 +85,7 @@ java_library(
     name = "platform_utils",
     srcs = ["PlatformUtils.java"],
     deps = [
+        ":platform",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/remote/options",
         "//src/main/protobuf:failure_details_java_proto",

src/main/java/com/google/devtools/build/lib/analysis/platform/PlatformUtils.java

@@ -67,6 +67,29 @@ public static Platform getPlatformProto(Spawn spawn, @Nullable RemoteOptions rem
     return getPlatformProto(spawn, remoteOptions, ImmutableMap.of());
   }
 
+  private static boolean shouldProducePlatformProto(
+      Spawn spawn, SortedMap<String, String> defaultExecProperties, Map<String, String> additionalProperties) {
+    PlatformInfo executionPlatform = spawn.getExecutionPlatform();
+    if (executionPlatform != null) {
+      if (!executionPlatform.execProperties().isEmpty()) {
+        return true;
+      }
+      if (!executionPlatform.remoteExecutionProperties().isEmpty()) {
+        return true;
+      }
+    }
+    if (!spawn.getCombinedExecProperties().isEmpty()) {
+      return true;
+    }
+    if (!defaultExecProperties.isEmpty()) {
+      return true;
+    }
+    if (!additionalProperties.isEmpty()) {
+      return true;
+    }
+    return false;
+  }
+
   @Nullable
   public static Platform getPlatformProto(
       Spawn spawn, @Nullable RemoteOptions remoteOptions, Map<String, String> additionalProperties)
@@ -76,10 +99,8 @@ public static Platform getPlatformProto(
             ? remoteOptions.getRemoteDefaultExecProperties()
             : ImmutableSortedMap.of();
 
-    if (spawn.getExecutionPlatform() == null
-        && spawn.getCombinedExecProperties().isEmpty()
-        && defaultExecProperties.isEmpty()
-        && additionalProperties.isEmpty()) {
+    if (!shouldProducePlatformProto(spawn, defaultExecProperties, additionalProperties)) {
+      // Execution platform is null or functionally empty
       return null;
     }
 

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<Label, List<String>> strategy : options.allowedStrategiesByExecPlatform) {
+      registryBuilder.addExecPlatformFilter(strategy.getKey(), strategy.getValue());
+    }
   }
 }

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

@@ -91,7 +91,9 @@ java_library(
     deps = [
         ":regex_filter_assignment_converter",
         "//src/main/java/com/google/devtools/build/lib/actions",
+        "//src/main/java/com/google/devtools/build/lib/analysis/config:core_option_converters",
         "//src/main/java/com/google/devtools/build/lib/analysis/config:per_label_options",
+        "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:cpu_resource_converter",
         "//src/main/java/com/google/devtools/build/lib/util:ram_resource_converter",
@@ -353,6 +355,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

@@ -17,6 +17,8 @@
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionContext.ShowSubcommands;
 import com.google.devtools.build.lib.analysis.config.PerLabelOptions;
+import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.analysis.config.CoreOptionConverters.LabelConverter;
 import com.google.devtools.build.lib.util.CpuResourceConverter;
 import com.google.devtools.build.lib.util.OptionsUtils;
 import com.google.devtools.build.lib.util.RamResourceConverter;
@@ -25,6 +27,7 @@
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.common.options.BoolOrEnumConverter;
 import com.google.devtools.common.options.Converters;
+import com.google.devtools.common.options.Converters.AssignmentToListOfValuesConverter;
 import com.google.devtools.common.options.Converters.CommaSeparatedNonEmptyOptionListConverter;
 import com.google.devtools.common.options.EnumConverter;
 import com.google.devtools.common.options.Option;
@@ -118,6 +121,31 @@ 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 = LabelToStringListConverter.class,
+      defaultValue = "null",
+      documentationCategory = OptionDocumentationCategory.EXECUTION_STRATEGY,
+      effectTags = {OptionEffectTag.EXECUTION},
+      help =
+          """
+          Filters spawn strategies by the execution platform without affecting order.
+          For example:
+          ```
+          common --spawn_strategy=remote,sandboxed,worker,local
+          common --strategy=Genrule=local
+          common --allowed_strategies_by_exec_platform=@platforms//host:host=local,sandboxed,worker
+          common --allowed_strategies_by_exec_platform=//:linux_amd64=remote
+          ```
+          With the above options;
+          - Actions configured for the host platform will be given `remote,sandboxed,worker`.
+          - Actions configured for the `//:linux_amd64` platform will be given `remote`.
+          - Actions configured for the `//:linux_amd64` platform with mnemonic `Genrule` will be
+            given no strategies and fail to spawn.
+          """)
+  public List<Map.Entry<Label, List<String>>> allowedStrategiesByExecPlatform;
+
   @Option(
       name = "materialize_param_files",
       defaultValue = "false",
@@ -641,4 +669,16 @@ public ShowSubcommandsConverter() {
           ShowSubcommands.class, "subcommand option", ShowSubcommands.TRUE, ShowSubcommands.FALSE);
     }
   }
+
+  private static class LabelToStringListConverter extends AssignmentToListOfValuesConverter<Label, String> {
+
+    public LabelToStringListConverter() {
+        super(new LabelConverter(), new Converters.StringConverter(), AllowEmptyKeys.NO);
+    }
+
+    @Override
+    public String getTypeDescription() {
+        return "a '<Label>=value[,value]' assignment";
+    }
+  }
 }

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,18 +80,21 @@ 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;
     this.remoteLocalFallbackStrategy = remoteLocalFallbackStrategy;
     logger.atInfo().log("Default strategies: %s", defaultStrategies);
-    logger.atInfo().log("Filter strategies: %s", strategyRegexFilter);
+    logger.atInfo().log("Regex filter strategies: %s", strategyRegexFilter);
+    logger.atInfo().log("Platform filter strategies: %s", strategyPlatformFilter);
     logger.atInfo().log("Mnemonic strategies: %s", mnemonicToStrategies);
     logger.atInfo().log("Remote strategies: %s", mnemonicToRemoteDynamicStrategies);
     logger.atInfo().log("Local strategies: %s", mnemonicToLocalDynamicStrategies);
@@ -105,7 +111,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 +123,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 not be used in Bazel.
    */
   public List<? extends SpawnStrategy> getStrategies(
       ActionExecutionMetadata resourceOwner, String mnemonic, @Nullable EventHandler reporter) {
@@ -154,18 +163,22 @@ 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("");
+      return strategyPlatformFilter.getStrategies(spawn, mnemonicToDynamicStrategies.get(""));
     }
     return ImmutableList.of();
   }
 
   @Nullable
   @Override
-  public AbstractSpawnStrategy getRemoteLocalFallbackStrategy() {
-    return remoteLocalFallbackStrategy;
+  public AbstractSpawnStrategy getRemoteLocalFallbackStrategy(Spawn spawn) {
+    var strategies = strategyPlatformFilter.getStrategies(spawn, Lists.newArrayList(remoteLocalFallbackStrategy));
+    if (strategies.isEmpty()) {
+      return null;
+    }
+    return strategies.getFirst();
   }
 
   /**
@@ -203,9 +216,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, Collection<SpawnStrategy>> entry :
+        strategyPlatformFilter.getFilterToStrategies().asMap().entrySet()) {
+      Collection<SpawnStrategy> value = entry.getValue();
+      logger.atInfo().log(
+          "FilterPlatformToStrategyImplementations: \"%s\" = [%s]",
+          entry.getKey().getCanonicalForm(), toImplementationNames(value));
+    }
 
     logger.atInfo().log(
         "DefaultStrategyImplementations: [%s]", toImplementationNames(defaultStrategies));
@@ -287,6 +307,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;
 
@@ -314,6 +335,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
@@ -444,6 +471,14 @@ public SpawnStrategyRegistry build() throws AbruptExitException {
         }
       }
 
+      ImmutableListMultimap.Builder<Label, SpawnStrategy> platformToStrategies = ImmutableListMultimap.builder();
+      for (Map.Entry<Label, List<String>> entry : execPlatformFilters.entrySet()) {
+        Label platform = entry.getKey();
+        platformToStrategies.putAll(
+          platform,
+          strategyMapper.toStrategies(entry.getValue(), "platform " + platform.getCanonicalForm()));
+      }
+
       ImmutableListMultimap.Builder<String, SpawnStrategy> mnemonicToStrategies =
           new ImmutableListMultimap.Builder<>();
       for (Map.Entry<String, List<String>> entry : mnemonicToIdentifiers.entrySet()) {
@@ -514,6 +549,7 @@ public SpawnStrategyRegistry build() throws AbruptExitException {
           mnemonicToStrategies.build(),
           new StrategyRegexFilter(
               strategyMapper, strategyPolicy, filterToIdentifiers, filterToStrategies),
+          new StrategyPlatformFilter(strategyMapper, platformToStrategies.build()),
           defaultStrategies,
           mnemonicToRemoteStrategies.build(),
           mnemonicToLocalStrategies.build(),
@@ -595,6 +631,64 @@ public String toString() {
     }
   }
 
+  private static class StrategyPlatformFilter {
+    private final StrategyMapper strategyMapper;
+    private final ImmutableListMultimap<Label, SpawnStrategy> platformToStrategies;
+
+    private StrategyPlatformFilter(
+      StrategyMapper strategyMapper,
+      ImmutableListMultimap<Label, SpawnStrategy> platformToStrategies) {
+      this.strategyMapper = strategyMapper;
+      this.platformToStrategies = platformToStrategies;
+    }
+
+    /**
+     * Gets strategies for the given spawn that are allowed by the execution platform.
+     * @param spawn Spawn to pick strategies for. Must have an execution platform.
+     * @param candidateStrategies Strategies ordered by priority to pick from. The contents of this
+     *                            list vary depending on the context but are always the result of
+     *                            an initial strategy selection pass.
+     *                            e.g. {@link SpawnStrategyRegistry#getStrategies(Spawn, EventHandler)},
+     *                            {@link SpawnStrategyRegistry#getDynamicSpawnActionContexts(Spawn, DynamicMode)}
+     *                            and {@link SpawnStrategyRegistry#getRemoteLocalFallbackStrategy(Spawn)}.
+     * @return A subset of {@code candidateStrategies} that are allowed by the spawn's execution
+     *         platform or all if no restrictions are in place.
+     *         Order from {@code candidateStrategies} is preserved.
+     */
+    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.");
+
+      if (platformToStrategies.containsKey(platformLabel)) {
+        var allowedStrategies = platformToStrategies.get(platformLabel);
+        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)));
+    }
+
+    ImmutableListMultimap<Label, SpawnStrategy> getFilterToStrategies() {
+      return platformToStrategies;
+    }
+
+    @Override
+    public String toString() {
+      return platformToStrategies.toString();
+    }
+  }
+
   /* Maps the strategy identifier (e.g. "local", "worker"..) to the real strategy. */
   private static class StrategyMapper {