bzlmod: use env.getValues() instead of getValue() in DiscoveryFunction

(#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
bazelbuild · Jul 2, 2021 · 1464393 · 1464393
1 parent 07a84ce
commit 1464393
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 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
@@ -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
@@ -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));