[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);
+ }
+ }
+}