Skip to content

Commit

Permalink
[7.2.0] Replace most Bzlmod events with a Skyframe graph lookup
Browse files Browse the repository at this point in the history 
Retaining `Postable`s across the graph increases the memory usage of every dependent Skyframe node.

`ModuleExtensionResolutionEvent` is replaced with a lookup of the "done" Skyframe nodes for extension evaluation in `BazelLockFileModule#afterCommand`.

`RootModuleFileFixupEvent` is replaced with direct storage in `SingleExtensionEvalValue`. The validation of repos imported from an extension is moved into a separate `SkyFunction`. This simplifies `bazel mod tidy`, which no longer needs to swallow certain exceptions, and allows cache hits when only the (invalid) imports of an extension are modified even without the lockfile.

The `BazelModuleResolutionEvent` is kept for now, but will become obsolete with the new lockfile format.

Work towards #20369

Closes #22058.

PiperOrigin-RevId: 628213907
Change-Id: I8ba19f5151a8183be5051c8d9280f93476db2272
  • Loading branch information
@fmeum @Wyverald
fmeum authored and Wyverald committed May 6, 2024
1 parent fed3603 commit af5fe43
Show file tree
Hide file tree
Showing 25 changed files with 564 additions and 328 deletions.
2 changes: 2 additions & 0 deletions src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@
import com.google.devtools.build.lib.bazel.bzlmod.RegistryFunction;
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.SingleExtensionFunction;
import com.google.devtools.build.lib.bazel.bzlmod.SingleExtensionUsagesFunction;
import com.google.devtools.build.lib.bazel.bzlmod.YankedVersionsFunction;
import com.google.devtools.build.lib.bazel.bzlmod.YankedVersionsUtil;
Expand Down Expand Up @@ -264,6 +265,7 @@ SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace
.addSkyFunction(SkyFunctions.BAZEL_MOD_TIDY, new BazelModTidyFunction())
.addSkyFunction(SkyFunctions.BAZEL_MODULE_INSPECTION, new BazelModuleInspectorFunction())
.addSkyFunction(SkyFunctions.BAZEL_MODULE_RESOLUTION, new BazelModuleResolutionFunction())
.addSkyFunction(SkyFunctions.SINGLE_EXTENSION, new SingleExtensionFunction())
.addSkyFunction(SkyFunctions.SINGLE_EXTENSION_EVAL, singleExtensionEvalFunction)
.addSkyFunction(SkyFunctions.SINGLE_EXTENSION_USAGES, new SingleExtensionUsagesFunction())
.addSkyFunction(
Expand Down
16 changes: 10 additions & 6 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 @@ -106,6 +106,7 @@ java_library(
"//src/main/java/com/google/devtools/build/lib/cmdline",
"//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
"//src/main/java/com/google/devtools/build/lib/vfs",
"//src/main/java/com/google/devtools/build/skyframe",
"//third_party:flogger",
"//third_party:guava",
"//third_party:jsr305",
Expand Down Expand Up @@ -133,7 +134,6 @@ java_library(
"ModuleExtensionEvalFactors.java",
"ModuleExtensionEvalStarlarkThreadContext.java",
"ModuleExtensionRepoMappingEntriesValue.java",
"ModuleExtensionResolutionEvent.java",
"ModuleFileGlobals.java",
"ModuleFileValue.java",
"ModuleOverride.java",
Expand All @@ -143,8 +143,8 @@ java_library(
"RegistryKey.java",
"RegistryOverride.java",
"RepoSpecKey.java",
"SingleExtensionEvalValue.java",
"SingleExtensionUsagesValue.java",
"SingleExtensionValue.java",
"SingleVersionOverride.java",
"YankedVersionsValue.java",
],
Expand All @@ -156,6 +156,7 @@ java_library(
":module_extension_metadata",
":registry",
":repo_rule_creator",
":root_module_file_fixup",
"//src/main/java/com/google/devtools/build/docgen/annot",
"//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
"//src/main/java/com/google/devtools/build/lib/cmdline",
Expand Down Expand Up @@ -199,6 +200,7 @@ java_library(
"RepoSpecFunction.java",
"Selection.java",
"SingleExtensionEvalFunction.java",
"SingleExtensionFunction.java",
"SingleExtensionUsagesFunction.java",
"StarlarkBazelModule.java",
"TypeCheckedTag.java",
Expand All @@ -214,6 +216,7 @@ java_library(
":registry",
":repo_rule_creator",
":resolution",
":root_module_file_fixup",
"//src/main/java/com/google/devtools/build/docgen/annot",
"//src/main/java/com/google/devtools/build/lib:runtime",
"//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
Expand Down Expand Up @@ -293,6 +296,7 @@ java_library(
],
deps = [
":resolution",
":root_module_file_fixup",
"//src/main/java/com/google/devtools/build/lib/bazel/repository:repository_options",
"//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
"//src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec:serialization-constant",
Expand All @@ -314,6 +318,7 @@ java_library(
":exception",
":resolution",
":resolution_impl",
":root_module_file_fixup",
":tidy",
"//src/main/java/com/google/devtools/build/lib/cmdline",
"//src/main/java/com/google/devtools/build/lib/rules:repository/repository_function",
Expand Down Expand Up @@ -371,7 +376,7 @@ java_library(
deps = [
":common",
":module_extension",
":module_file_fixup_event",
":root_module_file_fixup",
"//src/main/java/com/google/devtools/build/docgen/annot",
"//src/main/java/com/google/devtools/build/lib/cmdline",
"//src/main/java/com/google/devtools/build/lib/events",
Expand All @@ -386,13 +391,12 @@ java_library(
)

java_library(
name = "module_file_fixup_event",
name = "root_module_file_fixup",
srcs = [
"RootModuleFileFixupEvent.java",
"RootModuleFileFixup.java",
],
deps = [
":module_extension",
"//src/main/java/com/google/devtools/build/lib/events",
"//third_party:guava",
],
)
16 changes: 8 additions & 8 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelFetchAllFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,16 +57,16 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
// 2. Run every extension found in the modules & collect its generated repos
ImmutableSet<ModuleExtensionId> extensionIds =
depGraphValue.getExtensionUsagesTable().rowKeySet();
ImmutableSet<SkyKey> singleEvalKeys =
extensionIds.stream().map(SingleExtensionEvalValue::key).collect(toImmutableSet());
SkyframeLookupResult singleEvalValues = env.getValuesAndExceptions(singleEvalKeys);
for (SkyKey singleEvalKey : singleEvalKeys) {
SingleExtensionEvalValue singleEvalValue =
(SingleExtensionEvalValue) singleEvalValues.get(singleEvalKey);
if (singleEvalValue == null) {
ImmutableSet<SkyKey> singleExtensionKeys =
extensionIds.stream().map(SingleExtensionValue::key).collect(toImmutableSet());
SkyframeLookupResult singleExtensionValues = env.getValuesAndExceptions(singleExtensionKeys);
for (SkyKey singleExtensionKey : singleExtensionKeys) {
SingleExtensionValue singleExtensionValue =
(SingleExtensionValue) singleExtensionValues.get(singleExtensionKey);
if (singleExtensionValue == null) {
return null;
}
reposToFetch.addAll(singleEvalValue.getCanonicalRepoNameToInternalNames().keySet());
reposToFetch.addAll(singleExtensionValue.getCanonicalRepoNameToInternalNames().keySet());
}

// 3. If this is fetch configure, get repo rules and only collect repos marked as configure
Expand Down
98 changes: 57 additions & 41 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 @@ -16,7 +16,6 @@

import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSortedMap;
import com.google.common.eventbus.Subscribe;
Expand All @@ -31,10 +30,12 @@
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.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import javax.annotation.Nullable;

/**
Expand All @@ -43,15 +44,15 @@
*/
public class BazelLockFileModule extends BlazeModule {

private MemoizingEvaluator evaluator;
private Path workspaceRoot;
@Nullable private BazelModuleResolutionEvent moduleResolutionEvent;
private final Map<ModuleExtensionId, ModuleExtensionResolutionEvent>
extensionResolutionEventsMap = new HashMap<>();

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

@Override
public void beforeCommand(CommandEnvironment env) {
evaluator = env.getSkyframeExecutor().getEvaluator();
workspaceRoot = env.getWorkspace();
RepositoryOptions options = env.getOptions().getOptions(RepositoryOptions.class);
if (options.lockfileMode.equals(LockfileMode.UPDATE)) {
Expand All @@ -64,26 +65,38 @@ public void afterCommand() throws AbruptExitException {
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.
Preconditions.checkState(extensionResolutionEventsMap.isEmpty());
return;
}

// All nodes corresponding to module extensions that have been evaluated in the current build
// are done at this point. Look up entries by eval keys to record results even if validation
// later fails due to invalid imports.
// Note: This also picks up up-to-date result from previous builds that are not in the
// transitive closure of the current build. Since extension are potentially costly to evaluate,
// this is seen as an advantage. Full reproducibility can be ensured by running 'bazel shutdown'
// first if needed.
Map<ModuleExtensionId, LockFileModuleExtension.WithFactors> newExtensionInfos =
new ConcurrentHashMap<>();
evaluator
.getInMemoryGraph()
.parallelForEach(
entry -> {
if (entry.isDone()
&& entry.getKey() instanceof SingleExtensionValue.EvalKey key
// entry.getValue() can be null if the extension evaluation failed.
&& entry.getValue() instanceof SingleExtensionValue value) {
newExtensionInfos.put(key.argument(), value.getLockFileInfo().get());
}
});

BazelLockFileValue oldLockfile = moduleResolutionEvent.getOnDiskLockfileValue();
BazelLockFileValue newLockfile;
try {
// Create an updated version of the lockfile, keeping only the extension results from the old
// lockfile that are still up-to-date and adding the newly resolved extension results.
newLockfile =
moduleResolutionEvent.getResolutionOnlyLockfileValue().toBuilder()
.setModuleExtensions(combineModuleExtensions(oldLockfile))
.build();
} catch (ExternalDepsException e) {
logger.atSevere().withCause(e).log(
"Failed to read and parse the MODULE.bazel.lock file with error: %s."
+ " Try deleting it and rerun the build.",
e.getMessage());
return;
}
// Create an updated version of the lockfile, keeping only the extension results from the old
// lockfile that are still up-to-date and adding the newly resolved extension results.
BazelLockFileValue newLockfile =
moduleResolutionEvent.getResolutionOnlyLockfileValue().toBuilder()
.setModuleExtensions(
combineModuleExtensions(oldLockfile.getModuleExtensions(), newExtensionInfos))
.build();

// Write the new value to the file, but only if needed. This is not just a performance
// optimization: whenever the lockfile is updated, most Skyframe nodes will be marked as dirty
Expand All @@ -93,7 +106,6 @@ public void afterCommand() throws AbruptExitException {
updateLockfile(workspaceRoot, newLockfile);
}
this.moduleResolutionEvent = null;
this.extensionResolutionEventsMap.clear();
}

/**
Expand All @@ -102,12 +114,17 @@ public void afterCommand() throws AbruptExitException {
*/
private ImmutableMap<
ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
combineModuleExtensions(BazelLockFileValue oldLockfile) throws ExternalDepsException {
combineModuleExtensions(
ImmutableMap<
ModuleExtensionId,
ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
oldExtensionInfos,
Map<ModuleExtensionId, LockFileModuleExtension.WithFactors> newExtensionInfos) {
Map<ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
updatedExtensionMap = new HashMap<>();

// Keep old extensions if they are still valid.
for (var entry : oldLockfile.getModuleExtensions().entrySet()) {
for (var entry : oldExtensionInfos.entrySet()) {
var moduleExtensionId = entry.getKey();
var factorToLockedExtension = entry.getValue();
ModuleExtensionEvalFactors firstEntryFactors =
Expand All @@ -117,32 +134,36 @@ public void afterCommand() throws AbruptExitException {
// All entries for a single extension share the same usages digest, so it suffices to check
// the first entry.
if (shouldKeepExtension(
moduleExtensionId, firstEntryFactors, firstEntryExtension.getUsagesDigest())) {
moduleExtensionId,
firstEntryFactors,
firstEntryExtension.getUsagesDigest(),
newExtensionInfos.get(moduleExtensionId))) {
updatedExtensionMap.put(moduleExtensionId, factorToLockedExtension);
}
}

// Add the new resolved extensions
for (var event : extensionResolutionEventsMap.values()) {
LockFileModuleExtension extension = event.getModuleExtension();
for (var extensionIdAndInfo : newExtensionInfos.entrySet()) {
LockFileModuleExtension extension = extensionIdAndInfo.getValue().moduleExtension();
if (!extension.shouldLockExtension()) {
continue;
}

var oldExtensionEntries = updatedExtensionMap.get(event.getExtensionId());
var oldExtensionEntries = updatedExtensionMap.get(extensionIdAndInfo.getKey());
ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension> extensionEntries;
var factors = extensionIdAndInfo.getValue().extensionFactors();
if (oldExtensionEntries != null) {
// extension exists, add the new entry to the existing map
extensionEntries =
new ImmutableMap.Builder<ModuleExtensionEvalFactors, LockFileModuleExtension>()
.putAll(oldExtensionEntries)
.put(event.getExtensionFactors(), extension)
.put(factors, extension)
.buildKeepingLast();
} else {
// new extension
extensionEntries = ImmutableMap.of(event.getExtensionFactors(), extension);
extensionEntries = ImmutableMap.of(factors, extension);
}
updatedExtensionMap.put(event.getExtensionId(), extensionEntries);
updatedExtensionMap.put(extensionIdAndInfo.getKey(), extensionEntries);
}

// The order in which extensions are added to extensionResolutionEvents depends on the order
Expand All @@ -167,17 +188,18 @@ public void afterCommand() throws AbruptExitException {
private boolean shouldKeepExtension(
ModuleExtensionId extensionId,
ModuleExtensionEvalFactors lockedExtensionKey,
byte[] oldUsagesDigest) {
byte[] oldUsagesDigest,
@Nullable LockFileModuleExtension.WithFactors newExtensionInfo) {

// If there is a new event for this extension, compare it with the existing ones
ModuleExtensionResolutionEvent extEvent = extensionResolutionEventsMap.get(extensionId);
if (extEvent != null) {
boolean doNotLockExtension = !extEvent.getModuleExtension().shouldLockExtension();
if (newExtensionInfo != null) {
boolean doNotLockExtension = !newExtensionInfo.moduleExtension().shouldLockExtension();
boolean dependencyOnOsChanged =
lockedExtensionKey.getOs().isEmpty() != extEvent.getExtensionFactors().getOs().isEmpty();
lockedExtensionKey.getOs().isEmpty()
!= newExtensionInfo.extensionFactors().getOs().isEmpty();
boolean dependencyOnArchChanged =
lockedExtensionKey.getArch().isEmpty()
!= extEvent.getExtensionFactors().getArch().isEmpty();
!= newExtensionInfo.extensionFactors().getArch().isEmpty();
if (doNotLockExtension || dependencyOnOsChanged || dependencyOnArchChanged) {
return false;
}
Expand Down Expand Up @@ -228,10 +250,4 @@ public void bazelModuleResolved(BazelModuleResolutionEvent moduleResolutionEvent
// sent after the command has modified the module file.
this.moduleResolutionEvent = moduleResolutionEvent;
}

@Subscribe
public void moduleExtensionResolved(ModuleExtensionResolutionEvent extensionResolutionEvent) {
this.extensionResolutionEventsMap.put(
extensionResolutionEvent.getExtensionId(), extensionResolutionEvent);
}
}
Loading

0 comments on commit af5fe43

Please sign in to comment.