Add functionality to make certain SkyValues unshareable, meaning they are not serialized. Tag TestCompletionValue and any ActionExecutionValue coming from a NotifyOnActionCacheHit (i.e., tests) like that. To make such values really not shared, request the ActionExecutionValue from TestCompletionFunction as opposed to the ArtifactValue (propagating the unshareable bit up seemed like too much fuss, and I have a dream of getting rid of ArtifactValue anyway).

PiperOrigin-RevId: 200504358
diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionLookupData.java b/src/main/java/com/google/devtools/build/lib/actions/ActionLookupData.java
index 00aa657..b456de9 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionLookupData.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionLookupData.java
@@ -19,6 +19,7 @@
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.concurrent.BlazeInterners;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
+import com.google.devtools.build.skyframe.ShareabilityOfValue;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 
@@ -26,7 +27,9 @@
 @AutoCodec
 public class ActionLookupData implements SkyKey {
   private static final Interner<ActionLookupData> INTERNER = BlazeInterners.newWeakInterner();
-  public static final SkyFunctionName NAME = SkyFunctionName.create("ACTION_EXECUTION");
+  // Test actions are not shareable.
+  public static final SkyFunctionName NAME =
+      SkyFunctionName.create("ACTION_EXECUTION", ShareabilityOfValue.SOMETIMES);
 
   private final ActionLookupKey actionLookupKey;
   private final int actionIndex;
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
index 612632e..6f5044f 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
@@ -105,10 +105,6 @@
     int actionIndex = actionLookupData.getActionIndex();
     Action action = actionLookupValue.getAction(actionIndex);
     skyframeActionExecutor.noteActionEvaluationStarted(actionLookupData, action);
-    // TODO(bazel-team): Non-volatile NotifyOnActionCacheHit actions perform worse in Skyframe than
-    // legacy when they are not at the top of the action graph. In legacy, they are stored
-    // separately, so notifying non-dirty actions is cheap. In Skyframe, they depend on the
-    // BUILD_ID, forcing invalidation of upward transitive closure on each build.
     if ((action.isVolatile() && !(action instanceof SkyframeAwareAction))
         || action instanceof NotifyOnActionCacheHit) {
       // Volatile build actions may need to execute even if none of their known inputs have changed.
@@ -413,11 +409,12 @@
           "Error, we're not re-executing a "
               + "SkyframeAwareAction which should be re-executed unconditionally. Action: %s",
           action);
-      return new ActionExecutionValue(
+      return ActionExecutionValue.create(
           metadataHandler.getOutputArtifactData(),
           metadataHandler.getOutputTreeArtifactData(),
           metadataHandler.getAdditionalOutputData(),
-          /*outputSymlinks=*/ null);
+          /*outputSymlinks=*/ null,
+          action instanceof NotifyOnActionCacheHit);
     }
 
     // Delete the metadataHandler's cache of the action's outputs, since they are being deleted.
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
index 687a787..85664fd 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionValue.java
@@ -31,6 +31,7 @@
 import com.google.devtools.build.lib.actions.FilesetOutputSymlink;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.ThreadSafe;
+import com.google.devtools.build.lib.skyframe.serialization.UnshareableValue;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
@@ -86,7 +87,7 @@
    *     ActionExecutionValues.
    * @param outputSymlinks This represents the SymlinkTree which is the output of a fileset action.
    */
-  ActionExecutionValue(
+  private ActionExecutionValue(
       Map<Artifact, FileValue> artifactData,
       Map<Artifact, TreeArtifactValue> treeArtifactData,
       Map<Artifact, FileArtifactValue> additionalOutputData,
@@ -97,6 +98,19 @@
     this.outputSymlinks = outputSymlinks;
   }
 
+  static ActionExecutionValue create(
+      Map<Artifact, FileValue> artifactData,
+      Map<Artifact, TreeArtifactValue> treeArtifactData,
+      Map<Artifact, FileArtifactValue> additionalOutputData,
+      @Nullable ImmutableList<FilesetOutputSymlink> outputSymlinks,
+      boolean notifyOnActionCacheHitAction) {
+    return notifyOnActionCacheHitAction
+        ? new CrossServerUnshareableActionExecutionValue(
+            artifactData, treeArtifactData, additionalOutputData, outputSymlinks)
+        : new ActionExecutionValue(
+            artifactData, treeArtifactData, additionalOutputData, outputSymlinks);
+  }
+
   /**
    * Returns metadata for a given artifact, if that metadata cannot be inferred from the
    * corresponding {@link #getData} call for that Artifact. See {@link
@@ -153,7 +167,7 @@
    */
   @ThreadSafe
   @VisibleForTesting
-  public static SkyKey key(ActionLookupValue.ActionLookupKey lookupKey, int index) {
+  public static ActionLookupData key(ActionLookupValue.ActionLookupKey lookupKey, int index) {
     return ActionLookupData.create(lookupKey, index);
   }
 
@@ -171,7 +185,10 @@
     if (this == obj) {
       return true;
     }
-    if (!(obj instanceof ActionExecutionValue)) {
+    if (obj == null) {
+      return false;
+    }
+    if (!obj.getClass().equals(getClass())) {
       return false;
     }
     ActionExecutionValue o = (ActionExecutionValue) obj;
@@ -201,6 +218,21 @@
     return value;
   }
 
+  /**
+   * Marker subclass that indicates this value cannot be shared across servers. Note that this is
+   * unrelated to the concept of shared actions.
+   */
+  private static class CrossServerUnshareableActionExecutionValue extends ActionExecutionValue
+      implements UnshareableValue {
+    CrossServerUnshareableActionExecutionValue(
+        Map<Artifact, FileValue> artifactData,
+        Map<Artifact, TreeArtifactValue> treeArtifactData,
+        Map<Artifact, FileArtifactValue> additionalOutputData,
+        @Nullable ImmutableList<FilesetOutputSymlink> outputSymlinks) {
+      super(artifactData, treeArtifactData, additionalOutputData, outputSymlinks);
+    }
+  }
+
   private static <V> ImmutableMap<Artifact, V> transformKeys(
       ImmutableMap<Artifact, V> data, Map<OwnerlessArtifactWrapper, Artifact> newArtifactMap) {
     if (data.isEmpty()) {
@@ -223,10 +255,11 @@
             .collect(Collectors.toMap(OwnerlessArtifactWrapper::new, Function.identity()));
     // This is only called for shared actions, so we'll almost certainly have to transform all keys
     // in all sets.
-    return new ActionExecutionValue(
+    return create(
         transformKeys(artifactData, newArtifactMap),
         transformKeys(treeArtifactData, newArtifactMap),
         transformKeys(additionalOutputData, newArtifactMap),
-        outputSymlinks);
+        outputSymlinks,
+        this instanceof CrossServerUnshareableActionExecutionValue);
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
index fbf3ec0..e35be19 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ArtifactFunction.java
@@ -13,7 +13,6 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
@@ -318,7 +317,6 @@
     return Label.print(ArtifactSkyKey.artifact(skyKey).getOwner());
   }
 
-  @VisibleForTesting
   static ActionLookupKey getActionLookupKey(Artifact artifact) {
     ArtifactOwner artifactOwner = artifact.getArtifactOwner();
 
@@ -327,7 +325,7 @@
   }
 
   @Nullable
-  private static ActionLookupValue getActionLookupValue(
+  static ActionLookupValue getActionLookupValue(
       SkyKey actionLookupKey, SkyFunction.Environment env, Artifact artifact)
       throws InterruptedException {
     ActionLookupValue value = (ActionLookupValue) env.getValue(actionLookupKey);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
index 0b41400..aed4d5d 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyFunctions.java
@@ -15,6 +15,7 @@
 
 import com.google.common.base.Predicate;
 import com.google.devtools.build.lib.actions.ActionLookupData;
+import com.google.devtools.build.skyframe.ShareabilityOfValue;
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 
@@ -81,7 +82,8 @@
       SkyFunctionName.create("TARGET_COMPLETION");
   public static final SkyFunctionName ASPECT_COMPLETION =
       SkyFunctionName.create("ASPECT_COMPLETION");
-  public static final SkyFunctionName TEST_COMPLETION = SkyFunctionName.create("TEST_COMPLETION");
+  public static final SkyFunctionName TEST_COMPLETION =
+      SkyFunctionName.create("TEST_COMPLETION", ShareabilityOfValue.NEVER);
   public static final SkyFunctionName BUILD_CONFIGURATION =
       SkyFunctionName.create("BUILD_CONFIGURATION");
   public static final SkyFunctionName CONFIGURATION_FRAGMENT =
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
index 61c31d8..c515163 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -765,11 +765,12 @@
                 + " in an action which is not a SkyframeAwareAction. Action: %s\n symlinks:%s",
             action,
             actionExecutionContext.getOutputSymlinks());
-        return new ActionExecutionValue(
+        return ActionExecutionValue.create(
             metadataHandler.getOutputArtifactData(),
             metadataHandler.getOutputTreeArtifactData(),
             metadataHandler.getAdditionalOutputData(),
-            actionExecutionContext.getOutputSymlinks());
+            actionExecutionContext.getOutputSymlinks(),
+            action instanceof NotifyOnActionCacheHit);
       }
     }
   }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java
index cbbe9f9..9a7d7a6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/TestCompletionFunction.java
@@ -13,6 +13,11 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe;
 
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
+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.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.TopLevelArtifactContext;
@@ -21,6 +26,7 @@
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.util.Map;
 
 /**
  * TestCompletionFunction builds all relevant test artifacts of a {@link
@@ -45,14 +51,40 @@
 
     ConfiguredTarget ct = ctValue.getConfiguredTarget();
     if (key.exclusiveTesting()) {
-      // Request test artifacts iteratively if testing exclusively.
+      // Request test execution iteratively if testing exclusively.
       for (Artifact testArtifact : TestProvider.getTestStatusArtifacts(ct)) {
-        if (env.getValue(ArtifactSkyKey.key(testArtifact, /*isMandatory=*/ true)) == null) {
+        ActionLookupValue.ActionLookupKey actionLookupKey =
+            ArtifactFunction.getActionLookupKey(testArtifact);
+        ActionLookupValue actionLookupValue =
+            ArtifactFunction.getActionLookupValue(actionLookupKey, env, testArtifact);
+        if (actionLookupValue == null) {
+          return null;
+        }
+        env.getValue(getActionLookupData(testArtifact, actionLookupKey, actionLookupValue));
+        if (env.valuesMissing()) {
           return null;
         }
       }
     } else {
-      env.getValues(TestProvider.getTestStatusArtifacts(ct));
+      Multimap<ActionLookupValue.ActionLookupKey, Artifact> keyToArtifactMap =
+          Multimaps.index(
+              TestProvider.getTestStatusArtifacts(ct), ArtifactFunction::getActionLookupKey);
+      Map<SkyKey, SkyValue> actionLookupValues = env.getValues(keyToArtifactMap.keySet());
+      if (env.valuesMissing()) {
+        return null;
+      }
+      env.getValues(
+          keyToArtifactMap
+              .entries()
+              .stream()
+              .map(
+                  entry ->
+                      getActionLookupData(
+                          entry.getValue(),
+                          entry.getKey(),
+                          (ActionLookupValue) actionLookupValues.get(entry.getKey())))
+              .distinct()
+              .collect(ImmutableSet.toImmutableSet()));
       if (env.valuesMissing()) {
         return null;
       }
@@ -60,6 +92,14 @@
     return TestCompletionValue.TEST_COMPLETION_MARKER;
   }
 
+  private static ActionLookupData getActionLookupData(
+      Artifact artifact,
+      ActionLookupValue.ActionLookupKey actionLookupKey,
+      ActionLookupValue actionLookupValue) {
+    return ActionExecutionValue.key(
+        actionLookupKey, actionLookupValue.getGeneratingActionIndex(artifact));
+  }
+
   @Override
   public String extractTag(SkyKey skyKey) {
     return Label.print(((ConfiguredTargetKey) skyKey.argument()).getLabel());
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
index 96f9a87..1a01b9b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
@@ -19,6 +19,7 @@
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:registered-singleton",
         "//src/main/java/com/google/devtools/build/lib/unsafe:string",
         "//src/main/java/com/google/devtools/build/lib/unsafe:unsafe-provider",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//third_party:guava",
         "//third_party:jsr305",
         "//third_party/protobuf:protobuf_java",
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnshareableValue.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnshareableValue.java
new file mode 100644
index 0000000..24ebd0c
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/UnshareableValue.java
@@ -0,0 +1,20 @@
+// Copyright 2018 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.skyframe.serialization;
+
+import com.google.devtools.build.skyframe.SkyValue;
+
+/** Marker interface to indicate that this object contains unshareable data. */
+public interface UnshareableValue extends SkyValue {}
diff --git a/src/main/java/com/google/devtools/build/skyframe/BUILD b/src/main/java/com/google/devtools/build/skyframe/BUILD
index 34ad3a7..8a57fd0 100644
--- a/src/main/java/com/google/devtools/build/skyframe/BUILD
+++ b/src/main/java/com/google/devtools/build/skyframe/BUILD
@@ -5,10 +5,11 @@
 )
 
 SKYFRAME_OBJECT_SRCS = [
-    "SkyValue.java",
-    "SkyKey.java",
     "AbstractSkyKey.java",
+    "ShareabilityOfValue.java",
     "SkyFunctionName.java",
+    "SkyKey.java",
+    "SkyValue.java",
 ]
 
 java_library(
@@ -16,7 +17,6 @@
     srcs = SKYFRAME_OBJECT_SRCS,
     visibility = ["//visibility:public"],
     deps = [
-        "//src/main/java/com/google/devtools/build/lib/collect",
         "//third_party:guava",
     ],
 )
diff --git a/src/main/java/com/google/devtools/build/skyframe/ShareabilityOfValue.java b/src/main/java/com/google/devtools/build/skyframe/ShareabilityOfValue.java
new file mode 100644
index 0000000..d9f9899
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/skyframe/ShareabilityOfValue.java
@@ -0,0 +1,32 @@
+// Copyright 2018 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.skyframe;
+
+/**
+ * When the {@link NodeEntry#getValue} corresponding to a given {@link SkyFunctionName} is
+ * shareable: always, sometimes (depending on the specific key argument and/or value), or never.
+ *
+ * <p>Values may be unshareable because they are just not serializable, or because they contain data
+ * that cannot safely be re-used as-is by another invocation.
+ *
+ * <p>Unshareable data should not be serialized, since it will never be re-used. Attempts to fetch
+ * serialized data will check this value and only perform the fetch if the value is not {@link
+ * #NEVER}.
+ */
+public enum ShareabilityOfValue {
+  ALWAYS,
+  SOMETIMES,
+  NEVER
+}
diff --git a/src/main/java/com/google/devtools/build/skyframe/SkyFunctionName.java b/src/main/java/com/google/devtools/build/skyframe/SkyFunctionName.java
index af7474c..3d6ece1 100644
--- a/src/main/java/com/google/devtools/build/skyframe/SkyFunctionName.java
+++ b/src/main/java/com/google/devtools/build/skyframe/SkyFunctionName.java
@@ -13,30 +13,18 @@
 // limitations under the License.
 package com.google.devtools.build.skyframe;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
+import com.google.common.cache.Cache;
 import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-
 import java.io.Serializable;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
 
 /** An identifier for a {@code SkyFunction}. */
 public final class SkyFunctionName implements Serializable {
-
-  // In practice the number of unique SkyFunctionNames should be reasonably limited, use this cache
-  // to avoid accidentally creating many of the same.
-  private static final LoadingCache<String, SkyFunctionName> skyFunctionNameCache =
-      CacheBuilder.newBuilder()
-        .weakValues()
-        .build(
-            new CacheLoader<String, SkyFunctionName>() {
-              @Override
-              public SkyFunctionName load(String name) {
-                return new SkyFunctionName(name);
-              }
-            });
+  private static final Cache<NameOnlyWrapper, SkyFunctionName> interner =
+      CacheBuilder.newBuilder().build();
 
   /**
    * A well-known key type intended for testing only. The associated SkyKey should have a String
@@ -47,26 +35,45 @@
 
   /** Create a SkyFunctionName identified by {@code name}. */
   public static SkyFunctionName create(String name) {
+    return create(name, ShareabilityOfValue.ALWAYS);
+  }
+
+  public static SkyFunctionName create(String name, ShareabilityOfValue shareabilityOfValue) {
+    SkyFunctionName result = new SkyFunctionName(name, shareabilityOfValue);
+    SkyFunctionName cached;
     try {
-      return skyFunctionNameCache.get(name);
+      cached = interner.get(new NameOnlyWrapper(result), () -> result);
     } catch (ExecutionException e) {
-      throw new IllegalStateException("Unexpected exception creating SkyFunctionName", e);
+      throw new IllegalStateException(e);
     }
+    Preconditions.checkState(
+        result.equals(cached),
+        "Tried to create SkyFunctionName objects with same name but different properties: %s %s",
+        result,
+        cached);
+    return cached;
   }
 
   private final String name;
+  private final ShareabilityOfValue shareabilityOfValue;
 
-  private SkyFunctionName(String name) {
+  private SkyFunctionName(String name, ShareabilityOfValue shareabilityOfValue) {
     this.name = name;
+    this.shareabilityOfValue = shareabilityOfValue;
   }
 
   public String getName() {
     return name;
   }
 
+  public ShareabilityOfValue getShareabilityOfValue() {
+    return shareabilityOfValue;
+  }
+
   @Override
   public String toString() {
-    return name;
+    return name
+        + (shareabilityOfValue.equals(ShareabilityOfValue.ALWAYS) ? "" : " " + shareabilityOfValue);
   }
 
   @Override
@@ -78,11 +85,12 @@
       return false;
     }
     SkyFunctionName other = (SkyFunctionName) obj;
-    return name.equals(other.name);
+    return name.equals(other.name) && shareabilityOfValue.equals(other.shareabilityOfValue);
   }
 
   @Override
   public int hashCode() {
+    // Don't bother incorporating serializabilityOfValue into hashCode: should always be the same.
     return name.hashCode();
   }
 
@@ -90,23 +98,35 @@
    * A predicate that returns true for {@link SkyKey}s that have the given {@link SkyFunctionName}.
    */
   public static Predicate<SkyKey> functionIs(final SkyFunctionName functionName) {
-    return new Predicate<SkyKey>() {
-      @Override
-      public boolean apply(SkyKey skyKey) {
-        return functionName.equals(skyKey.functionName());
-      }
-    };
+    return skyKey -> functionName.equals(skyKey.functionName());
   }
 
   /**
    * A predicate that returns true for {@link SkyKey}s that have the given {@link SkyFunctionName}.
    */
   public static Predicate<SkyKey> functionIsIn(final Set<SkyFunctionName> functionNames) {
-    return new Predicate<SkyKey>() {
-      @Override
-      public boolean apply(SkyKey skyKey) {
-        return functionNames.contains(skyKey.functionName());
+    return skyKey -> functionNames.contains(skyKey.functionName());
+  }
+
+  private static class NameOnlyWrapper {
+    private final SkyFunctionName skyFunctionName;
+
+    private NameOnlyWrapper(SkyFunctionName skyFunctionName) {
+      this.skyFunctionName = skyFunctionName;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (!(obj instanceof NameOnlyWrapper)) {
+        return false;
       }
-    };
+      SkyFunctionName thatFunctionName = ((NameOnlyWrapper) obj).skyFunctionName;
+      return this.skyFunctionName.getName().equals(thatFunctionName.name);
+    }
+
+    @Override
+    public int hashCode() {
+      return skyFunctionName.getName().hashCode();
+    }
   }
 }