Clear all remote metadata if any of them are evicted from remote cache

src/main/java/com/google/devtools/build/lib/actions/cache/ActionCache.java

@@ -45,6 +45,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Predicate;
 import javax.annotation.Nullable;
 
 /**
@@ -76,6 +77,11 @@ public interface ActionCache {
    */
   void remove(String key);
 
+  /**
+   * Removes entry from cache that matches the predicate.
+   */
+  void removeIf(Predicate<ActionCache.Entry> predicate);
+
   /**
    * An entry in the ActionCache that contains all action input and output
    * artifact paths and their metadata plus action key itself.
@@ -246,7 +252,8 @@ public RemoteFileArtifactValue getOutputFile(Artifact output) {
       return outputFileMetadata.get(output.getExecPathString());
     }
 
-    Map<String, RemoteFileArtifactValue> getOutputFiles() {
+    /** Gets metadata of all output files */
+    public Map<String, RemoteFileArtifactValue> getOutputFiles() {
       return outputFileMetadata;
     }
 
@@ -273,7 +280,8 @@ public SerializableTreeArtifactValue getOutputTree(SpecialArtifact output) {
       return outputTreeMetadata.get(output.getExecPathString());
     }
 
-    Map<String, SerializableTreeArtifactValue> getOutputTrees() {
+    /** Gets metadata of all output trees */
+    public Map<String, SerializableTreeArtifactValue> getOutputTrees() {
       return outputTreeMetadata;
     }
 

src/main/java/com/google/devtools/build/lib/actions/cache/CompactPersistentActionCache.java

@@ -55,6 +55,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Predicate;
 import javax.annotation.Nullable;
 
 /**
@@ -342,6 +343,10 @@ public ActionCache.Entry get(String key) {
       return null;
     }
     byte[] data = map.get(index);
+    return get(data);
+  }
+
+  private ActionCache.Entry get(byte[] data) {
     try {
       return data != null ? decode(indexer, data) : null;
     } catch (IOException e) {
@@ -381,6 +386,11 @@ public void remove(String key) {
     map.remove(indexer.getIndex(key));
   }
 
+  @Override
+  public void removeIf(Predicate<Entry> predicate) {
+    map.entrySet().removeIf(entry -> predicate.test(get(entry.getValue())));
+  }
+
   @ThreadSafety.ThreadHostile
   @Override
   public long save() throws IOException {

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

@@ -233,6 +233,10 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions",
         "//src/main/java/com/google/devtools/build/lib/actions:action_lookup_data",
         "//src/main/java/com/google/devtools/build/lib/actions:artifacts",
+        "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:action_execution_value",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:tree_artifact_value",
         "//src/main/java/com/google/devtools/build/skyframe",
         "//third_party:jsr305",
     ],

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

@@ -13,15 +13,13 @@
 // limitations under the License.
 package com.google.devtools.build.lib.remote;
 
-import com.google.devtools.build.lib.actions.Action;
-import com.google.devtools.build.lib.actions.ActionCacheUtils;
 import com.google.devtools.build.lib.actions.ActionInput;
-import com.google.devtools.build.lib.actions.ActionLookupData;
-import com.google.devtools.build.lib.actions.ActionLookupValue;
-import com.google.devtools.build.lib.actions.Artifact;
+import com.google.devtools.build.lib.actions.FileArtifactValue;
 import com.google.devtools.build.lib.actions.cache.ActionCache;
+import com.google.devtools.build.lib.skyframe.ActionExecutionValue;
+import com.google.devtools.build.lib.skyframe.SkyFunctions;
+import com.google.devtools.build.lib.skyframe.TreeArtifactValue;
 import com.google.devtools.build.skyframe.MemoizingEvaluator;
-import java.util.HashMap;
 import java.util.Set;
 import javax.annotation.Nullable;
 
@@ -41,36 +39,39 @@ public void handleMissingInputs(Set<ActionInput> missingActionInputs) {
       return;
     }
 
-    var actions = new HashMap<ActionLookupData, Action>();
+    // If any outputs are evicted, remove all remote metadata from skyframe and local action cache.
 
-    try {
-      for (ActionInput actionInput : missingActionInputs) {
-        if (actionInput instanceof Artifact.DerivedArtifact) {
-          Artifact.DerivedArtifact output = (Artifact.DerivedArtifact) actionInput;
-          ActionLookupData actionLookupData = output.getGeneratingActionKey();
-          var actionLookupValue =
-              memoizingEvaluator.getExistingValue(actionLookupData.getActionLookupKey());
-          if (actionLookupValue instanceof ActionLookupValue) {
-            Action action =
-                ((ActionLookupValue) actionLookupValue)
-                    .getAction(actionLookupData.getActionIndex());
-            actions.put(actionLookupData, action);
+    memoizingEvaluator.delete(
+        key -> {
+          if (key.functionName().equals(SkyFunctions.ACTION_EXECUTION)) {
+            try {
+              var value = memoizingEvaluator.getExistingValue(key);
+              return value instanceof ActionExecutionValue
+                  && isRemote((ActionExecutionValue) value);
+            } catch (InterruptedException ignored) {
+              return false;
+            }
           }
-        }
-      }
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
+          return false;
+        });
+
+    if (actionCache != null) {
+      actionCache.removeIf(entry ->
+          !entry.getOutputFiles().isEmpty() || !entry.getOutputTrees().isEmpty());
     }
+  }
 
-    if (!actions.isEmpty()) {
-      var actionKeys = actions.keySet();
-      memoizingEvaluator.delete(key -> key instanceof ActionLookupData && actionKeys.contains(key));
+  private boolean isRemote(ActionExecutionValue value) {
+    return value.getAllFileValues().values().stream().anyMatch(FileArtifactValue::isRemote)
+        || value.getAllTreeArtifactValues().values().stream().anyMatch(this::isRemoteTree);
+  }
 
-      if (actionCache != null) {
-        for (var action : actions.values()) {
-          ActionCacheUtils.removeCacheEntry(actionCache, action);
-        }
-      }
-    }
+  private boolean isRemoteTree(TreeArtifactValue treeArtifactValue) {
+    return treeArtifactValue.getChildValues().values().stream()
+            .anyMatch(FileArtifactValue::isRemote)
+        || treeArtifactValue
+            .getArchivedRepresentation()
+            .map(ar -> ar.archivedFileValue().isRemote())
+            .orElse(false);
   }
 }

src/test/java/com/google/devtools/build/lib/actions/ActionCacheCheckerTest.java

@@ -73,6 +73,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Predicate;
 import javax.annotation.Nullable;
 import org.junit.After;
 import org.junit.Before;
@@ -1255,6 +1256,11 @@ public void remove(String key) {
       delegate.remove(key);
     }
 
+    @Override
+    public void removeIf(Predicate<Entry> predicate) {
+      delegate.removeIf(predicate);
+    }
+
     @Override
     public long save() throws IOException {
       return delegate.save();

src/test/java/com/google/devtools/build/lib/actions/util/ActionCacheTestHelper.java

@@ -17,6 +17,7 @@
 import com.google.devtools.build.lib.actions.cache.Protos.ActionCacheStatistics;
 import com.google.devtools.build.lib.actions.cache.Protos.ActionCacheStatistics.MissReason;
 import java.io.PrintStream;
+import java.util.function.Predicate;
 
 /**
  * Utilities for tests that use the action cache.
@@ -38,6 +39,9 @@ public Entry get(String fingerprint) {
         @Override
         public void remove(String key) {}
 
+        @Override
+        public void removeIf(Predicate<Entry> predicate) {}
+
         @Override
         public long save() {
           return -1;

src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java

@@ -580,6 +580,11 @@ public synchronized void remove(String key) {
       actionCache.remove(key);
     }
 
+    @Override
+    public void removeIf(java.util.function.Predicate<Entry> predicate) {
+      actionCache.values().removeIf(predicate);
+    }
+
     public synchronized void reset() {
       actionCache.clear();
     }