From 699bed35c6e918352a02586ed1dbc07953a77571 Mon Sep 17 00:00:00 2001
From: Anton Lopyrev <anton.lopyrev@snapchat.com>
Date: Thu, 2 Apr 2020 17:49:07 -0700
Subject: [PATCH] Hack to specificy the cache strategy per action mnemonic via
 command line (#13)

* [WIP] Hack to specificy the cache strategy per action mnemonic via command line

* Reset cache strategy
---
 .../devtools/build/lib/actions/Spawns.java    | 43 ++++++++++++++++---
 .../lib/bazel/rules/BazelStrategyModule.java  |  6 +++
 .../build/lib/exec/ExecutionOptions.java      | 11 +++++
 3 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/actions/Spawns.java b/src/main/java/com/google/devtools/build/lib/actions/Spawns.java
index 06e4b613497bc7..d853b78e4fa24c 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/Spawns.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/Spawns.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.lib.actions;
 
+import com.google.common.collect.LinkedHashMultimap;
 import com.google.devtools.build.lib.server.FailureDetails;
 import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
 import com.google.devtools.build.lib.server.FailureDetails.Spawn.Code;
@@ -23,25 +24,57 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 
 /** Helper methods relating to implementations of {@link Spawn}. */
 public final class Spawns {
   private Spawns() {}
 
+  /**
+   * This is a dirty hack by Anton. Sorry.
+   */
+  private static final LinkedHashMultimap<String, String> sCacheStrategyByMnemonicMap =
+      LinkedHashMultimap.create();
+
+  public static void resetCacheStategy() {
+    sCacheStrategyByMnemonicMap.clear();
+  }
+
+  /**
+   * Sets the cache strategy names for a given action mnemonic.
+   */
+  public static void addCacheStrategyByMnemonic(String mnemonic, List<String> strategies) {
+    sCacheStrategyByMnemonicMap.replaceValues(mnemonic, strategies);
+  }
+
   /**
    * Returns {@code true} if the result of {@code spawn} may be cached.
    */
   public static boolean mayBeCached(Spawn spawn) {
-    return !spawn.getExecutionInfo().containsKey(ExecutionRequirements.NO_CACHE)
-        && !spawn.getExecutionInfo().containsKey(ExecutionRequirements.LOCAL);
+    if (!spawn.getExecutionInfo().containsKey(ExecutionRequirements.LOCAL)) {
+      if (sCacheStrategyByMnemonicMap.containsKey(spawn.getMnemonic())) {
+        return !sCacheStrategyByMnemonicMap.get(spawn.getMnemonic()).contains(ExecutionRequirements.NO_CACHE);
+      } else {
+        return !spawn.getExecutionInfo().containsKey(ExecutionRequirements.NO_CACHE);
+      }
+    } else {
+      return false;
+    }
   }
 
   /** Returns {@code true} if the result of {@code spawn} may be cached remotely. */
   public static boolean mayBeCachedRemotely(Spawn spawn) {
-    return mayBeCached(spawn)
-        && !spawn.getExecutionInfo().containsKey(ExecutionRequirements.NO_REMOTE)
-        && !spawn.getExecutionInfo().containsKey(ExecutionRequirements.NO_REMOTE_CACHE);
+    if (mayBeCached(spawn)
+        && !spawn.getExecutionInfo().containsKey(ExecutionRequirements.NO_REMOTE)) {
+      if (sCacheStrategyByMnemonicMap.containsKey(spawn.getMnemonic())) {
+        return !sCacheStrategyByMnemonicMap.get(spawn.getMnemonic()).contains(ExecutionRequirements.NO_REMOTE_CACHE);
+      } else {
+        return !spawn.getExecutionInfo().containsKey(ExecutionRequirements.NO_REMOTE_CACHE);
+      }
+    } else {
+      return false;
+    }
   }
 
   /** Returns {@code true} if {@code spawn} may be executed remotely. */
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelStrategyModule.java b/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelStrategyModule.java
index f503180cada615..fe20d4ca5f7cca 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelStrategyModule.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/rules/BazelStrategyModule.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.bazel.rules;
 
 import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.actions.Spawns;
 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;
@@ -64,6 +65,11 @@ public void registerSpawnStrategies(
     ExecutionOptions options = env.getOptions().getOptions(ExecutionOptions.class);
     RemoteOptions remoteOptions = env.getOptions().getOptions(RemoteOptions.class);
 
+    Spawns.resetCacheStategy();
+    for (Map.Entry<String, List<String>> strategy : options.cacheStrategy) {
+      Spawns.addCacheStrategyByMnemonic(strategy.getKey(), strategy.getValue());
+    }
+
     List<String> spawnStrategies = new ArrayList<>(options.spawnStrategy);
 
     if (spawnStrategies.isEmpty()) {
diff --git a/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java b/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java
index a784b30c0a0d7e..5bfa0c5888bef3 100644
--- a/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java
+++ b/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java
@@ -59,6 +59,17 @@ public class ExecutionOptions extends OptionsBase {
 
   public static final ExecutionOptions DEFAULTS = Options.getDefaults(ExecutionOptions.class);
 
+  @Option(
+      name = "cache_strategy",
+      allowMultiple = true,
+      converter = Converters.StringToStringListConverter.class,
+      defaultValue = "null",
+      documentationCategory = OptionDocumentationCategory.UNCATEGORIZED,
+      effectTags = {OptionEffectTag.UNKNOWN},
+      help =
+          "Specify the cache policy for specific actions by mnemonic")
+  public List<Map.Entry<String, List<String>>> cacheStrategy;
+
   @Option(
       name = "spawn_strategy",
       defaultValue = "",