Use a persistent lockfile to speed up reproducible extensions

site/en/external/lockfile.md

@@ -70,6 +70,15 @@ The lockfile offers several benefits and can be utilized in various ways:
     preventing unexpected updates or breaking changes in external libraries. By
     locking the dependencies to specific versions, the risk of introducing bugs
     due to incompatible or untested updates is reduced.
+-
+
+### Persistent lockfile {:#persistent-lockfile}
+
+Bazel also maintains another lockfile at
+`"$(bazel info output_base)"/MODULE.bazel.lock`. The format and contents of this
+lockfile are explicitly unspecified. It is only used as a performance
+optimization. While it can be deleted together with the output base, any need to
+do so is a bug in either Bazel itself or a module extension.
 
 ## Lockfile Contents {:#lockfile-contents}
 

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

@@ -266,7 +266,8 @@ public void workspaceInit(
         .addSkyFunction(SkyFunctions.MODULE_FILE, moduleFileFunction)
         .addSkyFunction(SkyFunctions.BAZEL_DEP_GRAPH, new BazelDepGraphFunction())
         .addSkyFunction(
-            SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace()))
+            SkyFunctions.BAZEL_LOCK_FILE,
+            new BazelLockFileFunction(directories.getWorkspace(), directories.getOutputBase()))
         .addSkyFunction(SkyFunctions.BAZEL_FETCH_ALL, new BazelFetchAllFunction())
         .addSkyFunction(SkyFunctions.BAZEL_MOD_TIDY, new BazelModTidyFunction())
         .addSkyFunction(SkyFunctions.BAZEL_MODULE_INSPECTION, new BazelModuleInspectorFunction())

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

@@ -125,6 +125,7 @@ java_library(
         ":resolution_impl",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository:repository_options",
+        "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_cluster",

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

@@ -41,7 +41,11 @@
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
-/** Reads the contents of the lock file into its value. */
+/**
+ * Reads the contents of the lockfiles into {@link BazelLockFileValue}s.
+ *
+ * See {@link BazelLockFileValue} for more information.
+ */
 public class BazelLockFileFunction implements SkyFunction {
 
   public static final Precomputed<LockfileMode> LOCKFILE_MODE = new Precomputed<>("lockfile_mode");
@@ -55,29 +59,42 @@ public class BazelLockFileFunction implements SkyFunction {
   private static final BazelLockFileValue EMPTY_LOCKFILE = BazelLockFileValue.builder().build();
 
   private final Path rootDirectory;
+  private final Path outputBase;
 
-  public BazelLockFileFunction(Path rootDirectory) {
+  public BazelLockFileFunction(Path rootDirectory, Path outputBase) {
     this.rootDirectory = rootDirectory;
+    this.outputBase = outputBase;
   }
 
   @Override
   @Nullable
   public SkyValue compute(SkyKey skyKey, Environment env)
       throws BazelLockfileFunctionException, InterruptedException {
+    boolean forPersistentLockfile = skyKey == BazelLockFileValue.PERSISTENT_KEY;
     RootedPath lockfilePath =
-        RootedPath.toRootedPath(Root.fromPath(rootDirectory), LabelConstants.MODULE_LOCKFILE_NAME);
+        RootedPath.toRootedPath(
+            Root.fromPath(forPersistentLockfile ? outputBase : rootDirectory),
+            LabelConstants.MODULE_LOCKFILE_NAME);
 
-    // Add dependency on the lockfile to recognize changes to it
+    // Add dependency on the lockfiles to recognize changes to it
     if (env.getValue(FileValue.key(lockfilePath)) == null) {
       return null;
     }
 
-    try (SilentCloseable c = Profiler.instance().profile(ProfilerTask.BZLMOD, "parse lockfile")) {
-      return getLockfileValue(lockfilePath, LOCKFILE_MODE.get(env));
+    try (SilentCloseable c =
+        Profiler.instance()
+            .profile(
+                ProfilerTask.BZLMOD,
+                forPersistentLockfile ? "parse persistent lockfile" : "parse lockfile")) {
+      return getLockfileValue(
+          lockfilePath, forPersistentLockfile ? LockfileMode.UPDATE : LOCKFILE_MODE.get(env));
     } catch (IOException
         | JsonSyntaxException
         | NullPointerException
         | IllegalArgumentException e) {
+      if (forPersistentLockfile) {
+        return EMPTY_LOCKFILE;
+      }
       String actionSuffix;
       if (POSSIBLE_MERGE_CONFLICT_PATTERN.matcher(e.getMessage()).find()) {
         actionSuffix =

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

@@ -17,11 +17,16 @@
 import static com.google.devtools.build.lib.bazel.bzlmod.BazelLockFileFunction.LOCKFILE_MODE;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Sets;
 import com.google.common.flogger.GoogleLogger;
+import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
+import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.runtime.BlazeModule;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
@@ -35,7 +40,9 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Predicate;
 
 /**
  * Module collecting Bazel module and module extensions resolution results and updating the
@@ -45,13 +52,20 @@ public class BazelLockFileModule extends BlazeModule {
 
   private SkyframeExecutor executor;
   private Path workspaceRoot;
+  private Path outputBase;
+  private LockfileMode optionsLockfileMode;
 
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
+  private static final ImmutableSet<LockfileMode> ENABLED_IN_MODES =
+      Sets.immutableEnumSet(LockfileMode.UPDATE, LockfileMode.REFRESH);
+
   @Override
   public void beforeCommand(CommandEnvironment env) {
     executor = env.getSkyframeExecutor();
     workspaceRoot = env.getWorkspace();
+    outputBase = env.getOutputBase();
+    optionsLockfileMode = env.getOptions().getOptions(RepositoryOptions.class).lockfileMode;
   }
 
   @Override
@@ -60,29 +74,36 @@ public void afterCommand() {
     BazelModuleResolutionValue moduleResolutionValue;
     BazelDepGraphValue depGraphValue;
     BazelLockFileValue oldLockfile;
+    BazelLockFileValue oldPersistentLockfile;
     try {
       PrecomputedValue lockfileModeValue =
           (PrecomputedValue) evaluator.getExistingValue(LOCKFILE_MODE.getKey());
       if (lockfileModeValue == null) {
         // No command run on this server has triggered module resolution yet.
         return;
       }
-      LockfileMode lockfileMode = (LockfileMode) lockfileModeValue.get();
-      // Check the Skyframe value instead of the option since some commands (e.g. shutdown) don't
-      // propagate the options to Skyframe, but we can only operate on Skyframe values that were
-      // generated in UPDATE mode.
-      if (lockfileMode != LockfileMode.UPDATE && lockfileMode != LockfileMode.REFRESH) {
+      // Check the Skyframe value in addition to the option since some commands (e.g. shutdown)
+      // don't propagate the options to Skyframe, but we can only operate on Skyframe values that
+      // were generated in UPDATE mode.
+      LockfileMode skyframeLockfileMode = (LockfileMode) lockfileModeValue.get();
+      if (!(ENABLED_IN_MODES.contains(optionsLockfileMode)
+          && ENABLED_IN_MODES.contains(skyframeLockfileMode))) {
         return;
       }
       moduleResolutionValue =
           (BazelModuleResolutionValue) evaluator.getExistingValue(BazelModuleResolutionValue.KEY);
       depGraphValue = (BazelDepGraphValue) evaluator.getExistingValue(BazelDepGraphValue.KEY);
       oldLockfile = (BazelLockFileValue) evaluator.getExistingValue(BazelLockFileValue.KEY);
+      oldPersistentLockfile =
+          (BazelLockFileValue) evaluator.getExistingValue(BazelLockFileValue.PERSISTENT_KEY);
     } catch (InterruptedException e) {
       // Not thrown in Bazel.
       throw new IllegalStateException(e);
     }
-    if (moduleResolutionValue == null || depGraphValue == null || oldLockfile == null) {
+    if (moduleResolutionValue == null
+        || depGraphValue == null
+        || oldLockfile == null
+        || oldPersistentLockfile == null) {
       // An error during the actual build prevented the evaluation of these values and has already
       // been reported at this point.
       return;
@@ -109,51 +130,113 @@ public void afterCommand() {
                 newExtensionInfos.put(key.argument(), value.lockFileInfo().get());
               }
             });
-    var combinedExtensionInfos =
-        combineModuleExtensions(
-            oldLockfile.getModuleExtensions(), newExtensionInfos, depGraphValue);
-
-    // 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 =
-        BazelLockFileValue.builder()
-            .setRegistryFileHashes(
-                ImmutableSortedMap.copyOf(moduleResolutionValue.getRegistryFileHashes()))
-            .setSelectedYankedVersions(moduleResolutionValue.getSelectedYankedVersions())
-            .setModuleExtensions(combinedExtensionInfos)
-            .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
-    // on the next build, which breaks commands such as `bazel config` that rely on
-    // com.google.devtools.build.skyframe.MemoizingEvaluator#getDoneValues.
-    if (!newLockfile.equals(oldLockfile)) {
-      updateLockfile(workspaceRoot, newLockfile);
+
+    Thread updateLockfile =
+        Thread.startVirtualThread(
+            () -> {
+              var notReproducibleExtensionInfos =
+                  combineModuleExtensions(
+                      oldLockfile.getModuleExtensions(),
+                      newExtensionInfos,
+                      /* hasUsages= */ depGraphValue.getExtensionUsagesTable()::containsRow,
+                      /* reproducible= */ false);
+
+              // 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, as long as any of them are not known to be reproducible.
+              BazelLockFileValue newLockfile =
+                  BazelLockFileValue.builder()
+                      .setRegistryFileHashes(
+                          ImmutableSortedMap.copyOf(moduleResolutionValue.getRegistryFileHashes()))
+                      .setSelectedYankedVersions(moduleResolutionValue.getSelectedYankedVersions())
+                      .setModuleExtensions(notReproducibleExtensionInfos)
+                      .build();
+
+              // Write the new values to the files, but only if needed. This is not just a
+              // performance optimization: whenever the lockfile is updated, most Skyframe nodes
+              // will be marked as dirty on the next build, which breaks commands such as `bazel
+              // config` that rely on
+              // com.google.devtools.build.skyframe.MemoizingEvaluator#getDoneValues.
+              if (!newLockfile.equals(oldLockfile)) {
+                updateLockfile(workspaceRoot, newLockfile);
+              }
+            });
+
+    Thread updatePersistentLockfile =
+        Thread.startVirtualThread(
+            () -> {
+              // Results of reproducible extensions do not need to be stored for reproducibility,
+              // but avoiding reevaluations on server startups helps cold build performance.
+              var reproducibleExtensionInfos =
+                  combineModuleExtensions(
+                      oldLockfile.getModuleExtensions(),
+                      newExtensionInfos,
+                      /* hasUsages= */ depGraphValue.getExtensionUsagesTable()::containsRow,
+                      /* reproducible= */ true);
+
+              var persistentRegistryFileHashes =
+                  ImmutableMap.<String, Optional<Checksum>>builder()
+                      // While the hashes in the regular lockfile are trimmed down to those needed
+                      // for the current build, the persistent lockfile does not need to be a
+                      // deterministic function of the current build and can thus keep the extra
+                      // hashes. This speeds up resolution when switching branches and also provides
+                      // better (although not guaranteed) protection against retroactive changes to
+                      // registries.
+                      // TODO: Consider moving this information to the "true repository cache" once
+                      //  it is implemented as it can be reused across workspaces.
+                      .putAll(oldPersistentLockfile.getRegistryFileHashes())
+                      // Do not add "not found" entries to the persistent lockfile as they may
+                      // become invalid over time, which would violate the contract of the
+                      // persistent lockfile.
+                      .putAll(
+                          Maps.filterValues(
+                              moduleResolutionValue.getRegistryFileHashes(), Optional::isPresent))
+                      .buildKeepingLast();
+              BazelLockFileValue newPersistentLockfile =
+                  BazelLockFileValue.builder()
+                      .setRegistryFileHashes(
+                          ImmutableSortedMap.copyOf(persistentRegistryFileHashes))
+                      .setSelectedYankedVersions(ImmutableMap.of())
+                      .setModuleExtensions(reproducibleExtensionInfos)
+                      .build();
+
+              if (!newPersistentLockfile.equals(oldPersistentLockfile)) {
+                updateLockfile(outputBase, newPersistentLockfile);
+              }
+            });
+
+    try {
+      updateLockfile.join();
+      updatePersistentLockfile.join();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      logger.atSevere().withCause(e).log(
+          "Interrupted while updating MODULE.bazel.lock file: %s", e.getMessage());
     }
   }
 
   /**
    * Combines the old extensions stored in the lockfile -if they are still valid- with the new
    * extensions from the events (if any)
    */
-  private ImmutableMap<
+  @VisibleForTesting
+  static ImmutableMap<
           ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
       combineModuleExtensions(
-          ImmutableMap<
-                  ModuleExtensionId,
-                  ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
+          Map<ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
               oldExtensionInfos,
           Map<ModuleExtensionId, LockFileModuleExtension.WithFactors> newExtensionInfos,
-          BazelDepGraphValue depGraphValue) {
+          Predicate<ModuleExtensionId> hasUsages,
+          boolean reproducible) {
     Map<ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
         updatedExtensionMap = new HashMap<>();
 
     // Keep those per factor extension results that are still used according to the static
-    // information given in the extension declaration (dependence on os and arch, reproducibility).
+    // information given in the extension declaration (dependence on os and arch).
     // Other information such as transitive .bzl hash and usages hash are *not* checked here.
     for (var entry : oldExtensionInfos.entrySet()) {
       var moduleExtensionId = entry.getKey();
-      if (!depGraphValue.getExtensionUsagesTable().containsRow(moduleExtensionId)) {
+      if (!hasUsages.test(moduleExtensionId)) {
         // Extensions without any usages are not needed anymore.
         continue;
       }
@@ -163,15 +246,18 @@ public void afterCommand() {
         updatedExtensionMap.put(moduleExtensionId, entry.getValue());
         continue;
       }
-      if (!newExtensionInfo.moduleExtension().shouldLockExtension()) {
-        // Extension has become reproducible and should not be locked anymore.
-        continue;
-      }
       var newFactors = newExtensionInfo.extensionFactors();
+      // Factor results can be individually marked as reproducible.
       ImmutableSortedMap<ModuleExtensionEvalFactors, LockFileModuleExtension>
           perFactorResultsToKeep =
               ImmutableSortedMap.copyOf(
-                  Maps.filterKeys(entry.getValue(), newFactors::hasSameDependenciesAs));
+                  Maps.filterKeys(
+                      entry.getValue(),
+                      existingFactors ->
+                          newFactors.hasSameDependenciesAs(existingFactors)
+                              && !(newFactors.equals(existingFactors)
+                                  && newExtensionInfo.moduleExtension().isReproducible()
+                                      != reproducible)));
       if (perFactorResultsToKeep.isEmpty()) {
         continue;
       }
@@ -181,7 +267,7 @@ public void afterCommand() {
     // Add the new resolved extensions
     for (var extensionIdAndInfo : newExtensionInfos.entrySet()) {
       LockFileModuleExtension extension = extensionIdAndInfo.getValue().moduleExtension();
-      if (!extension.shouldLockExtension()) {
+      if (extension.isReproducible() != reproducible) {
         continue;
       }
 
@@ -213,12 +299,12 @@ public void afterCommand() {
   /**
    * Updates the data stored in the lockfile (MODULE.bazel.lock)
    *
-   * @param workspaceRoot Root of the workspace where the lockfile is located
+   * @param lockfileRoot Root under which the lockfile is located
    * @param updatedLockfile The updated lockfile data to save
    */
-  private static void updateLockfile(Path workspaceRoot, BazelLockFileValue updatedLockfile) {
+  private static void updateLockfile(Path lockfileRoot, BazelLockFileValue updatedLockfile) {
     RootedPath lockfilePath =
-        RootedPath.toRootedPath(Root.fromPath(workspaceRoot), LabelConstants.MODULE_LOCKFILE_NAME);
+        RootedPath.toRootedPath(Root.fromPath(lockfileRoot), LabelConstants.MODULE_LOCKFILE_NAME);
     try {
       FileSystemUtils.writeContent(
           lockfilePath.asPath(),