Fix

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

@@ -16,12 +16,14 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
+import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 
 /**
  * An abridged version of a {@link Module}, with a reduced set of information available, used for
  * module extension resolution.
  */
 @AutoValue
+@GenerateTypeAdapter
 public abstract class AbridgedModule {
   public abstract String getName();
 

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

@@ -89,13 +89,22 @@ public void afterCommand() throws AbruptExitException {
               }
             });
 
+    BazelDepGraphValue depGraphValue;
+    try {
+      depGraphValue = (BazelDepGraphValue) evaluator.getExistingValue(BazelDepGraphValue.KEY);
+    } catch (InterruptedException e) {
+      // Not thrown in Bazel.
+      throw new IllegalStateException(e);
+    }
+
     BazelLockFileValue oldLockfile = moduleResolutionEvent.getOnDiskLockfileValue();
     // 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))
+                combineModuleExtensions(
+                    oldLockfile.getModuleExtensions(), newExtensionInfos, depGraphValue))
             .build();
 
     // Write the new value to the file, but only if needed. This is not just a performance
@@ -119,7 +128,8 @@ public void afterCommand() throws AbruptExitException {
                   ModuleExtensionId,
                   ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
               oldExtensionInfos,
-          Map<ModuleExtensionId, LockFileModuleExtension.WithFactors> newExtensionInfos) {
+          Map<ModuleExtensionId, LockFileModuleExtension.WithFactors> newExtensionInfos,
+          BazelDepGraphValue depGraphValue) {
     Map<ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
         updatedExtensionMap = new HashMap<>();
 
@@ -137,7 +147,8 @@ public void afterCommand() throws AbruptExitException {
           moduleExtensionId,
           firstEntryFactors,
           firstEntryExtension.getUsagesDigest(),
-          newExtensionInfos.get(moduleExtensionId))) {
+          newExtensionInfos.get(moduleExtensionId),
+          depGraphValue)) {
         updatedExtensionMap.put(moduleExtensionId, factorToLockedExtension);
       }
     }
@@ -183,13 +194,17 @@ public void afterCommand() throws AbruptExitException {
    *
    * @param lockedExtensionKey object holding the old extension id and state of os and arch
    * @param oldUsagesDigest the digest of usages of this extension in the existing lockfile
+   * @param newExtensionInfo the in-memory lockfile entry produced by the most recent up-to-date
+   *     evaluation of this extension (if any)
+   * @param depGraphValue the dep graph value
    * @return True if this extension should still be in lockfile, false otherwise
    */
   private boolean shouldKeepExtension(
       ModuleExtensionId extensionId,
       ModuleExtensionEvalFactors lockedExtensionKey,
       byte[] oldUsagesDigest,
-      @Nullable LockFileModuleExtension.WithFactors newExtensionInfo) {
+      @Nullable LockFileModuleExtension.WithFactors newExtensionInfo,
+      BazelDepGraphValue depGraphValue) {
 
     // If there is a new event for this extension, compare it with the existing ones
     if (newExtensionInfo != null) {
@@ -210,10 +225,15 @@ private boolean shouldKeepExtension(
     // 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 usagesValue = SingleExtensionUsagesFunction.createValue(depGraphValue, extensionId);
+    if (usagesValue.isEmpty()) {
+      // The extension is no longer used anywhere, drop it.
+      return false;
+    }
     return Arrays.equals(
-        ModuleExtensionUsage.hashForEvaluation(
-            GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
-            moduleResolutionEvent.getExtensionUsagesById().row(extensionId)),
+        SingleExtensionUsagesValue.hashForEvaluation(
+            GsonTypeAdapterUtil.createSingleExtensionUsagesValueHashGson(),
+            SingleExtensionUsagesFunction.createValue(depGraphValue, extensionId).get()),
         oldUsagesDigest);
   }
 

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

@@ -482,7 +482,7 @@ public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
         .create();
   }
 
-  public static Gson createModuleExtensionUsagesHashGson() {
+  public static Gson createSingleExtensionUsagesValueHashGson() {
     return newGsonBuilder().create();
   }
 

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

@@ -94,22 +94,11 @@ public static Builder builder() {
     return new AutoValue_ModuleExtensionUsage.Builder();
   }
 
-  /**
-   * Turns the given collection of usages for a particular extension into a hash that can be
-   * compared for equality with another hash obtained in this way and compares equal only if the two
-   * original collections of usages are equivalent for the purpose of evaluating the extension.
-   */
-  static byte[] hashForEvaluation(Gson gson, ImmutableMap<ModuleKey, ModuleExtensionUsage> usages) {
-    ImmutableMap<ModuleKey, ModuleExtensionUsage> trimmedUsages =
-        ImmutableMap.copyOf(Maps.transformValues(usages, ModuleExtensionUsage::trimForEvaluation));
-    return Hashing.sha256().hashUnencodedChars(gson.toJson(trimmedUsages)).asBytes();
-  }
-
   /**
    * Returns a new usage with all information removed that does not influence the evaluation of the
    * extension.
    */
-  private ModuleExtensionUsage trimForEvaluation() {
+  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.

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

@@ -224,9 +224,9 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                   LockFileModuleExtension.builder()
                       .setBzlTransitiveDigest(extension.getBzlTransitiveDigest())
                       .setUsagesDigest(
-                          ModuleExtensionUsage.hashForEvaluation(
-                              GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
-                              usagesValue.getExtensionUsages()))
+                          SingleExtensionUsagesValue.hashForEvaluation(
+                              GsonTypeAdapterUtil.createSingleExtensionUsagesValueHashGson(),
+                              usagesValue))
                       .setRecordedFileInputs(moduleExtensionResult.getRecordedFileInputs())
                       .setRecordedDirentsInputs(moduleExtensionResult.getRecordedDirentsInputs())
                       .setEnvVariables(extension.getEnvVars())
@@ -281,9 +281,8 @@ private SingleExtensionValue tryGettingValueFromLockFile(
       // Check extension data in lockfile is still valid, disregarding usage information that is not
       // relevant for the evaluation of the extension.
       if (!Arrays.equals(
-          ModuleExtensionUsage.hashForEvaluation(
-              GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
-              usagesValue.getExtensionUsages()),
+          SingleExtensionUsagesValue.hashForEvaluation(
+              GsonTypeAdapterUtil.createSingleExtensionUsagesValueHashGson(), usagesValue),
           lockedExtension.getUsagesDigest())) {
         diffRecorder.record("The usages of the extension '" + extensionId + "' have changed");
       }

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

@@ -22,6 +22,7 @@
 import com.google.devtools.build.skyframe.SkyFunction;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.util.Optional;
 import javax.annotation.Nullable;
 
 /**
@@ -49,17 +50,30 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
     }
 
     ModuleExtensionId id = (ModuleExtensionId) skyKey.argument();
+    // We never request an extension without usages in Skyframe.
+    return createValue(bazelDepGraphValue, id).get();
+  }
+
+  public static Optional<SingleExtensionUsagesValue> createValue(
+      BazelDepGraphValue depGraphValue, ModuleExtensionId id) {
     ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> usagesTable =
-        bazelDepGraphValue.getExtensionUsagesTable();
-    return SingleExtensionUsagesValue.create(
-        usagesTable.row(id),
-        bazelDepGraphValue.getExtensionUniqueNames().get(id),
-        // Filter abridged modules down to only those that actually used this extension.
-        bazelDepGraphValue.getAbridgedModules().stream()
-            .filter(module -> usagesTable.contains(id, module.getKey()))
-            .collect(toImmutableList()),
-        // TODO(wyv): Maybe cache these mappings?
-        usagesTable.row(id).keySet().stream()
-            .collect(toImmutableMap(key -> key, bazelDepGraphValue::getFullRepoMapping)));
+        depGraphValue.getExtensionUsagesTable();
+    if (!usagesTable.containsRow(id)) {
+      return Optional.empty();
+    }
+    return Optional.of(
+        SingleExtensionUsagesValue.builder()
+            .setExtensionUsages(usagesTable.row(id))
+            .setExtensionUniqueName(depGraphValue.getExtensionUniqueNames().get(id))
+            .setAbridgedModules(
+                // Filter abridged modules down to only those that actually used this extension.
+                depGraphValue.getAbridgedModules().stream()
+                    .filter(module -> usagesTable.contains(id, module.getKey()))
+                    .collect(toImmutableList()))
+            .setRepoMappings(
+                // TODO(wyv): Maybe cache these mappings?
+                usagesTable.row(id).keySet().stream()
+                    .collect(toImmutableMap(key -> key, depGraphValue::getFullRepoMapping)))
+            .build());
   }
 }

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

@@ -18,6 +18,8 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import com.google.common.hash.Hashing;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.serialization.VisibleForSerialization;
@@ -26,9 +28,17 @@
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import com.google.gson.Gson;
+import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 
-/** The result of {@link SingleExtensionUsagesFunction}. */
+/**
+ * The result of {@link SingleExtensionUsagesFunction}.
+ *
+ * <p>When adding or exposing new fields to extensions, make sure to update {@link
+ * #trimForEvaluation()} as well.
+ */
 @AutoValue
+@GenerateTypeAdapter
 public abstract class SingleExtensionUsagesValue implements SkyValue {
   /** All usages of this extension, by the key of the module where the usage occurs. */
   // Note: Equality of SingleExtensionUsagesValue does not check for equality of the order of the
@@ -46,19 +56,63 @@ public abstract class SingleExtensionUsagesValue implements SkyValue {
   /** The repo mappings to use for each module that used this extension. */
   public abstract ImmutableMap<ModuleKey, RepositoryMapping> getRepoMappings();
 
-  public static SingleExtensionUsagesValue create(
-      ImmutableMap<ModuleKey, ModuleExtensionUsage> extensionUsages,
-      String extensionUniqueName,
-      ImmutableList<AbridgedModule> abridgedModules,
-      ImmutableMap<ModuleKey, RepositoryMapping> repoMappings) {
-    return new AutoValue_SingleExtensionUsagesValue(
-        extensionUsages, extensionUniqueName, abridgedModules, repoMappings);
+  abstract Builder toBuilder();
+
+  public static Builder builder() {
+    return new AutoValue_SingleExtensionUsagesValue.Builder();
+  }
+
+  /**
+   * Turns the given usages value for a particular extension into a hash that can be compared for
+   * equality with another hash obtained in this way and compares equal only if the two values are
+   * equivalent for the purpose of evaluating the extension.
+   */
+  static byte[] hashForEvaluation(Gson gson, SingleExtensionUsagesValue usagesValue) {
+    return Hashing.sha256()
+        .hashUnencodedChars(gson.toJson(usagesValue.trimForEvaluation()))
+        .asBytes();
+  }
+
+  /**
+   * Returns a new value with all information removed that does not influence the evaluation of the
+   * extension.
+   */
+  SingleExtensionUsagesValue trimForEvaluation() {
+    return toBuilder()
+        .setExtensionUsages(
+            ImmutableMap.copyOf(
+                Maps.transformValues(
+                    getExtensionUsages(), ModuleExtensionUsage::trimForEvaluation)))
+        // The unique name of the extension is not accessible to the extension's implementation
+        // function.
+        // TODO: Reconsider this when resolving #19055.
+        .setExtensionUniqueName("")
+        // The usage of repo mappings by the extension's implementation function is tracked on the
+        // level of individual entries and all label attributes are provided as `Label`, which
+        // exclusively reference canonical repository names.
+        .setRepoMappings(ImmutableMap.of())
+        .build();
   }
 
   public static Key key(ModuleExtensionId id) {
     return Key.create(id);
   }
 
+  /** Builder for {@link SingleExtensionUsagesValue}. */
+  @AutoValue.Builder
+  abstract static class Builder {
+
+    public abstract Builder setExtensionUsages(ImmutableMap<ModuleKey, ModuleExtensionUsage> value);
+
+    public abstract Builder setExtensionUniqueName(String value);
+
+    public abstract Builder setAbridgedModules(ImmutableList<AbridgedModule> value);
+
+    public abstract Builder setRepoMappings(ImmutableMap<ModuleKey, RepositoryMapping> value);
+
+    public abstract SingleExtensionUsagesValue build();
+  }
+
   @AutoCodec
   static class Key extends AbstractSkyKey<ModuleExtensionId> {
     private static final SkyKeyInterner<Key> interner = SkyKey.newInterner();

src/test/py/bazel/bzlmod/bazel_lockfile_test.py

@@ -2326,5 +2326,6 @@ def testModuleExtensionRerunsOnVersionChange(self):
     stderr = ''.join(stderr)
     self.assertIn('I am running the extension: 4.5.6', stderr)
 
+
 if __name__ == '__main__':
   absltest.main()