[6.4.0] Fetch RepoSpecs in parallel

src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java

@@ -46,6 +46,7 @@
 import com.google.devtools.build.lib.bazel.bzlmod.RegistryFactory;
 import com.google.devtools.build.lib.bazel.bzlmod.RegistryFactoryImpl;
 import com.google.devtools.build.lib.bazel.bzlmod.RepoSpec;
+import com.google.devtools.build.lib.bazel.bzlmod.RepoSpecFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionEvalFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionUsagesFunction;
 import com.google.devtools.build.lib.bazel.bzlmod.YankedVersionsUtil;
@@ -272,7 +273,8 @@ SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace
         .addSkyFunction(SkyFunctions.BAZEL_MODULE_INSPECTION, new BazelModuleInspectorFunction())
         .addSkyFunction(SkyFunctions.BAZEL_MODULE_RESOLUTION, new BazelModuleResolutionFunction())
         .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_EVAL, singleExtensionEvalFunction)
-        .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction());
+        .addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction())
+        .addSkyFunction(SkyFunctions.REPO_SPEC, new RepoSpecFunction(registryFactory));
     filesystem = runtime.getFileSystem();
 
     credentialModule = Preconditions.checkNotNull(runtime.getBlazeModule(CredentialModule.class));

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -24,6 +24,7 @@ java_library(
     ],
     deps = [
         "//src/main/java/com/google/devtools/build/lib/cmdline",
+        "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
         "//src/main/java/net/starlark/java/eval",
         "//third_party:auto_value",
         "//third_party:gson",
@@ -131,6 +132,7 @@ java_library(
         "MultipleVersionOverride.java",
         "NonRegistryOverride.java",
         "RegistryOverride.java",
+        "RepoSpecKey.java",
         "SingleExtensionEvalValue.java",
         "SingleExtensionUsagesValue.java",
         "SingleVersionOverride.java",
@@ -174,6 +176,7 @@ java_library(
         "ModuleExtensionContext.java",
         "ModuleFileFunction.java",
         "ModuleFileGlobals.java",
+        "RepoSpecFunction.java",
         "Selection.java",
         "SingleExtensionEvalFunction.java",
         "SingleExtensionUsagesFunction.java",

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java

@@ -15,10 +15,13 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.analysis.BlazeVersionInfo;
 import com.google.devtools.build.lib.bazel.BazelVersion;
@@ -29,6 +32,7 @@
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
+import com.google.devtools.build.lib.events.StoredEventHandler;
 import com.google.devtools.build.lib.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
@@ -39,7 +43,7 @@
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
-import java.io.IOException;
+import com.google.devtools.build.skyframe.SkyframeLookupResult;
 import java.util.Map;
 import java.util.Objects;
 import javax.annotation.Nullable;
@@ -64,12 +68,61 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (root == null) {
       return null;
     }
+
+    var state = env.getState(ModuleResolutionComputeState::new);
+    try {
+      if (state.selectionResult == null) {
+        state.storedEventHandler = new StoredEventHandler();
+        state.selectionResult = discoverAndSelect(env, root, state.storedEventHandler);
+        if (state.selectionResult == null) {
+          return null;
+        }
+      }
+    } finally {
+      state.storedEventHandler.replayOn(env.getListener());
+    }
+
+    ImmutableSet<RepoSpecKey> repoSpecKeys =
+        state.selectionResult.getResolvedDepGraph().values().stream()
+            // Modules with a null registry have a non-registry override. We don't need to
+            // fetch or store the repo spec in this case.
+            .filter(module -> module.getRegistry() != null)
+            .map(RepoSpecKey::of)
+            .collect(toImmutableSet());
+    SkyframeLookupResult repoSpecResults = env.getValuesAndExceptions(repoSpecKeys);
+    ImmutableMap.Builder<ModuleKey, RepoSpec> remoteRepoSpecs = ImmutableMap.builder();
+    for (RepoSpecKey repoSpecKey : repoSpecKeys) {
+      RepoSpec repoSpec = (RepoSpec) repoSpecResults.get(repoSpecKey);
+      if (repoSpec == null) {
+        return null;
+      }
+      remoteRepoSpecs.put(repoSpecKey.getModuleKey(), repoSpec);
+    }
+
+    ImmutableMap<ModuleKey, Module> finalDepGraph;
+    try (SilentCloseable c =
+        Profiler.instance().profile(ProfilerTask.BZLMOD, "compute final dep graph")) {
+      finalDepGraph =
+          computeFinalDepGraph(
+              state.selectionResult.getResolvedDepGraph(),
+              root.getOverrides(),
+              remoteRepoSpecs.buildOrThrow());
+    }
+
+    return BazelModuleResolutionValue.create(
+        finalDepGraph, state.selectionResult.getUnprunedDepGraph());
+  }
+
+  @Nullable
+  private static Selection.Result discoverAndSelect(
+      Environment env, RootModuleFileValue root, ExtendedEventHandler eventHandler)
+      throws BazelModuleResolutionFunctionException, InterruptedException {
     ImmutableMap<ModuleKey, InterimModule> initialDepGraph;
     try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "discovery")) {
       initialDepGraph = Discovery.run(env, root);
-      if (initialDepGraph == null) {
-        return null;
-      }
+    }
+    if (initialDepGraph == null) {
+      return null;
     }
 
     Selection.Result selectionResult;
@@ -86,32 +139,23 @@ public SkyValue compute(SkyKey skyKey, Environment env)
           initialDepGraph.get(ModuleKey.ROOT),
           resolvedDepGraph.get(ModuleKey.ROOT),
           Objects.requireNonNull(CHECK_DIRECT_DEPENDENCIES.get(env)),
-          env.getListener());
+          eventHandler);
     }
 
     try (SilentCloseable c =
         Profiler.instance().profile(ProfilerTask.BZLMOD, "check bazel compatibility")) {
       checkBazelCompatibility(
           resolvedDepGraph.values(),
           Objects.requireNonNull(BAZEL_COMPATIBILITY_MODE.get(env)),
-          env.getListener());
+          eventHandler);
     }
 
     try (SilentCloseable c =
         Profiler.instance().profile(ProfilerTask.BZLMOD, "check no yanked versions")) {
       checkNoYankedVersions(resolvedDepGraph);
     }
 
-    ImmutableMap<ModuleKey, Module> finalDepGraph;
-    try (SilentCloseable c =
-        Profiler.instance().profile(ProfilerTask.BZLMOD, "compute final dep graph")) {
-      finalDepGraph =
-          computeFinalDepGraph(resolvedDepGraph, root.getOverrides(), env.getListener());
-    }
-
-    Profiler.instance().profile(ProfilerTask.BZLMOD, "module resolution completed").close();
-
-    return BazelModuleResolutionValue.create(finalDepGraph, selectionResult.getUnprunedDepGraph());
+    return selectionResult;
   }
 
   private static void verifyRootModuleDirectDepsAreAccurate(
@@ -210,7 +254,8 @@ private static void checkNoYankedVersions(ImmutableMap<ModuleKey, InterimModule>
     }
   }
 
-  private static RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOverride override) {
+  private static RepoSpec maybeAppendAdditionalPatches(
+      @Nullable RepoSpec repoSpec, @Nullable ModuleOverride override) {
     if (!(override instanceof SingleVersionOverride)) {
       return repoSpec;
     }
@@ -230,44 +275,15 @@ private static RepoSpec maybeAppendAdditionalPatches(RepoSpec repoSpec, ModuleOv
         .build();
   }
 
-  @Nullable
-  private static RepoSpec computeRepoSpec(
-      InterimModule interimModule, ModuleOverride override, ExtendedEventHandler eventHandler)
-      throws BazelModuleResolutionFunctionException, InterruptedException {
-    if (interimModule.getRegistry() == null) {
-      // This module has a non-registry override. We don't need to store the repo spec in this case.
-      return null;
-    }
-    try {
-      RepoSpec moduleRepoSpec =
-          interimModule
-              .getRegistry()
-              .getRepoSpec(
-                  interimModule.getKey(), interimModule.getCanonicalRepoName(), eventHandler);
-      return maybeAppendAdditionalPatches(moduleRepoSpec, override);
-    } catch (IOException e) {
-      throw new BazelModuleResolutionFunctionException(
-          ExternalDepsException.withMessage(
-              Code.ERROR_ACCESSING_REGISTRY,
-              "Unable to get module repo spec from registry: %s",
-              e.getMessage()),
-          Transience.PERSISTENT);
-    }
-  }
-
   /**
    * Builds a {@link Module} from an {@link InterimModule}, discarding unnecessary fields and adding
    * extra necessary ones (such as the repo spec).
+   *
+   * @param remoteRepoSpec the {@link RepoSpec} for the module obtained from a registry or null if
+   *     the module has a non-registry override
    */
   static Module moduleFromInterimModule(
-      InterimModule interim, ModuleOverride override, ExtendedEventHandler eventHandler)
-      throws BazelModuleResolutionFunctionException, InterruptedException {
-    RepoSpec repoSpec;
-    try (SilentCloseable c =
-        Profiler.instance()
-            .profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + interim.getKey())) {
-      repoSpec = computeRepoSpec(interim, override, eventHandler);
-    }
+      InterimModule interim, @Nullable ModuleOverride override, @Nullable RepoSpec remoteRepoSpec) {
     return Module.builder()
         .setName(interim.getName())
         .setVersion(interim.getVersion())
@@ -276,26 +292,32 @@ static Module moduleFromInterimModule(
         .setExecutionPlatformsToRegister(interim.getExecutionPlatformsToRegister())
         .setToolchainsToRegister(interim.getToolchainsToRegister())
         .setDeps(ImmutableMap.copyOf(Maps.transformValues(interim.getDeps(), DepSpec::toModuleKey)))
-        .setRepoSpec(repoSpec)
+        .setRepoSpec(maybeAppendAdditionalPatches(remoteRepoSpec, override))
         .setExtensionUsages(interim.getExtensionUsages())
         .build();
   }
 
   private static ImmutableMap<ModuleKey, Module> computeFinalDepGraph(
       ImmutableMap<ModuleKey, InterimModule> resolvedDepGraph,
       ImmutableMap<String, ModuleOverride> overrides,
-      ExtendedEventHandler eventHandler)
-      throws BazelModuleResolutionFunctionException, InterruptedException {
+      ImmutableMap<ModuleKey, RepoSpec> remoteRepoSpecs) {
     ImmutableMap.Builder<ModuleKey, Module> finalDepGraph = ImmutableMap.builder();
     for (Map.Entry<ModuleKey, InterimModule> entry : resolvedDepGraph.entrySet()) {
       finalDepGraph.put(
           entry.getKey(),
           moduleFromInterimModule(
-              entry.getValue(), overrides.get(entry.getKey().getName()), eventHandler));
+              entry.getValue(),
+              overrides.get(entry.getKey().getName()),
+              remoteRepoSpecs.get(entry.getKey())));
     }
     return finalDepGraph.buildOrThrow();
   }
 
+  private static class ModuleResolutionComputeState implements Environment.SkyKeyComputeState {
+    Selection.Result selectionResult;
+    StoredEventHandler storedEventHandler;
+  }
+
   static class BazelModuleResolutionFunctionException extends SkyFunctionException {
     BazelModuleResolutionFunctionException(ExternalDepsException e, Transience transience) {
       super(e, transience);

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpec.java

@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
+import com.google.devtools.build.skyframe.SkyValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import javax.annotation.Nullable;
 
@@ -24,7 +25,7 @@
  */
 @AutoValue
 @GenerateTypeAdapter
-public abstract class RepoSpec {
+public abstract class RepoSpec implements SkyValue {
 
   /**
    * The label string for the bzl file this repository rule is defined in, empty for native rule.

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecFunction.java

@@ -0,0 +1,73 @@
+// Copyright 2021 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.bazel.bzlmod;
+
+import com.google.devtools.build.lib.profiler.Profiler;
+import com.google.devtools.build.lib.profiler.ProfilerTask;
+import com.google.devtools.build.lib.profiler.SilentCloseable;
+import com.google.devtools.build.lib.server.FailureDetails;
+import com.google.devtools.build.skyframe.SkyFunction;
+import com.google.devtools.build.skyframe.SkyFunctionException;
+import com.google.devtools.build.skyframe.SkyKey;
+import com.google.devtools.build.skyframe.SkyValue;
+import java.io.IOException;
+import java.net.URISyntaxException;
+import javax.annotation.Nullable;
+
+/**
+ * A simple SkyFunction that computes a {@link RepoSpec} for the given {@link InterimModule} by
+ * fetching required information from its {@link Registry}.
+ */
+public class RepoSpecFunction implements SkyFunction {
+  private final RegistryFactory registryFactory;
+
+  public RepoSpecFunction(RegistryFactory registryFactory) {
+    this.registryFactory = registryFactory;
+  }
+
+  @Override
+  @Nullable
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws InterruptedException, RepoSpecException {
+    RepoSpecKey key = (RepoSpecKey) skyKey.argument();
+    try (SilentCloseable c =
+        Profiler.instance()
+            .profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + key.getModuleKey())) {
+      return registryFactory
+          .getRegistryWithUrl(key.getRegistryUrl())
+          .getRepoSpec(
+              key.getModuleKey(), key.getModuleKey().getCanonicalRepoName(), env.getListener());
+    } catch (IOException e) {
+      throw new RepoSpecException(
+          ExternalDepsException.withCauseAndMessage(
+              FailureDetails.ExternalDeps.Code.ERROR_ACCESSING_REGISTRY,
+              e,
+              "Unable to get module repo spec for %s from registry",
+              key.getModuleKey()));
+    } catch (URISyntaxException e) {
+      // This should never happen since we obtain the registry URL from an already constructed
+      // registry.
+      throw new IllegalStateException(e);
+    }
+  }
+
+  static final class RepoSpecException extends SkyFunctionException {
+
+    RepoSpecException(ExternalDepsException cause) {
+      super(cause, Transience.TRANSIENT);
+    }
+  }
+}