Optimize GC usage of iterating over all elements of GroupedLists when we don't care about the group structure, and simplify the logic for prefetching old deps.

PiperOrigin-RevId: 187681887

3 files changed, 29 insertions, 17 deletions

diff --git a/src/main/java/com/google/devtools/build/skyframe/InMemoryNodeEntry.java b/src/main/java/com/google/devtools/build/skyframe/InMemoryNodeEntry.java
index 76e8100f32..24dfeab6b4 100644
--- a/src/main/java/com/google/devtools/build/skyframe/InMemoryNodeEntry.java
+++ b/src/main/java/com/google/devtools/build/skyframe/InMemoryNodeEntry.java
@@ -18,7 +18,6 @@ import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Iterables;
 import com.google.devtools.build.lib.util.GroupedList;
 import com.google.devtools.build.lib.util.GroupedList.GroupedListHelper;
 import com.google.devtools.build.skyframe.KeyToConsolidate.Op;
@@ -220,7 +219,7 @@ public class InMemoryNodeEntry implements NodeEntry {
 
   @Override
   public synchronized Iterable<SkyKey> getDirectDeps() {
-    return getGroupedDirectDeps().toSet();
+    return getGroupedDirectDeps().getAllElementsAsIterable();
   }
 
   /**
@@ -520,7 +519,7 @@ public class InMemoryNodeEntry implements NodeEntry {
       throws InterruptedException {
     Preconditions.checkState(!isDone(), this);
     if (!isDirty()) {
-      return Iterables.concat(getTemporaryDirectDeps());
+      return getTemporaryDirectDeps().getAllElementsAsIterable();
     } else {
       // There may be duplicates here. Make sure everything is unique.
       ImmutableSet.Builder<SkyKey> result = ImmutableSet.builder();
diff --git a/src/main/java/com/google/devtools/build/skyframe/QueryableGraph.java b/src/main/java/com/google/devtools/build/skyframe/QueryableGraph.java
index 26422dd418..8c5af09ded 100644
--- a/src/main/java/com/google/devtools/build/skyframe/QueryableGraph.java
+++ b/src/main/java/com/google/devtools/build/skyframe/QueryableGraph.java
@@ -50,6 +50,15 @@ public interface QueryableGraph {
           throws InterruptedException;
 
   /**
+   * A prefetch call may be used to hint to the graph that we may call {@link #getBatch} on the
+   * specified keys later.
+   */
+  default void prefetchBatch(
+      @Nullable SkyKey requestor, Reason reason, Iterable<? extends SkyKey> keys) {
+    // Do nothing.
+  }
+
+  /**
    * Examines all the given keys. Returns an iterable of keys whose corresponding nodes are
    * currently available to be fetched.
    *
diff --git a/src/main/java/com/google/devtools/build/skyframe/SkyFunctionEnvironment.java b/src/main/java/com/google/devtools/build/skyframe/SkyFunctionEnvironment.java
index 9646b07567..d8c8f7ac28 100644
--- a/src/main/java/com/google/devtools/build/skyframe/SkyFunctionEnvironment.java
+++ b/src/main/java/com/google/devtools/build/skyframe/SkyFunctionEnvironment.java
@@ -20,8 +20,6 @@ import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicates;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -151,20 +149,26 @@ class SkyFunctionEnvironment extends AbstractSkyFunctionEnvironment {
       boolean assertDone,
       SkyKey keyForDebugging)
       throws InterruptedException {
-    Iterable<SkyKey> depKeysAsIterable = Iterables.concat(depKeys);
-    Iterable<SkyKey> keysToPrefetch = depKeysAsIterable;
+    Set<SkyKey> depKeysAsSet = null;
     if (PREFETCH_OLD_DEPS) {
-      ImmutableSet.Builder<SkyKey> keysToPrefetchBuilder = ImmutableSet.builder();
-      keysToPrefetchBuilder.addAll(depKeysAsIterable).addAll(oldDeps);
-      keysToPrefetch = keysToPrefetchBuilder.build();
+      if (!oldDeps.isEmpty()) {
+        // Create a set here so that filtering the old deps below is fast. Once we create this set,
+        // we may as well use it for the call to evaluatorContext#getBatchValues since we've
+        // precomputed the size.
+        depKeysAsSet = depKeys.toSet();
+        evaluatorContext
+            .getGraph()
+            .prefetchBatch(
+                requestor,
+                Reason.PREFETCH,
+                Iterables.filter(oldDeps, Predicates.not(Predicates.in(depKeysAsSet))));
+      }
     }
     Map<SkyKey, ? extends NodeEntry> batchMap =
-        evaluatorContext.getBatchValues(requestor, Reason.PREFETCH, keysToPrefetch);
-    if (PREFETCH_OLD_DEPS) {
-      batchMap =
-          ImmutableMap.<SkyKey, NodeEntry>copyOf(
-              Maps.filterKeys(batchMap, Predicates.in(ImmutableSet.copyOf(depKeysAsIterable))));
-    }
+        evaluatorContext.getBatchValues(
+            requestor,
+            Reason.PREFETCH,
+            depKeysAsSet == null ? depKeys.getAllElementsAsIterable() : depKeysAsSet);
     if (batchMap.size() != depKeys.numElements()) {
       throw new IllegalStateException(
           "Missing keys for "
@@ -316,7 +320,7 @@ class SkyFunctionEnvironment extends AbstractSkyFunctionEnvironment {
             keySize - directDeps.size() - (bubbleErrorInfo == null ? 0 : bubbleErrorInfo.size()),
             0);
     ArrayList<SkyKey> missingKeys = new ArrayList<>(expectedMissingKeySize);
-    for (SkyKey key : Iterables.concat(depKeys)) {
+    for (SkyKey key : depKeys.getAllElementsAsIterable()) {
       SkyValue value = maybeGetValueFromErrorOrDeps(key);
       if (value == null) {
         missingKeys.add(key);