Basic strategy-policy implementation

Basic foundation for limiting strategies used in bazel a la invocation policy.
Implemented as an explicit policy for fine-grained control, though for now we
just implement an allowlist (in the future it should be extensible enough for
more creativity).

Naming of the fields in the proto definition is biased towards more ergonomic
textproto configuration.

Alternatives considered include extending invocation policy to support
key/value flags, but this seemed like too much power being added to the flag
system for too narrow of a use case.

Note that this policy applies universally, regardless of whether or not
strategies come by flag or are hard-coded in bazel. This is intentional - it
allows *full* control over where/how bazel is executing actions.

--strategy_regexp is presently ignored, it requires larger changes to implement, will be done in the next CLs.

PiperOrigin-RevId: 635733824
Change-Id: Ic43cf76c2b58b0fcce089119b8ce43a0116c8490

src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java

@@ -162,7 +162,8 @@ public class ExecutionTool {
     ExecutorBuilder executorBuilder = new ExecutorBuilder();
     ModuleActionContextRegistry.Builder actionContextRegistryBuilder =
         ModuleActionContextRegistry.builder();
-    SpawnStrategyRegistry.Builder spawnStrategyRegistryBuilder = SpawnStrategyRegistry.builder();
+    SpawnStrategyRegistry.Builder spawnStrategyRegistryBuilder =
+        SpawnStrategyRegistry.builder(env.getInvocationPolicy().getStrategyPolicy());
     actionContextRegistryBuilder.register(SpawnStrategyResolver.class, new SpawnStrategyResolver());
     actionContextRegistryBuilder.register(
         ImportantOutputHandler.class, ImportantOutputHandler.NO_OP);

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

@@ -337,18 +337,29 @@ java_library(
     ],
 )
 
+java_library(
+    name = "spawn_strategy_policy",
+    srcs = ["SpawnStrategyPolicy.java"],
+    deps = [
+        "//src/main/protobuf:strategy_policy_java_proto",
+        "//third_party:guava",
+    ],
+)
+
 java_library(
     name = "spawn_strategy_registry",
     srcs = ["SpawnStrategyRegistry.java"],
     deps = [
         ":abstract_spawn_strategy",
         ":remote_local_fallback_registry",
+        ":spawn_strategy_policy",
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//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/util:detailed_exit_code",
         "//src/main/protobuf:failure_details_java_proto",
+        "//src/main/protobuf:strategy_policy_java_proto",
         "//third_party:auto_value",
         "//third_party:flogger",
         "//third_party:guava",

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

@@ -0,0 +1,80 @@
+// Copyright 2022 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.exec;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.runtime.proto.MnemonicPolicy;
+import com.google.devtools.build.lib.runtime.proto.StrategiesForMnemonic;
+import java.util.List;
+
+/** Policy for filtering spawn strategies. */
+public interface SpawnStrategyPolicy {
+
+  /** Returns result of applying policy to per-mnemonic strategies. */
+  ImmutableList<String> apply(String mnemonic, List<String> strategies);
+
+  /** Creates new policy from proto descriptor. Empty proto policy implies everything allowed. */
+  static SpawnStrategyPolicy create(MnemonicPolicy policy) {
+    if (MnemonicPolicy.getDefaultInstance().equals(policy)) {
+      return new AllowAllStrategiesPolicy();
+    }
+
+    ImmutableMap.Builder<String, ImmutableSet<String>> perMnemonicAllowList =
+        ImmutableMap.builder();
+    for (StrategiesForMnemonic strategiesForMnemonic : policy.getStrategyAllowlistList()) {
+      perMnemonicAllowList.put(
+          strategiesForMnemonic.getMnemonic(),
+          ImmutableSet.copyOf(strategiesForMnemonic.getStrategyList()));
+    }
+    return new SpawnStrategyPolicyImpl(
+        perMnemonicAllowList.buildKeepingLast(),
+        ImmutableSet.copyOf(policy.getDefaultAllowlistList()));
+  }
+
+  /** Allows all strategies - effectively a no-op strategy. */
+  class AllowAllStrategiesPolicy implements SpawnStrategyPolicy {
+
+    private AllowAllStrategiesPolicy() {}
+
+    @Override
+    public ImmutableList<String> apply(String mnemonic, List<String> strategies) {
+      return ImmutableList.copyOf(strategies);
+    }
+  }
+
+  /** Enforces a real strategy policy based on provided config. */
+  class SpawnStrategyPolicyImpl implements SpawnStrategyPolicy {
+
+    private final ImmutableMap<String, ImmutableSet<String>> perMnemonicAllowList;
+    private final ImmutableSet<String> defaultAllowList;
+
+    private SpawnStrategyPolicyImpl(
+        ImmutableMap<String, ImmutableSet<String>> perMnemonicAllowList,
+        ImmutableSet<String> defaultAllowList) {
+      this.perMnemonicAllowList = perMnemonicAllowList;
+      this.defaultAllowList = defaultAllowList;
+    }
+
+    @Override
+    public ImmutableList<String> apply(String mnemonic, List<String> strategies) {
+      ImmutableSet<String> allowList =
+          perMnemonicAllowList.getOrDefault(mnemonic, defaultAllowList);
+      return strategies.stream().filter(allowList::contains).collect(toImmutableList());
+    }
+  }
+}

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

@@ -36,6 +36,8 @@
 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.runtime.proto.MnemonicPolicy;
+import com.google.devtools.build.lib.runtime.proto.StrategyPolicy;
 import com.google.devtools.build.lib.server.FailureDetails;
 import com.google.devtools.build.lib.server.FailureDetails.ExecutionOptions.Code;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
@@ -61,6 +63,9 @@ public final class SpawnStrategyRegistry
     implements DynamicStrategyRegistry, ActionContext, RemoteLocalFallbackRegistry {
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
+  private static final SpawnStrategyPolicy ALLOW_ALL_STRATEGIES =
+      SpawnStrategyPolicy.create(MnemonicPolicy.getDefaultInstance());
+
   private final ImmutableListMultimap<String, SpawnStrategy> mnemonicToStrategies;
   private final ImmutableListMultimap<RegexFilter, SpawnStrategy> filterToStrategies;
   private final ImmutableList<? extends SpawnStrategy> defaultStrategies;
@@ -232,8 +237,20 @@ private static String toImplementationNames(Collection<?> strategies) {
   }
 
   /** Returns a new {@link Builder} suitable for creating instances of SpawnStrategyRegistry. */
+  @VisibleForTesting
   public static Builder builder() {
-    return new Builder();
+    return new Builder(
+        /* strategyPolicy= */ ALLOW_ALL_STRATEGIES,
+        /* dynamicRemotePolicy= */ ALLOW_ALL_STRATEGIES,
+        /* dynamicLocalPolicy= */ ALLOW_ALL_STRATEGIES);
+  }
+
+  /** Returns a new {@link Builder} suitable for creating instances of SpawnStrategyRegistry. */
+  public static Builder builder(StrategyPolicy strategyPolicyProto) {
+    return new Builder(
+        SpawnStrategyPolicy.create(strategyPolicyProto.getMnemonicPolicy()),
+        SpawnStrategyPolicy.create(strategyPolicyProto.getDynamicRemotePolicy()),
+        SpawnStrategyPolicy.create(strategyPolicyProto.getDynamicLocalPolicy()));
   }
 
   /**
@@ -257,6 +274,9 @@ public static final class Builder {
 
     private ImmutableList<String> explicitDefaultStrategies = ImmutableList.of();
 
+    private final SpawnStrategyPolicy strategyPolicy;
+    private final SpawnStrategyPolicy dynamicRemotePolicy;
+    private final SpawnStrategyPolicy dynamicLocalPolicy;
     // TODO(schmitt): Using a list and autovalue so as to be able to reverse order while legacy sort
     //  is supported. Can be converted to same as mnemonics once legacy behavior is removed.
     private final List<FilterAndIdentifiers> filterAndIdentifiers = new ArrayList<>();
@@ -269,6 +289,15 @@ public static final class Builder {
 
     @Nullable private String remoteLocalFallbackStrategyIdentifier;
 
+    private Builder(
+        SpawnStrategyPolicy strategyPolicy,
+        SpawnStrategyPolicy dynamicRemotePolicy,
+        SpawnStrategyPolicy dynamicLocalPolicy) {
+      this.strategyPolicy = strategyPolicy;
+      this.dynamicRemotePolicy = dynamicRemotePolicy;
+      this.dynamicLocalPolicy = dynamicLocalPolicy;
+    }
+
     /**
      * Adds a filter limiting any spawn whose {@linkplain
      * com.google.devtools.build.lib.actions.ActionExecutionMetadata#getProgressMessage() owner's
@@ -292,17 +321,45 @@ public Builder addDescriptionFilter(RegexFilter filter, List<String> identifiers
      * command-line identifiers, in order.
      *
      * <p>If the same mnemonic is registered multiple times the last such call will take precedence.
+     * Or in other words, last one wins.
      *
      * <p>Note that if a spawn matches a {@linkplain #addDescriptionFilter registered description
      * filter} that filter will take precedence over any mnemonic-based filters.
      */
-    // last one wins
     @CanIgnoreReturnValue
     public Builder addMnemonicFilter(String mnemonic, List<String> identifiers) {
       mnemonicToIdentifiers.put(mnemonic, identifiers);
       return this;
     }
 
+    /**
+     * Sets the strategy names to use in the remote branch of dynamic execution for a set of action
+     * mnemonics.
+     *
+     * <p>During execution, each strategy is {@linkplain SpawnStrategy#canExec(Spawn,
+     * ActionContextRegistry) asked} whether it can execute a given Spawn. The first strategy in the
+     * list that says so will get the job.
+     */
+    @CanIgnoreReturnValue
+    public Builder addDynamicRemoteStrategies(Map<String, List<String>> strategies) {
+      mnemonicToRemoteDynamicIdentifiers.putAll(strategies);
+      return this;
+    }
+
+    /**
+     * Sets the strategy names to use in the local branch of dynamic execution for a number of
+     * action mnemonics.
+     *
+     * <p>During execution, each strategy is {@linkplain SpawnStrategy#canExec(Spawn,
+     * ActionContextRegistry) asked} whether it can execute a given Spawn. The first strategy in the
+     * list that says so will get the job.
+     */
+    @CanIgnoreReturnValue
+    public Builder addDynamicLocalStrategies(Map<String, List<String>> strategies) {
+      mnemonicToLocalDynamicIdentifiers.putAll(strategies);
+      return this;
+    }
+
     /**
      * Registers a strategy implementation with this collector, distinguishing it from other
      * strategies with the given command-line identifiers (of which at least one is required).
@@ -348,34 +405,6 @@ public Builder resetDefaultStrategies() {
       return this;
     }
 
-    /**
-     * Sets the strategy names to use in the remote branch of dynamic execution for a set of action
-     * mnemonics.
-     *
-     * <p>During execution, each strategy is {@linkplain SpawnStrategy#canExec(Spawn,
-     * ActionContextRegistry) asked} whether it can execute a given Spawn. The first strategy in the
-     * list that says so will get the job.
-     */
-    @CanIgnoreReturnValue
-    public Builder addDynamicRemoteStrategies(Map<String, List<String>> strategies) {
-      mnemonicToRemoteDynamicIdentifiers.putAll(strategies);
-      return this;
-    }
-
-    /**
-     * Sets the strategy names to use in the local branch of dynamic execution for a number of
-     * action mnemonics.
-     *
-     * <p>During execution, each strategy is {@linkplain SpawnStrategy#canExec(Spawn,
-     * ActionContextRegistry) asked} whether it can execute a given Spawn. The first strategy in the
-     * list that says so will get the job.
-     */
-    @CanIgnoreReturnValue
-    public Builder addDynamicLocalStrategies(Map<String, List<String>> strategies) {
-      mnemonicToLocalDynamicIdentifiers.putAll(strategies);
-      return this;
-    }
-
     /**
      * Sets the commandline identifier of the strategy to be used when falling back from remote to
      * local execution.
@@ -412,24 +441,31 @@ public SpawnStrategyRegistry build() throws AbruptExitException {
       ImmutableListMultimap.Builder<String, SpawnStrategy> mnemonicToStrategies =
           new ImmutableListMultimap.Builder<>();
       for (Map.Entry<String, List<String>> entry : mnemonicToIdentifiers.entrySet()) {
+        String mnemonic = entry.getKey();
+        ImmutableList<String> sanitizedStrategies =
+            strategyPolicy.apply(mnemonic, entry.getValue());
         mnemonicToStrategies.putAll(
-            entry.getKey(), toStrategies(entry.getValue(), "mnemonic " + entry.getKey()));
+            mnemonic, toStrategies(sanitizedStrategies, "mnemonic " + mnemonic));
       }
 
       ImmutableListMultimap.Builder<String, SandboxedSpawnStrategy> mnemonicToLocalStrategies =
           new ImmutableListMultimap.Builder<>();
       for (Map.Entry<String, List<String>> entry : mnemonicToLocalDynamicIdentifiers.entrySet()) {
+        String mnemonic = entry.getKey();
+        ImmutableList<String> sanitizedStrategies =
+            dynamicLocalPolicy.apply(mnemonic, entry.getValue());
         mnemonicToLocalStrategies.putAll(
-            entry.getKey(),
-            toSandboxedStrategies(entry.getValue(), "local mnemonic " + entry.getKey()));
+            mnemonic, toSandboxedStrategies(sanitizedStrategies, "local mnemonic " + mnemonic));
       }
 
       ImmutableListMultimap.Builder<String, SandboxedSpawnStrategy> mnemonicToRemoteStrategies =
           new ImmutableListMultimap.Builder<>();
       for (Map.Entry<String, List<String>> entry : mnemonicToRemoteDynamicIdentifiers.entrySet()) {
+        String mnemonic = entry.getKey();
+        ImmutableList<String> sanitizedStrategies =
+            dynamicRemotePolicy.apply(mnemonic, entry.getValue());
         mnemonicToRemoteStrategies.putAll(
-            entry.getKey(),
-            toSandboxedStrategies(entry.getValue(), "remote mnemonic " + entry.getKey()));
+            mnemonic, toSandboxedStrategies(sanitizedStrategies, "remote mnemonic " + mnemonic));
       }
 
       AbstractSpawnStrategy remoteLocalFallbackStrategy = null;
@@ -515,10 +551,6 @@ private Iterable<? extends SandboxedSpawnStrategy> toSandboxedStrategies(
     }
   }
 
-  public ImmutableListMultimap<String, SpawnStrategy> getMnemonicToStrategies() {
-    return mnemonicToStrategies;
-  }
-
   private static AbruptExitException createExitException(String message, Code detailedCode) {
     return new AbruptExitException(
         DetailedExitCode.of(

src/main/java/com/google/devtools/build/lib/runtime/BlazeCommandDispatcher.java

@@ -348,6 +348,7 @@ private BlazeCommandResult execExclusively(
         workspace.initCommand(
             commandAnnotation,
             options,
+            invocationPolicy,
             commandEnvWarnings,
             waitTimeInMs,
             firstContactTime,

src/main/java/com/google/devtools/build/lib/runtime/BlazeWorkspace.java

@@ -31,6 +31,7 @@
 import com.google.devtools.build.lib.profiler.AutoProfiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.memory.AllocationTracker;
+import com.google.devtools.build.lib.runtime.proto.InvocationPolicyOuterClass.InvocationPolicy;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import com.google.devtools.build.lib.util.io.CommandExtensionReporter;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
@@ -205,6 +206,7 @@ public Range<Long> getLastExecutionTimeRange() {
   public CommandEnvironment initCommand(
       Command command,
       OptionsParsingResult options,
+      InvocationPolicy invocationPolicy,
       List<String> warnings,
       long waitTimeInMs,
       long commandStartTime,
@@ -221,6 +223,7 @@ public CommandEnvironment initCommand(
             Thread.currentThread(),
             command,
             options,
+            invocationPolicy,
             syscallCache,
             quiescingExecutors,
             warnings,

src/main/java/com/google/devtools/build/lib/runtime/CommandEnvironment.java

@@ -108,6 +108,7 @@ public class CommandEnvironment {
   private final Thread commandThread;
   private final Command command;
   private final OptionsParsingResult options;
+  private final InvocationPolicy invocationPolicy;
   private final PathPackageLocator packageLocator;
   private final Path workingDirectory;
   private final PathFragment relativeWorkingDirectory;
@@ -193,6 +194,7 @@ public void exit(AbruptExitException exception) {
       Thread commandThread,
       Command command,
       OptionsParsingResult options,
+      InvocationPolicy invocationPolicy,
       SyscallCache syscallCache,
       QuiescingExecutors quiescingExecutors,
       List<String> warnings,
@@ -212,6 +214,7 @@ public void exit(AbruptExitException exception) {
     this.commandThread = commandThread;
     this.command = command;
     this.options = options;
+    this.invocationPolicy = invocationPolicy;
     this.shutdownReasonConsumer = shutdownReasonConsumer;
     this.syscallCache = syscallCache;
     this.quiescingExecutors = quiescingExecutors;
@@ -448,6 +451,10 @@ public OptionsParsingResult getOptions() {
     return options;
   }
 
+  public InvocationPolicy getInvocationPolicy() {
+    return invocationPolicy;
+  }
+
   public void setInvocationPolicyFlags(ImmutableList<OptionAndRawValue> invocationPolicyFlags) {
     this.invocationPolicyFlags = invocationPolicyFlags;
   }

src/main/java/com/google/devtools/build/lib/runtime/commands/CanonicalizeCommand.java

@@ -62,9 +62,9 @@
     mustRunInWorkspace = false,
     shortDescription = "Canonicalizes a list of %{product} options.",
     help =
-        "This command canonicalizes a list of %{product} options. Don't forget to prepend "
-            + " '--' to end option parsing before the flags to canonicalize.\n"
-            + "%{options}")
+        "This command canonicalizes a list of %{product} options. Don't forget to prepend  '--' to"
+            + " end option parsing before the flags to canonicalize. This command doesn't support"
+            + " the strategy policies under --invocation_policy flag.\n%{options}")
 public final class CanonicalizeCommand implements BlazeCommand {
 
   public static class Options extends OptionsBase {