Skip to content

Commit

Permalink
Fetch RepoSpecs in parallel
Browse files Browse the repository at this point in the history
Computing `RepoSpec`s for all selected Bazel modules is on the critical path for the computation of the main repository mapping and thus benefits from parallelized downloads.

On my machine, this change has the following effect on `bazel build //src:bazel-dev --enable_bzlmod --nobuild`:

```
before
compute main repo mapping: 8s 127ms

after
compute main repo mapping: 4s 226ms
```

Closes #19294.

PiperOrigin-RevId: 559819452
Change-Id: Ieef957fcfe402c909d2863ba4a4ca3540781a56d
  • Loading branch information
fmeum authored and Copybara-Service committed Aug 24, 2023
1 parent 32fa5c2 commit 8a68310
Show file tree
Hide file tree
Showing 16 changed files with 215 additions and 55 deletions.
Expand Up @@ -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;
Expand Down Expand Up @@ -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(registryFactory));
filesystem = runtime.getFileSystem();

credentialModule = Preconditions.checkNotNull(runtime.getBlazeModule(CredentialModule.class));
Expand Down
Expand Up @@ -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",
Expand Down Expand Up @@ -131,6 +132,7 @@ java_library(
"MultipleVersionOverride.java",
"NonRegistryOverride.java",
"RegistryOverride.java",
"RepoSpecKey.java",
"SingleExtensionEvalValue.java",
"SingleExtensionUsagesValue.java",
"SingleVersionOverride.java",
Expand Down Expand Up @@ -173,6 +175,7 @@ java_library(
"ModuleExtensionContext.java",
"ModuleFileFunction.java",
"ModuleFileGlobals.java",
"RepoSpecFunction.java",
"Selection.java",
"SingleExtensionEvalFunction.java",
"SingleExtensionUsagesFunction.java",
Expand Down
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -39,7 +41,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;
Expand All @@ -64,12 +66,55 @@ 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());
}

return BazelModuleResolutionValue.create(
finalDepGraph, state.selectionResult.getUnprunedDepGraph());
}

@Nullable
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;
Expand Down Expand Up @@ -102,16 +147,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(
Expand Down Expand Up @@ -210,7 +246,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;
}
Expand All @@ -230,44 +267,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())
Expand All @@ -276,26 +284,31 @@ 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;
}

static class BazelModuleResolutionFunctionException extends SkyFunctionException {
BazelModuleResolutionFunctionException(ExternalDepsException e, Transience transience) {
super(e, transience);
Expand Down
Expand Up @@ -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;

Expand All @@ -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.
Expand Down
@@ -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);
}
}
}
@@ -0,0 +1,55 @@
// 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.auto.value.AutoValue;
import com.google.common.base.Preconditions;
import com.google.devtools.build.lib.skyframe.SkyFunctions;
import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
import com.google.devtools.build.skyframe.SkyFunctionName;
import com.google.devtools.build.skyframe.SkyKey;

/** The key for {@link RepoSpecFunction}. */
@AutoCodec
@AutoValue
abstract class RepoSpecKey implements SkyKey {
private static final SkyKeyInterner<RepoSpecKey> interner = SkyKey.newInterner();

static RepoSpecKey of(InterimModule module) {
Preconditions.checkNotNull(
module.getRegistry(), "module must not have a non-registry override");
return create(module.getKey(), module.getRegistry().getUrl());
}

abstract ModuleKey getModuleKey();

abstract String getRegistryUrl();

@AutoCodec.Instantiator
static RepoSpecKey create(ModuleKey moduleKey, String registryUrl) {
return interner.intern(new AutoValue_RepoSpecKey(moduleKey, registryUrl));
}

@Override
public SkyFunctionName functionName() {
return SkyFunctions.REPO_SPEC;
}

@Override
public SkyKeyInterner<RepoSpecKey> getSkyKeyInterner() {
return interner;
}
}
Expand Up @@ -158,6 +158,7 @@ public final class SkyFunctions {
SkyFunctionName.createHermetic("BAZEL_DEP_GRAPH");
public static final SkyFunctionName BAZEL_LOCK_FILE =
SkyFunctionName.createHermetic("BAZEL_LOCK_FILE");
public static final SkyFunctionName REPO_SPEC = SkyFunctionName.createNonHermetic("REPO_SPEC");

public static Predicate<SkyKey> isSkyFunction(SkyFunctionName functionName) {
return key -> key.functionName().equals(functionName);
Expand Down

0 comments on commit 8a68310

Please sign in to comment.