From 1464393e6230c39f83d6b1f753134862f40ba684 Mon Sep 17 00:00:00 2001
From: wyv <wyv@google.com>
Date: Fri, 2 Jul 2021 02:37:33 -0700
Subject: [PATCH] bzlmod: use env.getValues() instead of getValue() in
 DiscoveryFunction

(https://github.com/bazelbuild/bazel/issues/13316)

Before we were requesting the ModuleFileValues one by one. This does not result in things being slow per se (since the individual ModuleFileFunction calls should be parallelized by Skyframe) but technically using getValues() can make SkyValue dependencies cleaner (since there is no dependency between the same "level" of SkyKeys being requested here).

PiperOrigin-RevId: 382700697
---
 .../lib/bazel/bzlmod/DiscoveryFunction.java   | 30 ++++++++++++-------
 .../lib/bazel/bzlmod/ModuleFileValue.java     |  2 --
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DiscoveryFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DiscoveryFunction.java
index e855d3b6dbfd51..580dd1c0e66136 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DiscoveryFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/DiscoveryFunction.java
@@ -22,8 +22,10 @@
 import com.google.devtools.build.skyframe.SkyValue;
 import java.util.ArrayDeque;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Queue;
+import java.util.Set;
 import javax.annotation.Nullable;
 
 /**
@@ -48,22 +50,28 @@ public SkyValue compute(SkyKey skyKey, Environment env)
         rootModuleKey, rewriteDepKeys(root.getModule(), overrides, rootModuleKey.getName()));
     Queue<ModuleKey> unexpanded = new ArrayDeque<>();
     unexpanded.add(rootModuleKey);
-    // TODO(wyv): currently we expand the "unexpanded" keys one by one. We should try to expand them
-    //   all at once, using `env.getValues`.
     while (!unexpanded.isEmpty()) {
-      Module module = depGraph.get(unexpanded.remove());
-      for (ModuleKey depKey : module.getDeps().values()) {
-        if (depGraph.containsKey(depKey)) {
-          continue;
+      Set<SkyKey> unexpandedSkyKeys = new HashSet<>();
+      while (!unexpanded.isEmpty()) {
+        Module module = depGraph.get(unexpanded.remove());
+        for (ModuleKey depKey : module.getDeps().values()) {
+          if (depGraph.containsKey(depKey)) {
+            continue;
+          }
+          unexpandedSkyKeys.add(ModuleFileValue.key(depKey, overrides.get(depKey.getName())));
         }
-        ModuleFileValue dep =
-            (ModuleFileValue)
-                env.getValue(ModuleFileValue.key(depKey, overrides.get(depKey.getName())));
-        if (dep == null) {
+      }
+      Map<SkyKey, SkyValue> result = env.getValues(unexpandedSkyKeys);
+      for (Map.Entry<SkyKey, SkyValue> entry : result.entrySet()) {
+        ModuleKey depKey = ((ModuleFileValue.Key) entry.getKey()).getModuleKey();
+        ModuleFileValue moduleFileValue = (ModuleFileValue) entry.getValue();
+        if (moduleFileValue == null) {
           // Don't return yet. Try to expand any other unexpanded nodes before returning.
           depGraph.put(depKey, null);
         } else {
-          depGraph.put(depKey, rewriteDepKeys(dep.getModule(), overrides, rootModuleKey.getName()));
+          depGraph.put(
+              depKey,
+              rewriteDepKeys(moduleFileValue.getModule(), overrides, rootModuleKey.getName()));
           unexpanded.add(depKey);
         }
       }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileValue.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileValue.java
index 37a80add67948e..d39262b7372121 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileValue.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileValue.java
@@ -65,7 +65,6 @@ public static Key keyForRootModule() {
   }
 
   /** {@link SkyKey} for {@link ModuleFileValue} computation. */
-  @AutoCodec.VisibleForSerialization
   @AutoCodec
   @AutoValue
   abstract static class Key implements SkyKey {
@@ -76,7 +75,6 @@ abstract static class Key implements SkyKey {
     @Nullable
     abstract ModuleOverride getOverride();
 
-    @AutoCodec.VisibleForSerialization
     @AutoCodec.Instantiator
     static Key create(ModuleKey moduleKey, @Nullable ModuleOverride override) {
       return interner.intern(new AutoValue_ModuleFileValue_Key(moduleKey, override));