Remove stale extension entries from lockfile

When changes to module files result in module resolution rerunning,
extension results from the lockfile may become stale due to changed
usages even if the current build doesn't evaluate the affected
extension.

We now compare the current usages against the usages recorded in the
lockfile and drop all extensions that became stale.

Along the way, also trim tag locations when comparing usages for
staleness detections as they do not affect the evaluation result.

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

@@ -94,6 +94,8 @@ java_library(
     name = "bazel_lockfile_module",
     srcs = ["BazelLockFileModule.java"],
     deps = [
+        ":common",
+        ":exception",
         ":resolution",
         ":resolution_impl",
         "//src/main/java/com/google/devtools/build/lib:runtime",

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

@@ -128,6 +128,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
         calculateUniqueNameForUsedExtensionId(extensionUsagesById);
 
     if (!lockfileMode.equals(LockfileMode.OFF)) {
+      // This will keep all module extension evaluation results, some of which may be stale due to
+      // changed usages. They will be removed in BazelLockFileModule.
       BazelLockFileValue updateLockfile =
           lockfile.toBuilder()
               .setLockFileVersion(BazelLockFileValue.LOCK_FILE_VERSION)

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

@@ -18,6 +18,8 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.ImmutableTable;
+import com.google.common.collect.Maps;
 import com.google.common.eventbus.Subscribe;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
@@ -67,25 +69,42 @@ public void afterCommand() throws AbruptExitException {
     RootedPath lockfilePath =
         RootedPath.toRootedPath(Root.fromPath(workspaceRoot), LabelConstants.MODULE_LOCKFILE_NAME);
 
+    // Read the existing lockfile (if none exists, will get an empty lockfile value) and get its
+    // module extension usages. This information is needed to determine which extension results are
+    // now stale and need to be removed.
+    BazelLockFileValue oldLockfile;
+    try {
+      oldLockfile = BazelLockFileFunction.getLockfileValue(lockfilePath);
+    } catch (IOException | JsonSyntaxException | NullPointerException 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;
+    }
+    ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages;
+    try {
+      oldExtensionUsages =
+          BazelDepGraphFunction.getExtensionUsagesById(oldLockfile.getModuleDepGraph());
+    } 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 with the events updates
     BazelLockFileValue lockfile;
     if (moduleResolutionEvent != null) {
       lockfile = moduleResolutionEvent.getLockfileValue();
     } else {
-      // Read the existing lockfile (if none exists, will get an empty lockfile value)
-      try {
-        lockfile = BazelLockFileFunction.getLockfileValue(lockfilePath);
-      } catch (IOException | JsonSyntaxException | NullPointerException 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;
-      }
+      lockfile = oldLockfile;
     }
     lockfile =
         lockfile.toBuilder()
-            .setModuleExtensions(combineModuleExtensions(lockfile.getModuleExtensions()))
+            .setModuleExtensions(
+                combineModuleExtensions(lockfile.getModuleExtensions(), oldExtensionUsages))
             .build();
 
     // Write the new value to the file
@@ -99,14 +118,16 @@ public void afterCommand() throws AbruptExitException {
    * extensions from the events (if any)
    *
    * @param oldModuleExtensions Module extensions stored in the current lockfile
+   * @param oldExtensionUsages Module extension usages stored in the current lockfile
    */
   private ImmutableMap<
           ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
       combineModuleExtensions(
           ImmutableMap<
                   ModuleExtensionId,
                   ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
-              oldModuleExtensions) {
+              oldModuleExtensions,
+          ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages) {
 
     Map<ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
         updatedExtensionMap = new HashMap<>();
@@ -115,7 +136,7 @@ public void afterCommand() throws AbruptExitException {
     oldModuleExtensions.forEach(
         (moduleExtensionId, innerMap) -> {
           ModuleExtensionEvalFactors firstEntryKey = innerMap.keySet().iterator().next();
-          if (shouldKeepExtension(moduleExtensionId, firstEntryKey)) {
+          if (shouldKeepExtension(moduleExtensionId, firstEntryKey, oldExtensionUsages)) {
             updatedExtensionMap.put(moduleExtensionId, innerMap);
           }
         });
@@ -146,14 +167,21 @@ public void afterCommand() throws AbruptExitException {
   }
 
   /**
-   * Decide whether to keep this extension or not depending on both: 1. If its dependency on os &
-   * arch didn't change 2. If it is still has a usage in the module
+   * Decide whether to keep this extension or not depending on all of:
+   *
+   * <ol>
+   *   <li>If its dependency on os & arch didn't change
+   *   <li>If its usages haven't changed
+   * </ol>
    *
    * @param lockedExtensionKey object holding the old extension id and state of os and arch
+   * @param oldExtensionUsages the usages of this extension in the existing lockfile
    * @return True if this extension should still be in lockfile, false otherwise
    */
   private boolean shouldKeepExtension(
-      ModuleExtensionId extensionId, ModuleExtensionEvalFactors lockedExtensionKey) {
+      ModuleExtensionId extensionId,
+      ModuleExtensionEvalFactors lockedExtensionKey,
+      ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages) {
 
     // If there is a new event for this extension, compare it with the existing ones
     ModuleExtensionResolutionEvent extEvent = extensionResolutionEventsMap.get(extensionId);
@@ -168,11 +196,24 @@ private boolean shouldKeepExtension(
       }
     }
 
-    // If moduleResolutionEvent is null, then no usage has changed. So we don't need this check
-    if (moduleResolutionEvent != null) {
-      return moduleResolutionEvent.getExtensionUsagesById().containsRow(extensionId);
+    // If moduleResolutionEvent is null, then no usage has changed and all locked extension
+    // resolutions are still up-to-date.
+    if (moduleResolutionEvent == null) {
+      return true;
     }
-    return true;
+    // Otherwise, compare the current usages of this extension with the ones in the lockfile. We
+    // trim the usages to only the information that influences the evaluation of the extension so
+    // that irrelevant changes (e.g. locations or imports) don't cause the extension to be removed.
+    // Note: Extension results can still be stale for other reasons, e.g. because their transitive
+    // bzl hash changed, but such changes will be detected in SingleExtensionEvalFunction.
+    var currentTrimmedUsages =
+        Maps.transformValues(
+            moduleResolutionEvent.getExtensionUsagesById().row(extensionId),
+            ModuleExtensionUsage::trimForEvaluation);
+    var lockedTrimmedUsages =
+        Maps.transformValues(
+            oldExtensionUsages.row(extensionId), ModuleExtensionUsage::trimForEvaluation);
+    return currentTrimmedUsages.equals(lockedTrimmedUsages);
   }
 
   /**

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

@@ -14,6 +14,8 @@
 
 package com.google.devtools.build.lib.bazel.bzlmod;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
@@ -26,6 +28,8 @@
 /**
  * Represents one usage of a module extension in one MODULE.bazel file. This class records all the
  * information pertinent to the proxy object returned from the {@code use_extension} call.
+ *
+ * <p>When adding new fields, make sure to update {@link #trimForEvaluation()} as well.
  */
 @AutoValue
 @GenerateTypeAdapter
@@ -86,6 +90,26 @@ public static Builder builder() {
     return new AutoValue_ModuleExtensionUsage.Builder();
   }
 
+  /**
+   * Returns a new usage with all information removed that does not influence the evaluation of the
+   * extension.
+   */
+  ModuleExtensionUsage trimForEvaluation() {
+    // We start with the full usage and selectively remove information that does not influence the
+    // evaluation of the extension. Compared to explicitly copying over the parts that do, this
+    // preserves correctness in case new fields are added without updating this code.
+    return toBuilder()
+        .setTags(getTags().stream().map(Tag::trimForEvaluation).collect(toImmutableList()))
+        // Locations are only used for error reporting and thus don't influence whether the
+        // evaluation of the extension is successful and what its result is in case of success.
+        .setLocation(Location.BUILTIN)
+        // Extension implementation functions do not see the imports, they are only validated
+        // against the set of generated repos in a validation step that comes afterward.
+        .setImports(ImmutableBiMap.of())
+        .setDevImports(ImmutableSet.of())
+        .build();
+  }
+
   /** Builder for {@link ModuleExtensionUsage}. */
   @AutoValue.Builder
   public abstract static class Builder {

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

@@ -18,14 +18,13 @@
 import static com.google.common.base.StandardSystemProperty.OS_ARCH;
 import static com.google.common.collect.ImmutableBiMap.toImmutableBiMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.Maps.transformValues;
 
 import com.google.auto.value.AutoValue;
-import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableTable;
-import com.google.common.collect.Maps;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
@@ -274,8 +273,13 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
 
     // Check extension data in lockfile is still valid, disregarding usage information that is not
     // relevant for the evaluation of the extension.
-    var trimmedLockedUsages = trimUsagesForEvaluation(lockedExtensionUsages);
-    var trimmedUsages = trimUsagesForEvaluation(usagesValue.getExtensionUsages());
+    var trimmedLockedUsages =
+        ImmutableMap.copyOf(
+            transformValues(lockedExtensionUsages, ModuleExtensionUsage::trimForEvaluation));
+    var trimmedUsages =
+        ImmutableMap.copyOf(
+            transformValues(
+                usagesValue.getExtensionUsages(), ModuleExtensionUsage::trimForEvaluation));
     if (!filesChanged
         && Arrays.equals(bzlTransitiveDigest, lockedExtension.getBzlTransitiveDigest())
         && trimmedUsages.equals(trimmedLockedUsages)
@@ -414,33 +418,6 @@ private SingleExtensionEvalValue validateAndCreateSingleExtensionEvalValue(
                     Function.identity())));
   }
 
-  /**
-   * Returns usages with all information removed that does not influence the evaluation of the
-   * extension.
-   */
-  private static ImmutableMap<ModuleKey, ModuleExtensionUsage> trimUsagesForEvaluation(
-      Map<ModuleKey, ModuleExtensionUsage> usages) {
-    return ImmutableMap.copyOf(
-        Maps.transformValues(
-            usages,
-            usage ->
-                // We start with the full usage and selectively remove information that does not
-                // influence the evaluation of the extension. Compared to explicitly copying over
-                // the parts that do, this preserves correctness in case new fields are added to
-                // ModuleExtensionUsage without updating this code.
-                usage.toBuilder()
-                    // Locations are only used for error reporting and thus don't influence whether
-                    // the evaluation of the extension is successful and what its result is
-                    // in case of success.
-                    .setLocation(Location.BUILTIN)
-                    // Extension implementation functions do not see the imports, they are only
-                    // validated against the set of generated repos in a validation step that comes
-                    // afterward.
-                    .setImports(ImmutableBiMap.of())
-                    .setDevImports(ImmutableSet.of())
-                    .build()));
-  }
-
   private BzlLoadValue loadBzlFile(
       Label bzlFileLabel,
       Location sampleUsageLocation,

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

@@ -40,10 +40,27 @@ public abstract class Tag {
   /** The source location in the module file where this tag was created. */
   public abstract Location getLocation();
 
+  public abstract Builder toBuilder();
+
   public static Builder builder() {
     return new AutoValue_Tag.Builder();
   }
 
+  /**
+   * Returns a new tag with all information removed that does not influence the evaluation of the
+   * extension defining the tag.
+   */
+  Tag trimForEvaluation() {
+    // We start with the full usage and selectively remove information that does not influence the
+    // evaluation of the extension. Compared to explicitly copying over the parts that do, this
+    // preserves correctness in case new fields are added without updating this code.
+    return toBuilder()
+        // Locations are only used for error reporting and thus don't influence whether the
+        // evaluation of the extension is successful and what its result is in case of success.
+        .setLocation(Location.BUILTIN)
+        .build();
+  }
+
   /** Builder for {@link Tag}. */
   @AutoValue.Builder
   public abstract static class Builder {