Skip to content

Commit

Permalink
WIP: Missing registry deduplication
Browse files Browse the repository at this point in the history
  • Loading branch information
@fmeum
fmeum committed Apr 25, 2024
1 parent 6f1aaf7 commit 189282d
Show file tree
Hide file tree
Showing 17 changed files with 215 additions and 610 deletions.
513 changes: 0 additions & 513 deletions MODULE.bazel.lock
View file

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ java_library(
"RegistryKey.java",
"RegistryOverride.java",
"RepoSpecKey.java",
"RepoSpecValue.java",
"SingleExtensionUsagesValue.java",
"SingleExtensionValue.java",
"SingleVersionOverride.java",
Expand Down
58 changes: 11 additions & 47 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@

package com.google.devtools.build.lib.bazel.bzlmod;

import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.collect.ImmutableMap;
Expand All @@ -27,17 +26,14 @@
import com.google.devtools.build.lib.cmdline.LabelConstants;
import com.google.devtools.build.lib.runtime.BlazeModule;
import com.google.devtools.build.lib.runtime.CommandEnvironment;
import com.google.devtools.build.lib.util.AbruptExitException;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.Root;
import com.google.devtools.build.lib.vfs.RootedPath;
import com.google.devtools.build.skyframe.MemoizingEvaluator;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
Expand All @@ -52,7 +48,6 @@ public class BazelLockFileModule extends BlazeModule {
private MemoizingEvaluator evaluator;
private Path workspaceRoot;
@Nullable private BazelModuleResolutionEvent moduleResolutionEvent;
private final List<RegistryFileDownloadEvent> downloadEvents = new ArrayList<>();

private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();

Expand All @@ -67,16 +62,7 @@ public void beforeCommand(CommandEnvironment env) {
}

@Override
public void afterCommand() throws AbruptExitException {
try {
doAfterCommand();
} finally {
moduleResolutionEvent = null;
downloadEvents.clear();
}
}

public void doAfterCommand() {
public void afterCommand() {
if (moduleResolutionEvent == null) {
// Command does not use Bazel modules or the lockfile mode is not update.
// Since Skyframe caches events, they are replayed even when nothing has changed.
Expand All @@ -86,22 +72,20 @@ public void doAfterCommand() {
BazelLockFileValue oldLockfile = moduleResolutionEvent.getOnDiskLockfileValue();

ImmutableMap<String, Optional<Checksum>> fileHashes;
if (downloadEvents.isEmpty()) {
BazelModuleResolutionValue moduleResolutionValue;
try {
moduleResolutionValue =
(BazelModuleResolutionValue) evaluator.getExistingValue(BazelModuleResolutionValue.KEY);
} catch (InterruptedException e) {
// Not thrown in Bazel.
throw new IllegalStateException(e);
}
if (moduleResolutionValue == null) {
// BazelDepGraphFunction took the dep graph from the lockfile and didn't cause evaluation of
// BazelModuleResolutionFunction. The file hashes in the lockfile are still up-to-date.
fileHashes = oldLockfile.getFileHashes();
} else {
// BazelModuleResolutionFunction ran and emitted events with the up-to-date list of accessed
// registry files.
fileHashes =
downloadEvents.stream()
.collect(
toImmutableMap(
RegistryFileDownloadEvent::getUri,
RegistryFileDownloadEvent::getChecksum,
// This should never be called since we take care to download the same file
// only once per build.
BazelLockFileModule::failOnDifferentChecksums));
fileHashes = moduleResolutionValue.getRegistryFileHashes();
}

// All nodes corresponding to module extensions that have been evaluated in the current build
Expand Down Expand Up @@ -144,20 +128,6 @@ public void doAfterCommand() {
}
}

/** A toMap merge function that throws an unchecked exception if the checksums differ. */
private static Optional<Checksum> failOnDifferentChecksums(
Optional<Checksum> checksum1, Optional<Checksum> checksum2) {
if (checksum1.isEmpty() && checksum2.isEmpty()) {
return Optional.empty();
}
if (checksum1.isPresent()
&& checksum2.isPresent()
&& checksum1.get().toString().equals(checksum2.get().toString())) {
return checksum1;
}
throw new IllegalStateException("Checksums differ: " + checksum1 + " vs " + checksum2);
}

/**
* Combines the old extensions stored in the lockfile -if they are still valid- with the new
* extensions from the events (if any)
Expand Down Expand Up @@ -301,10 +271,4 @@ public void bazelModuleResolved(BazelModuleResolutionEvent moduleResolutionEvent
// sent after the command has modified the module file.
this.moduleResolutionEvent = moduleResolutionEvent;
}

@SuppressWarnings("unused")
@Subscribe
public void registryFileDownloaded(RegistryFileDownloadEvent downloadEvent) {
downloadEvents.add(downloadEvent);
}
}
41 changes: 34 additions & 7 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import com.google.common.base.Joiner;
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.Sets;
Expand All @@ -31,6 +32,7 @@
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.BazelCompatibilityMode;
import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.CheckDirectDepsMode;
import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
import com.google.devtools.build.lib.events.Event;
import com.google.devtools.build.lib.events.EventHandler;
import com.google.devtools.build.lib.profiler.Profiler;
Expand All @@ -46,10 +48,16 @@
import com.google.devtools.build.skyframe.SkyKey;
import com.google.devtools.build.skyframe.SkyValue;
import com.google.devtools.build.skyframe.SkyframeLookupResult;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Stream;
import javax.annotation.Nullable;

/**
Expand Down Expand Up @@ -92,10 +100,14 @@ public SkyValue compute(SkyKey skyKey, Environment env)

var state = env.getState(ModuleResolutionComputeState::new);
if (state.selectionResult == null) {
state.selectionResult = discoverAndSelect(env, root, allowedYankedVersions);
ImmutableList.Builder<ImmutableMap<String, Optional<Checksum>>>
discoveryNestedRegistryFileHashes = ImmutableList.builder();
state.selectionResult =
discoverAndSelect(env, root, allowedYankedVersions, discoveryNestedRegistryFileHashes);
if (state.selectionResult == null) {
return null;
}
state.discoveryNestedRegistryFileHashes = discoveryNestedRegistryFileHashes.build();
}

ImmutableSet<RepoSpecKey> repoSpecKeys =
Expand All @@ -107,12 +119,15 @@ public SkyValue compute(SkyKey skyKey, Environment env)
.collect(toImmutableSet());
SkyframeLookupResult repoSpecResults = env.getValuesAndExceptions(repoSpecKeys);
ImmutableMap.Builder<ModuleKey, RepoSpec> remoteRepoSpecs = ImmutableMap.builder();
List<ImmutableMap<String, Optional<Checksum>>> repoSpecNestedRegistryFileHashes =
new ArrayList<>();
for (RepoSpecKey repoSpecKey : repoSpecKeys) {
RepoSpec repoSpec = (RepoSpec) repoSpecResults.get(repoSpecKey);
if (repoSpec == null) {
RepoSpecValue repoSpecValue = (RepoSpecValue) repoSpecResults.get(repoSpecKey);
if (repoSpecValue == null) {
return null;
}
remoteRepoSpecs.put(repoSpecKey.getModuleKey(), repoSpec);
remoteRepoSpecs.put(repoSpecKey.getModuleKey(), repoSpecValue.repoSpec());
repoSpecNestedRegistryFileHashes.add(repoSpecValue.registryFileHashes());
}

ImmutableMap<ModuleKey, Module> finalDepGraph;
Expand All @@ -125,19 +140,30 @@ public SkyValue compute(SkyKey skyKey, Environment env)
remoteRepoSpecs.buildOrThrow());
}

var registryFileHashes =
Stream.concat(
state.discoveryNestedRegistryFileHashes.stream(),
repoSpecNestedRegistryFileHashes.stream())
.flatMap(map -> map.entrySet().stream())
.collect(
toImmutableMap(
Map.Entry::getKey,
Map.Entry::getValue,
RegistryFileDownloadEvent::failOnDifferentChecksums));
return BazelModuleResolutionValue.create(
finalDepGraph, state.selectionResult.getUnprunedDepGraph());
finalDepGraph, state.selectionResult.getUnprunedDepGraph(), registryFileHashes);
}

@Nullable
private static Selection.Result discoverAndSelect(
Environment env,
RootModuleFileValue root,
Optional<ImmutableSet<ModuleKey>> allowedYankedVersions)
Optional<ImmutableSet<ModuleKey>> allowedYankedVersions,
ImmutableList.Builder<ImmutableMap<String, Optional<Checksum>>> nestedRegistryFileHashes)
throws BazelModuleResolutionFunctionException, InterruptedException {
ImmutableMap<ModuleKey, InterimModule> initialDepGraph;
try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "discovery")) {
initialDepGraph = Discovery.run(env, root);
initialDepGraph = Discovery.run(env, root, nestedRegistryFileHashes);
} catch (ExternalDepsException e) {
throw new BazelModuleResolutionFunctionException(e, Transience.PERSISTENT);
}
Expand Down Expand Up @@ -338,6 +364,7 @@ private static ImmutableMap<ModuleKey, Module> computeFinalDepGraph(
}

private static class ModuleResolutionComputeState implements Environment.SkyKeyComputeState {
ImmutableList<ImmutableMap<String, Optional<Checksum>>> discoveryNestedRegistryFileHashes;
Selection.Result selectionResult;
}

Expand Down
13 changes: 11 additions & 2 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelModuleResolutionValue.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,12 @@

import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
import com.google.devtools.build.lib.skyframe.SkyFunctions;
import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
import com.google.devtools.build.skyframe.SkyKey;
import com.google.devtools.build.skyframe.SkyValue;
import java.util.Optional;

/**
* The result of the selection process, containing both the pruned and the un-pruned dependency
Expand All @@ -46,9 +48,16 @@ abstract class BazelModuleResolutionValue implements SkyValue {
*/
abstract ImmutableMap<ModuleKey, InterimModule> getUnprunedDepGraph();

/**
* Hashes of files obtained (or known to be missing) from registries while performing resolution.
*/
abstract ImmutableMap<String, Optional<Checksum>> getRegistryFileHashes();

static BazelModuleResolutionValue create(
ImmutableMap<ModuleKey, Module> resolvedDepGraph,
ImmutableMap<ModuleKey, InterimModule> unprunedDepGraph) {
return new AutoValue_BazelModuleResolutionValue(resolvedDepGraph, unprunedDepGraph);
ImmutableMap<ModuleKey, InterimModule> unprunedDepGraph,
ImmutableMap<String, Optional<Checksum>> registryFileHashes) {
return new AutoValue_BazelModuleResolutionValue(
resolvedDepGraph, unprunedDepGraph, registryFileHashes);
}
}
9 changes: 8 additions & 1 deletion src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Discovery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,11 @@

import static java.util.stream.Collectors.joining;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.bazel.bzlmod.InterimModule.DepSpec;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
import com.google.devtools.build.lib.server.FailureDetails;
import com.google.devtools.build.skyframe.SkyFunction.Environment;
import com.google.devtools.build.skyframe.SkyKey;
Expand All @@ -31,6 +33,7 @@
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Queue;
import java.util.Set;
import javax.annotation.Nullable;
Expand All @@ -49,7 +52,9 @@ private Discovery() {}
*/
@Nullable
public static ImmutableMap<ModuleKey, InterimModule> run(
Environment env, RootModuleFileValue root)
Environment env,
RootModuleFileValue root,
ImmutableList.Builder<ImmutableMap<String, Optional<Checksum>>> nestedRegistryFileHashes)
throws InterruptedException, ExternalDepsException {
String rootModuleName = root.getModule().getName();
ImmutableMap<String, ModuleOverride> overrides = root.getOverrides();
Expand All @@ -60,6 +65,7 @@ public static ImmutableMap<ModuleKey, InterimModule> run(
.withDepSpecsTransformed(InterimModule.applyOverrides(overrides, rootModuleName)));
Queue<ModuleKey> unexpanded = new ArrayDeque<>();
Map<ModuleKey, ModuleKey> predecessors = new HashMap<>();
nestedRegistryFileHashes.add(root.getRegistryFileHashes());
unexpanded.add(ModuleKey.ROOT);
while (!unexpanded.isEmpty()) {
Set<SkyKey> unexpandedSkyKeys = new HashSet<>();
Expand Down Expand Up @@ -109,6 +115,7 @@ public static ImmutableMap<ModuleKey, InterimModule> run(
.getModule()
.withDepSpecsTransformed(
InterimModule.applyOverrides(overrides, rootModuleName)));
nestedRegistryFileHashes.add(moduleFileValue.getRegistryFileHashes());
unexpanded.add(depKey);
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ private Optional<byte[]> grabFile(
maybeContent = Optional.empty();
}
if (knownFileHashesMode == KnownFileHashesMode.USE_AND_UPDATE && useChecksum) {
eventHandler.post(new RegistryFileDownloadEvent(url, maybeContent));
eventHandler.post(RegistryFileDownloadEvent.create(url, maybeContent));
}
return maybeContent;
}
Expand Down
8 changes: 4 additions & 4 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
import com.google.devtools.build.lib.bazel.bzlmod.CompiledModuleFile.IncludeStatement;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.NonRootModuleFileValue;
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.cmdline.LabelConstants;
import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
Expand Down Expand Up @@ -159,8 +158,6 @@ public SkyValue compute(SkyKey skyKey, Environment env)
if (getModuleFileResult == null) {
return null;
}
// Download events are retained by Skyframe and must thus be emitted in the exact same form
// across restarts.
getModuleFileResult.downloadEvents.replayOn(env.getListener());
String moduleFileHash =
new Fingerprint().addBytes(getModuleFileResult.moduleFile.getContent()).hexDigestAndReset();
Expand Down Expand Up @@ -222,7 +219,10 @@ public SkyValue compute(SkyKey skyKey, Environment env)
module.getVersion());
}

return NonRootModuleFileValue.create(module, moduleFileHash);
return NonRootModuleFileValue.create(
module,
moduleFileHash,
RegistryFileDownloadEvent.collectToMap(getModuleFileResult.downloadEvents.getPosts()));
}

@Nullable
Expand Down
Loading

0 comments on commit 189282d

Please sign in to comment.