Store usages hash in module extension lockfile entries

Work towards bazelbuild#20369

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

@@ -211,16 +211,6 @@ static BzlmodFlagsAndEnvVars getFlagsAndEnvVars(Environment env) throws Interrup
  compatabilityMode);
  }
 
- /**
- * For each extension usage, we resolve (i.e. canonicalize) its bzl file label. Then we can group
- * all usages by the label + name (the ModuleExtensionId).
- */
- public static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
- getExtensionUsagesById(ImmutableMap<ModuleKey, Module> depGraph)
- throws ExternalDepsException {
- return getExtensionUsagesById(depGraph, computeCanonicalRepoNameLookup(depGraph).inverse());
- }
-
  private static ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage>
  getExtensionUsagesById(
  ImmutableMap<ModuleKey, Module> depGraph,

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

@@ -19,7 +19,6 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
-import com.google.common.collect.ImmutableTable;
 import com.google.common.eventbus.Subscribe;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
@@ -33,6 +32,7 @@
 import com.google.devtools.build.lib.vfs.Root;
 import com.google.devtools.build.lib.vfs.RootedPath;
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
@@ -107,15 +107,19 @@ public void afterCommand() throws AbruptExitException {
  updatedExtensionMap = new HashMap<>();
 
  // Keep old extensions if they are still valid.
- ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages =
- BazelDepGraphFunction.getExtensionUsagesById(oldLockfile.getModuleDepGraph());
  for (var entry : oldLockfile.getModuleExtensions().entrySet()) {
  var moduleExtensionId = entry.getKey();
- var factorToLockedExtension = entry.getValue();
- ModuleExtensionEvalFactors firstEntryFactors =
- factorToLockedExtension.keySet().iterator().next();
- if (shouldKeepExtension(moduleExtensionId, firstEntryFactors, oldExtensionUsages)) {
- updatedExtensionMap.put(moduleExtensionId, factorToLockedExtension);
+ ImmutableMap.Builder<ModuleExtensionEvalFactors, LockFileModuleExtension>
+ factorsToKeepBuilder = new ImmutableMap.Builder<>();
+ for (var factorAndExtension : entry.getValue().entrySet()) {
+ if (shouldKeepExtension(
+ moduleExtensionId, factorAndExtension.getKey(), factorAndExtension.getValue())) {
+ factorsToKeepBuilder.put(factorAndExtension.getKey(), factorAndExtension.getValue());
+ }
+ }
+ var factorsToKeep = factorsToKeepBuilder.build();
+ if (!factorsToKeep.isEmpty()) {
+ updatedExtensionMap.put(moduleExtensionId, factorsToKeep);
  }
  }
 
@@ -164,7 +168,7 @@ public void afterCommand() throws AbruptExitException {
  private boolean shouldKeepExtension(
  ModuleExtensionId extensionId,
  ModuleExtensionEvalFactors lockedExtensionKey,
- ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> oldExtensionUsages) {
+ LockFileModuleExtension lockedExtension) {
 
  // If there is a new event for this extension, compare it with the existing ones
  ModuleExtensionResolutionEvent extEvent = extensionResolutionEventsMap.get(extensionId);
@@ -185,9 +189,11 @@ 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.
- return ModuleExtensionUsage.trimForEvaluation(
- moduleResolutionEvent.getExtensionUsagesById().row(extensionId))
- .equals(ModuleExtensionUsage.trimForEvaluation(oldExtensionUsages.row(extensionId)));
+ return Arrays.equals(
+ ModuleExtensionUsage.hashForEvaluation(
+ GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
+ moduleResolutionEvent.getExtensionUsagesById().row(extensionId)),
+ lockedExtension.getUsagesDigest());
  }
 
  /**

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

@@ -37,7 +37,7 @@
 @GenerateTypeAdapter
 public abstract class BazelLockFileValue implements SkyValue, Postable {
 
- public static final int LOCK_FILE_VERSION = 6;
+ public static final int LOCK_FILE_VERSION = 7;
 
  @SerializationConstant public static final SkyKey KEY = () -> SkyFunctions.BAZEL_LOCK_FILE;
 

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

@@ -476,8 +476,18 @@ public RepoRecordedInput.Dirents read(JsonReader jsonReader) throws IOException
  };
 
  public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
- return new GsonBuilder()
+ return newGsonBuilder()
  .setPrettyPrinting()
+ .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath, workspaceRoot))
+ .create();
+ }
+
+ public static Gson createModuleExtensionUsagesHashGson() {
+ return newGsonBuilder().create();
+ }
+
+ private static GsonBuilder newGsonBuilder() {
+ return new GsonBuilder()
  .disableHtmlEscaping()
  .enableComplexMapKeySerialization()
  .registerTypeAdapterFactory(GenerateTypeAdapter.FACTORY)
@@ -488,7 +498,6 @@ public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
  .registerTypeAdapterFactory(IMMUTABLE_SET)
  .registerTypeAdapterFactory(OPTIONAL)
  .registerTypeAdapterFactory(IMMUTABLE_TABLE)
- .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath, workspaceRoot))
  .registerTypeAdapter(Label.class, LABEL_TYPE_ADAPTER)
  .registerTypeAdapter(RepositoryName.class, REPOSITORY_NAME_TYPE_ADAPTER)
  .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
@@ -501,8 +510,7 @@ public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
  .registerTypeAdapter(byte[].class, BYTE_ARRAY_TYPE_ADAPTER)
  .registerTypeAdapter(RepoRecordedInput.File.class, REPO_RECORDED_INPUT_FILE_TYPE_ADAPTER)
  .registerTypeAdapter(
- RepoRecordedInput.Dirents.class, REPO_RECORDED_INPUT_DIRENTS_TYPE_ADAPTER)
- .create();
+ RepoRecordedInput.Dirents.class, REPO_RECORDED_INPUT_DIRENTS_TYPE_ADAPTER);
  }
 
  private GsonTypeAdapterUtil() {}

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

@@ -41,6 +41,9 @@ public static Builder builder() {
  @SuppressWarnings("mutable")
  public abstract byte[] getBzlTransitiveDigest();
 
+ @SuppressWarnings("mutable")
+ public abstract byte[] getUsagesDigest();
+
  public abstract ImmutableMap<RepoRecordedInput.File, String> getRecordedFileInputs();
 
  public abstract ImmutableMap<RepoRecordedInput.Dirents, String> getRecordedDirentsInputs();
@@ -67,6 +70,8 @@ public abstract static class Builder {
 
  public abstract Builder setBzlTransitiveDigest(byte[] digest);
 
+ public abstract Builder setUsagesDigest(byte[] digest);
+
  public abstract Builder setRecordedFileInputs(
  ImmutableMap<RepoRecordedInput.File, String> value);
 

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

@@ -22,9 +22,10 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Maps;
+import com.google.common.hash.Hashing;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import com.google.gson.Gson;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
-import java.util.Map;
 import java.util.Optional;
 import net.starlark.java.syntax.Location;
 
@@ -94,16 +95,15 @@ public static Builder builder() {
  }
 
  /**
- * Turns the given collection of usages for a particular extension into an object that can be
- * compared for equality with another object obtained in this way and compares equal only if the
- * two original collections of usages are equivalent for the purpose of evaluating the extension.
+ * 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 ImmutableList<Map.Entry<ModuleKey, ModuleExtensionUsage>> trimForEvaluation(
- ImmutableMap<ModuleKey, ModuleExtensionUsage> usages) {
- // ImmutableMap#equals doesn't compare the order of entries, but it matters for the evaluation
- // of the extension.
- return ImmutableList.copyOf(
- Maps.transformValues(usages, ModuleExtensionUsage::trimForEvaluation).entrySet());
+ static byte[] hashForEvaluation(Gson gson, ImmutableMap<ModuleKey, ModuleExtensionUsage> usages) {
+ ImmutableMap<ModuleKey, ModuleExtensionUsage> trimmedUsages =
+ ImmutableMap.copyOf(Maps.transformValues(usages, ModuleExtensionUsage::trimForEvaluation));
+ String usagesJson = gson.toJson(trimmedUsages);
+ return Hashing.sha256().hashUnencodedChars(usagesJson).asBytes();
  }
 
  /**

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

@@ -163,7 +163,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
  try {
  SingleExtensionEvalValue singleExtensionEvalValue =
  tryGettingValueFromLockFile(
- env, extensionId, extension, usagesValue, lockfile, lockedExtension);
+ env, extensionId, extension, usagesValue, lockedExtension);
  if (singleExtensionEvalValue != null) {
  return singleExtensionEvalValue;
  }
@@ -217,6 +217,10 @@ public SkyValue compute(SkyKey skyKey, Environment env)
  extension.getEvalFactors(),
  LockFileModuleExtension.builder()
  .setBzlTransitiveDigest(extension.getBzlTransitiveDigest())
+ .setUsagesDigest(
+ ModuleExtensionUsage.hashForEvaluation(
+ GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
+ usagesValue.getExtensionUsages()))
  .setRecordedFileInputs(moduleExtensionResult.getRecordedFileInputs())
  .setRecordedDirentsInputs(moduleExtensionResult.getRecordedDirentsInputs())
  .setEnvVariables(extension.getEnvVars())
@@ -243,24 +247,11 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
  ModuleExtensionId extensionId,
  RunnableExtension extension,
  SingleExtensionUsagesValue usagesValue,
- BazelLockFileValue lockfile,
  LockFileModuleExtension lockedExtension)
  throws SingleExtensionEvalFunctionException,
  InterruptedException,
  NeedsSkyframeRestartException {
  LockfileMode lockfileMode = BazelLockFileFunction.LOCKFILE_MODE.get(env);
-
- ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages;
- try {
- // TODO(salmasamy) might be nicer to precompute this table when we construct
- // BazelLockFileValue, without adding it to the json file
- ImmutableTable<ModuleExtensionId, ModuleKey, ModuleExtensionUsage> extensionUsagesById =
- BazelDepGraphFunction.getExtensionUsagesById(lockfile.getModuleDepGraph());
- lockedExtensionUsages = extensionUsagesById.row(extensionId);
- } catch (ExternalDepsException e) {
- throw new SingleExtensionEvalFunctionException(e, Transience.PERSISTENT);
- }
-
  DiffRecorder diffRecorder =
  new DiffRecorder(/* recordMessages= */ lockfileMode.equals(LockfileMode.ERROR));
  try {
@@ -280,8 +271,11 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
  }
  // Check extension data in lockfile is still valid, disregarding usage information that is not
  // relevant for the evaluation of the extension.
- if (!ModuleExtensionUsage.trimForEvaluation(usagesValue.getExtensionUsages())
- .equals(ModuleExtensionUsage.trimForEvaluation(lockedExtensionUsages))) {
+ if (!Arrays.equals(
+ ModuleExtensionUsage.hashForEvaluation(
+ GsonTypeAdapterUtil.createModuleExtensionUsagesHashGson(),
+ usagesValue.getExtensionUsages()),
+ lockedExtension.getUsagesDigest())) {
  diffRecorder.record("The usages of the extension '" + extensionId + "' have changed");
  }
  if (didRepoMappingsChange(env, lockedExtension.getRecordedRepoMappingEntries())) {