Use a SkyFunction

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

@@ -47,6 +47,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;
@@ -278,7 +279,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());
     filesystem = runtime.getFileSystem();
 
     credentialModule = Preconditions.checkNotNull(runtime.getBlazeModule(CredentialModule.class));

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

@@ -25,6 +25,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",
@@ -173,6 +175,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;
@@ -28,7 +31,6 @@
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.CheckDirectDepsMode;
 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.profiler.Profiler;
 import com.google.devtools.build.lib.profiler.ProfilerTask;
 import com.google.devtools.build.lib.profiler.SilentCloseable;
@@ -39,18 +41,11 @@
 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 java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionException;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
 import javax.annotation.Nullable;
 
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
-
 /**
  * Discovers the whole dependency graph and runs selection algorithm on it to produce the pruned
  * dependency graph and runs checks on it.
@@ -71,12 +66,56 @@ public SkyValue compute(SkyKey skyKey, Environment env)
     if (root == null) {
       return null;
     }
+
+    var state = env.getState(ModuleResolutionComputeState::new);
+    if (state.selectionResult == null) {
+      state.selectionResult = discoverAndSelect(env, root);
+      if (state.selectionResult == null) {
+        return null;
+      }
+    }
+
+    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());
+    }
+
+    Profiler.instance().profile(ProfilerTask.BZLMOD, "module resolution completed").close();
+
+    return BazelModuleResolutionValue.create(
+        finalDepGraph, state.selectionResult.getUnprunedDepGraph());
+  }
+
+  private static Selection.Result discoverAndSelect(Environment env, RootModuleFileValue root)
+      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;
@@ -109,16 +148,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       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(
@@ -237,44 +267,12 @@ 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).
    */
   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, ModuleOverride override, RepoSpec remoteRepoSpec) {
     return Module.builder()
         .setName(interim.getName())
         .setVersion(interim.getVersion())
@@ -283,51 +281,29 @@ 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 {
-    ExecutorService executorService = Executors.newWorkStealingPool(8);
-    var future =
-        executorService.submit(
-            () ->
-                resolvedDepGraph.entrySet().stream()
-                    .parallel()
-                    .map(
-                        e -> {
-                          try {
-                            return Maps.immutableEntry(
-                                e.getKey(),
-                                moduleFromInterimModule(
-                                    e.getValue(),
-                                    overrides.get(e.getKey().getName()),
-                                    eventHandler));
-                          } catch (BazelModuleResolutionFunctionException ex) {
-                            throw new CompletionException(ex);
-                          } catch (InterruptedException ex) {
-                            throw new CompletionException(ex);
-                          }
-                        })
-                    .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue)));
-    try {
-      var result = future.get();
-      executorService.shutdown();
-      return result;
-    } catch (ExecutionException e) {
-      if (e.getCause() instanceof BazelModuleResolutionFunctionException) {
-        throw (BazelModuleResolutionFunctionException) e.getCause();
-      } else if (e.getCause() instanceof InterruptedException) {
-        throw (InterruptedException) e.getCause();
-      } else {
-        throw new IllegalStateException(e);
-      }
+      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()),
+              remoteRepoSpecs.get(entry.getKey())));
     }
+    return finalDepGraph.buildOrThrow();
+  }
+
+  private static class ModuleResolutionComputeState implements Environment.SkyKeyComputeState {
+    Selection.Result selectionResult;
   }
 
   static class BazelModuleResolutionFunctionException extends SkyFunctionException {

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

@@ -17,6 +17,7 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import com.google.common.base.Objects;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -65,6 +66,23 @@ public IndexRegistry(URI uri, DownloadManager downloadManager, Map<String, Strin
             .create();
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    IndexRegistry that = (IndexRegistry) o;
+    return Objects.equal(uri, that.uri);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(uri);
+  }
+
   @Override
   public String getUrl() {
     return uri.toString();

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,61 @@
+// 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 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 {
+  @Override
+  @Nullable
+  public SkyValue compute(SkyKey skyKey, Environment env)
+      throws InterruptedException, RepoSpecException {
+    RepoSpecKey module = (RepoSpecKey) skyKey.argument();
+    try (SilentCloseable c =
+        Profiler.instance()
+            .profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + module.getModuleKey())) {
+      return module
+          .getRegistry()
+          .getRepoSpec(module.getModuleKey(), module.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",
+              module.getModuleKey()));
+    }
+  }
+
+  static final class RepoSpecException extends SkyFunctionException {
+
+    RepoSpecException(ExternalDepsException cause) {
+      super(cause, Transience.TRANSIENT);
+    }
+  }
+}