From bbeb87efbccfc70f19c2a9f37211b77dd95d42b6 Mon Sep 17 00:00:00 2001
From: Xdng Yng <wyverald@gmail.com>
Date: Thu, 4 Jan 2024 15:27:19 -0800
Subject: [PATCH] Let module extensions track calls to `Label()`

... that use repo mapping. This is a rather obscure case of the lockfile being stale; if the `Label()` constructor is called in an extension impl function, and that call uses repo mapping of any form (i.e. the argument looks like `@foo//bar`), then we need to be ready to rerun the extension if `@foo` were to suddenly map to something else.

I also did a minor refactoring in `SingleExtensionEvalFunction` around the logic to decide whether the locked extension is up-to-date. Right now we perform a "diff" between the locked extension and what we expect to be up-to-date, and if a "diff" is found *and* `--lockfile_mode=error`, we basically perform a diff again. We also don't short circuit; that is, if the transitive bzl digest has changed, there's no point in seeing whether any files have changed, but we do right now.

A follow-up will be sent to fix the analogous bug for repo rules.

Fixes #20721.

Closes #20742.

PiperOrigin-RevId: 595818144
Change-Id: Id660b7a659a7f2e4dde19c16784c2ab18a9ceb69
---
 .../starlark/StarlarkRuleClassFunctions.java  |   5 +-
 .../devtools/build/lib/bazel/bzlmod/BUILD     |   1 +
 .../lib/bazel/bzlmod/BazelLockFileValue.java  |  33 ----
 .../lib/bazel/bzlmod/GsonTypeAdapterUtil.java |  84 ++++++++-
 .../bazel/bzlmod/LockFileModuleExtension.java |  11 +-
 .../bzlmod/SingleExtensionEvalFunction.java   | 166 ++++++++++++++----
 .../devtools/build/lib/cmdline/Label.java     |  35 ++++
 .../devtools/build/lib/cmdline/LabelTest.java |  37 ++++
 .../py/bazel/bzlmod/bazel_lockfile_test.py    |  80 +++++++++
 9 files changed, 378 insertions(+), 74 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
index 4d5d2ebfa57c8d..31488ce1f47a1c 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkRuleClassFunctions.java
@@ -1277,7 +1277,10 @@ public Label label(Object input, StarlarkThread thread) throws EvalException {
     // environment across .bzl files. Hence, we opt for stack inspection.
     BazelModuleContext moduleContext = BazelModuleContext.ofInnermostBzlOrFail(thread, "Label()");
     try {
-      return Label.parseWithPackageContext((String) input, moduleContext.packageContext());
+      return Label.parseWithPackageContext(
+          (String) input,
+          moduleContext.packageContext(),
+          thread.getThreadLocal(Label.RepoMappingRecorder.class));
     } catch (LabelSyntaxException e) {
       throw Starlark.errorf("invalid label in Label(): %s", e.getMessage());
     }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
index d97a0305b5ce2c..0083f8057b24ad 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD
@@ -220,6 +220,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe:client_environment_function",
         "//src/main/java/com/google/devtools/build/lib/skyframe:client_environment_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_value",
+        "//src/main/java/com/google/devtools/build/lib/skyframe:repository_mapping_value",
         "//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_cluster",
         "//src/main/java/com/google/devtools/build/lib/util",
         "//src/main/java/com/google/devtools/build/lib/util:os",
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java
index 21b02c479b76e7..8ac7fb38e77d64 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BazelLockFileValue.java
@@ -25,7 +25,6 @@
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
-import java.util.Arrays;
 import java.util.Map;
 
 /**
@@ -118,36 +117,4 @@ public ImmutableList<String> getModuleAndFlagsDiff(
     }
     return moduleDiff.build();
   }
-
-  /** Returns the differences between an extension and its locked data */
-  public ImmutableList<String> getModuleExtensionDiff(
-      ModuleExtensionId extensionId,
-      LockFileModuleExtension lockedExtension,
-      byte[] transitiveDigest,
-      boolean filesChanged,
-      ImmutableMap<String, String> envVariables,
-      ImmutableList<Map.Entry<ModuleKey, ModuleExtensionUsage>> extensionUsages,
-      ImmutableList<Map.Entry<ModuleKey, ModuleExtensionUsage>> lockedExtensionUsages) {
-
-    ImmutableList.Builder<String> extDiff = new ImmutableList.Builder<>();
-    if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {
-        extDiff.add(
-            "The implementation of the extension '"
-                + extensionId
-                + "' or one of its transitive .bzl files has changed");
-    }
-    if (filesChanged) {
-      extDiff.add("One or more files the extension '" + extensionId + "' is using have changed");
-    }
-    if (!extensionUsages.equals(lockedExtensionUsages)) {
-      extDiff.add("The usages of the extension '" + extensionId + "' have changed");
-    }
-    if (!envVariables.equals(lockedExtension.getEnvVariables())) {
-      extDiff.add(
-          "The environment variables the extension '"
-              + extensionId
-              + "' depends on (or their values) have changed");
-    }
-    return extDiff.build();
-  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java
index 05a54e2e053cc7..8ef5920da28bed 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/GsonTypeAdapterUtil.java
@@ -23,9 +23,12 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableTable;
+import com.google.common.collect.Table;
 import com.google.devtools.build.lib.bazel.bzlmod.Version.ParseException;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
@@ -92,7 +95,7 @@ public ModuleKey read(JsonReader jsonReader) throws IOException {
       };
 
   public static final TypeAdapter<Label> LABEL_TYPE_ADAPTER =
-      new TypeAdapter<Label>() {
+      new TypeAdapter<>() {
         @Override
         public void write(JsonWriter jsonWriter, Label label) throws IOException {
           jsonWriter.value(label.getUnambiguousCanonicalForm());
@@ -104,6 +107,19 @@ public Label read(JsonReader jsonReader) throws IOException {
         }
       };
 
+  public static final TypeAdapter<RepositoryName> REPOSITORY_NAME_TYPE_ADAPTER =
+      new TypeAdapter<>() {
+        @Override
+        public void write(JsonWriter jsonWriter, RepositoryName repoName) throws IOException {
+          jsonWriter.value(repoName.getName());
+        }
+
+        @Override
+        public RepositoryName read(JsonReader jsonReader) throws IOException {
+          return RepositoryName.createUnvalidated(jsonReader.nextString());
+        }
+      };
+
   public static final TypeAdapter<ModuleExtensionId> MODULE_EXTENSION_ID_TYPE_ADAPTER =
       new TypeAdapter<>() {
         @Override
@@ -283,6 +299,70 @@ public Optional<T> read(JsonReader jsonReader) throws IOException {
     }
   }
 
+  /**
+   * Converts Guava tables into a JSON array of 3-tuples (one per cell). Each 3-tuple is a JSON
+   * array itself (rowKey, columnKey, value). For example, a JSON snippet could be: {@code [
+   * ["row1", "col1", "value1"], ["row2", "col2", "value2"], ... ]}
+   */
+  public static final TypeAdapterFactory IMMUTABLE_TABLE =
+      new TypeAdapterFactory() {
+        @Nullable
+        @Override
+        @SuppressWarnings("unchecked")
+        public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> typeToken) {
+          if (typeToken.getRawType() != ImmutableTable.class) {
+            return null;
+          }
+          Type type = typeToken.getType();
+          if (!(type instanceof ParameterizedType)) {
+            return null;
+          }
+          Type[] typeArgs = ((ParameterizedType) typeToken.getType()).getActualTypeArguments();
+          if (typeArgs.length != 3) {
+            return null;
+          }
+          var rowTypeAdapter = (TypeAdapter<Object>) gson.getAdapter(TypeToken.get(typeArgs[0]));
+          var colTypeAdapter = (TypeAdapter<Object>) gson.getAdapter(TypeToken.get(typeArgs[1]));
+          var valTypeAdapter = (TypeAdapter<Object>) gson.getAdapter(TypeToken.get(typeArgs[2]));
+          if (rowTypeAdapter == null || colTypeAdapter == null || valTypeAdapter == null) {
+            return null;
+          }
+          return (TypeAdapter<T>)
+              new TypeAdapter<ImmutableTable<Object, Object, Object>>() {
+                @Override
+                public void write(JsonWriter jsonWriter, ImmutableTable<Object, Object, Object> t)
+                    throws IOException {
+                  jsonWriter.beginArray();
+                  for (Table.Cell<Object, Object, Object> cell : t.cellSet()) {
+                    jsonWriter.beginArray();
+                    rowTypeAdapter.write(jsonWriter, cell.getRowKey());
+                    colTypeAdapter.write(jsonWriter, cell.getColumnKey());
+                    valTypeAdapter.write(jsonWriter, cell.getValue());
+                    jsonWriter.endArray();
+                  }
+                  jsonWriter.endArray();
+                }
+
+                @Override
+                public ImmutableTable<Object, Object, Object> read(JsonReader jsonReader)
+                    throws IOException {
+                  var builder = ImmutableTable.builder();
+                  jsonReader.beginArray();
+                  while (jsonReader.peek() != JsonToken.END_ARRAY) {
+                    jsonReader.beginArray();
+                    builder.put(
+                        rowTypeAdapter.read(jsonReader),
+                        colTypeAdapter.read(jsonReader),
+                        valTypeAdapter.read(jsonReader));
+                    jsonReader.endArray();
+                  }
+                  jsonReader.endArray();
+                  return builder.buildOrThrow();
+                }
+              };
+        }
+      };
+
   /**
    * A variant of {@link Location} that converts the absolute path to the root module file to a
    * constant and back.
@@ -371,8 +451,10 @@ public static Gson createLockFileGson(Path moduleFilePath) {
         .registerTypeAdapterFactory(IMMUTABLE_BIMAP)
         .registerTypeAdapterFactory(IMMUTABLE_SET)
         .registerTypeAdapterFactory(OPTIONAL)
+        .registerTypeAdapterFactory(IMMUTABLE_TABLE)
         .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath))
         .registerTypeAdapter(Label.class, LABEL_TYPE_ADAPTER)
+        .registerTypeAdapter(RepositoryName.class, REPOSITORY_NAME_TYPE_ADAPTER)
         .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleKey.class, MODULE_KEY_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleExtensionId.class, MODULE_EXTENSION_ID_TYPE_ADAPTER)
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java
index 0dedeb023a8ab7..fea4715f7de5e8 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/LockFileModuleExtension.java
@@ -16,7 +16,9 @@
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableTable;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.Postable;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
 import java.util.Optional;
@@ -32,7 +34,8 @@ public abstract class LockFileModuleExtension implements Postable {
 
   public static Builder builder() {
     return new AutoValue_LockFileModuleExtension.Builder()
-        .setModuleExtensionMetadata(Optional.empty());
+        .setModuleExtensionMetadata(Optional.empty())
+        .setRecordedRepoMappingEntries(ImmutableTable.of());
   }
 
   @SuppressWarnings("mutable")
@@ -46,6 +49,9 @@ public static Builder builder() {
 
   public abstract Optional<ModuleExtensionMetadata> getModuleExtensionMetadata();
 
+  public abstract ImmutableTable<RepositoryName, String, RepositoryName>
+      getRecordedRepoMappingEntries();
+
   public abstract Builder toBuilder();
 
   /** Builder type for {@link LockFileModuleExtension}. */
@@ -62,6 +68,9 @@ public abstract static class Builder {
 
     public abstract Builder setModuleExtensionMetadata(Optional<ModuleExtensionMetadata> value);
 
+    public abstract Builder setRecordedRepoMappingEntries(
+        ImmutableTable<RepositoryName, String, RepositoryName> value);
+
     public abstract LockFileModuleExtension build();
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
index 4cd81e7eeed0ec..8838dff04fd571 100644
--- a/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/bazel/bzlmod/SingleExtensionEvalFunction.java
@@ -29,6 +29,7 @@
 import com.google.common.collect.ImmutableTable;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
+import com.google.common.collect.Table;
 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;
@@ -57,6 +58,7 @@
 import com.google.devtools.build.lib.skyframe.BzlLoadFunction.BzlLoadFailedException;
 import com.google.devtools.build.lib.skyframe.BzlLoadValue;
 import com.google.devtools.build.lib.skyframe.PrecomputedValue;
+import com.google.devtools.build.lib.skyframe.RepositoryMappingValue;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.vfs.Path;
@@ -158,14 +160,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
       }
       try {
         SingleExtensionEvalValue singleExtensionEvalValue =
-            tryGettingValueFromLockFile(
-                env,
-                extensionId,
-                extension.getEvalFactors(),
-                extension.getEnvVars(),
-                usagesValue,
-                extension.getBzlTransitiveDigest(),
-                lockfile);
+            tryGettingValueFromLockFile(env, extensionId, extension, usagesValue, lockfile);
         if (singleExtensionEvalValue != null) {
           return singleExtensionEvalValue;
         }
@@ -204,6 +199,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
                       .setEnvVariables(extension.getEnvVars())
                       .setGeneratedRepoSpecs(generatedRepoSpecs)
                       .setModuleExtensionMetadata(moduleExtensionMetadata)
+                      .setRecordedRepoMappingEntries(
+                          moduleExtensionResult.getRecordedRepoMappingEntries())
                       .build()));
     }
     return validateAndCreateSingleExtensionEvalValue(
@@ -221,10 +218,8 @@ public SkyValue compute(SkyKey skyKey, Environment env)
   private SingleExtensionEvalValue tryGettingValueFromLockFile(
       Environment env,
       ModuleExtensionId extensionId,
-      ModuleExtensionEvalFactors extensionFactors,
-      ImmutableMap<String, String> envVariables,
+      RunnableExtension extension,
       SingleExtensionUsagesValue usagesValue,
-      byte[] bzlTransitiveDigest,
       BazelLockFileValue lockfile)
       throws SingleExtensionEvalFunctionException,
           InterruptedException,
@@ -233,7 +228,7 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
 
     var lockedExtensionMap = lockfile.getModuleExtensions().get(extensionId);
     LockFileModuleExtension lockedExtension =
-        lockedExtensionMap == null ? null : lockedExtensionMap.get(extensionFactors);
+        lockedExtensionMap == null ? null : lockedExtensionMap.get(extension.getEvalFactors());
     if (lockedExtension == null) {
       if (lockfileMode.equals(LockfileMode.ERROR)) {
         throw new SingleExtensionEvalFunctionException(
@@ -241,7 +236,9 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
                 Code.BAD_MODULE,
                 "The module extension '%s'%s does not exist in the lockfile",
                 extensionId,
-                extensionFactors.isEmpty() ? "" : " for platform " + extensionFactors),
+                extension.getEvalFactors().isEmpty()
+                    ? ""
+                    : " for platform " + extension.getEvalFactors()),
             Transience.PERSISTENT);
       }
       return null;
@@ -258,43 +255,122 @@ private SingleExtensionEvalValue tryGettingValueFromLockFile(
       throw new SingleExtensionEvalFunctionException(e, Transience.PERSISTENT);
     }
 
-    boolean filesChanged = didFilesChange(env, lockedExtension.getAccumulatedFileDigests());
-    // Check extension data in lockfile is still valid, disregarding usage information that is not
-    // relevant for the evaluation of the extension.
-    var trimmedLockedUsages = ModuleExtensionUsage.trimForEvaluation(lockedExtensionUsages);
-    var trimmedUsages = ModuleExtensionUsage.trimForEvaluation(usagesValue.getExtensionUsages());
-    if (!filesChanged
-        && Arrays.equals(bzlTransitiveDigest, lockedExtension.getBzlTransitiveDigest())
-        && trimmedUsages.equals(trimmedLockedUsages)
-        && envVariables.equals(lockedExtension.getEnvVariables())) {
+    DiffRecorder diffRecorder =
+        new DiffRecorder(/* recordMessages= */ lockfileMode.equals(LockfileMode.ERROR));
+    try {
+      // Put faster diff detections earlier, so that we can short-circuit in UPDATE mode.
+      if (!Arrays.equals(
+          extension.getBzlTransitiveDigest(), lockedExtension.getBzlTransitiveDigest())) {
+        diffRecorder.record(
+            "The implementation of the extension '"
+                + extensionId
+                + "' or one of its transitive .bzl files has changed");
+      }
+      if (!extension.getEnvVars().equals(lockedExtension.getEnvVariables())) {
+        diffRecorder.record(
+            "The environment variables the extension '"
+                + extensionId
+                + "' depends on (or their values) have changed");
+      }
+      // 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))) {
+        diffRecorder.record("The usages of the extension '" + extensionId + "' have changed");
+      }
+      if (didRepoMappingsChange(env, lockedExtension.getRecordedRepoMappingEntries())) {
+        diffRecorder.record(
+            "The repo mappings of certain repos used by the extension '"
+                + extensionId
+                + "' have changed");
+      }
+      if (didFilesChange(env, lockedExtension.getAccumulatedFileDigests())) {
+        diffRecorder.record(
+            "One or more files the extension '" + extensionId + "' is using have changed");
+      }
+    } catch (DiffFoundEarlyExitException ignored) {
+      // ignored
+    }
+    if (!diffRecorder.anyDiffsDetected()) {
       return validateAndCreateSingleExtensionEvalValue(
           lockedExtension.getGeneratedRepoSpecs(),
           lockedExtension.getModuleExtensionMetadata(),
           extensionId,
           usagesValue,
           env);
-    } else if (lockfileMode.equals(LockfileMode.ERROR)) {
-      ImmutableList<String> extDiff =
-          lockfile.getModuleExtensionDiff(
-              extensionId,
-              lockedExtension,
-              bzlTransitiveDigest,
-              filesChanged,
-              envVariables,
-              trimmedUsages,
-              trimmedLockedUsages);
+    }
+    if (lockfileMode.equals(LockfileMode.ERROR)) {
       throw new SingleExtensionEvalFunctionException(
           ExternalDepsException.withMessage(
               Code.BAD_MODULE,
               "MODULE.bazel.lock is no longer up-to-date because: %s. "
                   + "Please run `bazel mod deps --lockfile_mode=update` to update your lockfile.",
-              String.join(", ", extDiff)),
+              diffRecorder.getRecordedDiffMessages()),
           Transience.PERSISTENT);
     }
     return null;
   }
 
-  private boolean didFilesChange(
+  private static final class DiffFoundEarlyExitException extends Exception {}
+
+  private static final class DiffRecorder {
+    private boolean diffDetected = false;
+    private final ImmutableList.Builder<String> diffMessages;
+
+    DiffRecorder(boolean recordMessages) {
+      diffMessages = recordMessages ? ImmutableList.builder() : null;
+    }
+
+    private void record(String message) throws DiffFoundEarlyExitException {
+      diffDetected = true;
+      if (diffMessages != null) {
+        diffMessages.add(message);
+      } else {
+        throw new DiffFoundEarlyExitException();
+      }
+    }
+
+    public boolean anyDiffsDetected() {
+      return diffDetected;
+    }
+
+    public String getRecordedDiffMessages() {
+      return String.join(",", diffMessages.build());
+    }
+  }
+
+  private static boolean didRepoMappingsChange(
+      Environment env, ImmutableTable<RepositoryName, String, RepositoryName> recordedRepoMappings)
+      throws InterruptedException, NeedsSkyframeRestartException {
+    SkyframeLookupResult result =
+        env.getValuesAndExceptions(
+            recordedRepoMappings.rowKeySet().stream()
+                .map(RepositoryMappingValue::key)
+                .collect(toImmutableSet()));
+    if (env.valuesMissing()) {
+      // This shouldn't really happen, since the RepositoryMappingValues of any recorded repos
+      // should have already been requested by the time we load the .bzl for the extension. And this
+      // method is only called if the transitive .bzl digest hasn't changed.
+      // However, we pretend it could happen anyway because we're good citizens.
+      throw new NeedsSkyframeRestartException();
+    }
+    for (Table.Cell<RepositoryName, String, RepositoryName> cell : recordedRepoMappings.cellSet()) {
+      RepositoryMappingValue repoMappingValue =
+          (RepositoryMappingValue) result.get(RepositoryMappingValue.key(cell.getRowKey()));
+      if (repoMappingValue == null) {
+        // Again, this shouldn't happen. But anyway.
+        throw new NeedsSkyframeRestartException();
+      }
+      if (!cell.getValue()
+          .equals(repoMappingValue.getRepositoryMapping().get(cell.getColumnKey()))) {
+        // Wee woo wee woo -- diff detected!
+        return true;
+      }
+    }
+    return false;
+  }
+
+  private static boolean didFilesChange(
       Environment env, ImmutableMap<Label, String> accumulatedFileDigests)
       throws InterruptedException, NeedsSkyframeRestartException {
     // Turn labels into FileValue keys & get those values
@@ -679,7 +755,10 @@ public RunModuleExtensionResult run(
         generatedRepoSpecs.put(name, repoSpec);
       }
       return RunModuleExtensionResult.create(
-          ImmutableMap.of(), generatedRepoSpecs.buildOrThrow(), Optional.empty());
+          ImmutableMap.of(),
+          generatedRepoSpecs.buildOrThrow(),
+          Optional.empty(),
+          ImmutableTable.of());
     }
   }
 
@@ -779,12 +858,16 @@ public RunModuleExtensionResult run(
               env.getListener());
       ModuleExtensionContext moduleContext;
       Optional<ModuleExtensionMetadata> moduleExtensionMetadata;
+      var repoMappingRecorder = new Label.RepoMappingRecorder();
       try (Mutability mu =
           Mutability.create("module extension", usagesValue.getExtensionUniqueName())) {
         StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics);
         thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
         moduleContext = createContext(env, usagesValue, starlarkSemantics, extensionId);
         threadContext.storeInThread(thread);
+        // This is used by the `Label()` constructor in Starlark, to record any attempts to resolve
+        // apparent repo names to canonical repo names. See #20721 for why this is necessary.
+        thread.setThreadLocal(Label.RepoMappingRecorder.class, repoMappingRecorder);
         try (SilentCloseable c =
             Profiler.instance()
                 .profile(
@@ -842,7 +925,8 @@ public RunModuleExtensionResult run(
       return RunModuleExtensionResult.create(
           moduleContext.getAccumulatedFileDigests(),
           threadContext.getGeneratedRepoSpecs(),
-          moduleExtensionMetadata);
+          moduleExtensionMetadata,
+          repoMappingRecorder.recordedEntries());
     }
 
     private ModuleExtensionContext createContext(
@@ -905,12 +989,18 @@ abstract static class RunModuleExtensionResult {
 
     abstract Optional<ModuleExtensionMetadata> getModuleExtensionMetadata();
 
+    abstract ImmutableTable<RepositoryName, String, RepositoryName> getRecordedRepoMappingEntries();
+
     static RunModuleExtensionResult create(
         ImmutableMap<Label, String> accumulatedFileDigests,
         ImmutableMap<String, RepoSpec> generatedRepoSpecs,
-        Optional<ModuleExtensionMetadata> moduleExtensionMetadata) {
+        Optional<ModuleExtensionMetadata> moduleExtensionMetadata,
+        ImmutableTable<RepositoryName, String, RepositoryName> recordedRepoMappingEntries) {
       return new AutoValue_SingleExtensionEvalFunction_RunModuleExtensionResult(
-          accumulatedFileDigests, generatedRepoSpecs, moduleExtensionMetadata);
+          accumulatedFileDigests,
+          generatedRepoSpecs,
+          moduleExtensionMetadata,
+          recordedRepoMappingEntries);
     }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/cmdline/Label.java b/src/main/java/com/google/devtools/build/lib/cmdline/Label.java
index 82e3ed16db6dfc..d0a033d4b23bdc 100644
--- a/src/main/java/com/google/devtools/build/lib/cmdline/Label.java
+++ b/src/main/java/com/google/devtools/build/lib/cmdline/Label.java
@@ -14,11 +14,15 @@
 package com.google.devtools.build.lib.cmdline;
 
 import static com.google.devtools.build.lib.cmdline.LabelParser.validateAndProcessTargetName;
+import static java.util.Comparator.naturalOrder;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ComparisonChain;
+import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableTable;
+import com.google.common.collect.Table;
 import com.google.common.util.concurrent.Striped;
 import com.google.devtools.build.docgen.annot.DocCategory;
 import com.google.devtools.build.lib.actions.CommandLineItem;
@@ -34,6 +38,7 @@
 import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.devtools.build.skyframe.SkyKey;
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import javax.annotation.Nullable;
@@ -193,7 +198,23 @@ public static Label parseWithRepoContext(String raw, RepoContext repoContext)
    */
   public static Label parseWithPackageContext(String raw, PackageContext packageContext)
       throws LabelSyntaxException {
+    return parseWithPackageContextInternal(Parts.parse(raw), packageContext);
+  }
+
+  public static Label parseWithPackageContext(
+      String raw, PackageContext packageContext, @Nullable RepoMappingRecorder repoMappingRecorder)
+      throws LabelSyntaxException {
     Parts parts = Parts.parse(raw);
+    Label parsed = parseWithPackageContextInternal(parts, packageContext);
+    if (repoMappingRecorder != null && parts.repo() != null && !parts.repoIsCanonical()) {
+      repoMappingRecorder.entries.put(
+          packageContext.currentRepo(), parts.repo(), parsed.getRepository());
+    }
+    return parsed;
+  }
+
+  private static Label parseWithPackageContextInternal(Parts parts, PackageContext packageContext)
+      throws LabelSyntaxException {
     parts.checkPkgDoesNotEndWithTripleDots();
     // pkg is either absolute or empty
     if (!parts.pkg().isEmpty()) {
@@ -205,6 +226,20 @@ public static Label parseWithPackageContext(String raw, PackageContext packageCo
     return createUnvalidated(PackageIdentifier.create(repoName, pkgFragment), parts.target());
   }
 
+  /** Records repo mapping entries used by {@link #parseWithPackageContext}. */
+  public static final class RepoMappingRecorder {
+    /** {@code <fromRepo, apparentRepoName, canonicalRepoName> } */
+    Table<RepositoryName, String, RepositoryName> entries = HashBasedTable.create();
+
+    public ImmutableTable<RepositoryName, String, RepositoryName> recordedEntries() {
+      return ImmutableTable.<RepositoryName, String, RepositoryName>builder()
+          .orderRowsBy(Comparator.comparing(RepositoryName::getName))
+          .orderColumnsBy(naturalOrder())
+          .putAll(entries)
+          .buildOrThrow();
+    }
+  }
+
   /**
    * Factory for Labels from separate components.
    *
diff --git a/src/test/java/com/google/devtools/build/lib/cmdline/LabelTest.java b/src/test/java/com/google/devtools/build/lib/cmdline/LabelTest.java
index bdfc93e0ddc456..a31d185cc8d8ee 100644
--- a/src/test/java/com/google/devtools/build/lib/cmdline/LabelTest.java
+++ b/src/test/java/com/google/devtools/build/lib/cmdline/LabelTest.java
@@ -17,9 +17,11 @@
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Tables;
 import com.google.common.testing.EqualsTester;
 import com.google.devtools.build.lib.cmdline.Label.PackageContext;
 import com.google.devtools.build.lib.cmdline.Label.RepoContext;
+import com.google.devtools.build.lib.cmdline.Label.RepoMappingRecorder;
 import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import java.util.regex.Pattern;
@@ -501,4 +503,39 @@ public void starlarkStr_ambiguousAndCanonicalLabelLiteral() throws Exception {
     assertThat(Starlark.str(Label.parseCanonical("//x"), semantics)).isEqualTo("//x:x");
     assertThat(Starlark.str(Label.parseCanonical("@x//y"), semantics)).isEqualTo("@@x//y:y");
   }
+
+  @Test
+  public void parseWithPackageContext_recordingRepoMapping() throws Exception {
+    RepositoryName foo = RepositoryName.createUnvalidated("foo");
+    RepositoryName bar = RepositoryName.createUnvalidated("bar");
+    RepositoryName quux = RepositoryName.createUnvalidated("quux");
+    PackageContext fooPackageContext =
+        PackageContext.of(
+            PackageIdentifier.create(foo, PathFragment.create("hah")),
+            RepositoryMapping.create(ImmutableMap.of("bar", quux, "quux", bar), foo));
+    PackageContext barPackageContext =
+        PackageContext.of(
+            PackageIdentifier.create(bar, PathFragment.EMPTY_FRAGMENT),
+            RepositoryMapping.create(ImmutableMap.of("foo", quux, "quux", foo), bar));
+
+    RepoMappingRecorder recorder = new RepoMappingRecorder();
+    // Not recorded: no repo part
+    var unused = Label.parseWithPackageContext("//foo/bar", fooPackageContext, recorder);
+    // Recorded: <foo, bar, quux>
+    unused = Label.parseWithPackageContext("@bar//foo/bar", fooPackageContext, recorder);
+    // Recorded: <bar, quux, foo>
+    unused = Label.parseWithPackageContext("@quux//foo/bar", barPackageContext, recorder);
+    // Recorded: <bar, foo, quux>
+    unused = Label.parseWithPackageContext("@foo//foo/bar", barPackageContext, recorder);
+    // Not recorded: canonical repo name
+    unused = Label.parseWithPackageContext("@@quux//foo/bar", fooPackageContext, recorder);
+
+    // Recorded entries are sorted by row and then column
+    assertThat(recorder.recordedEntries().cellSet())
+        .containsExactly(
+            Tables.immutableCell(bar, "foo", quux),
+            Tables.immutableCell(bar, "quux", foo),
+            Tables.immutableCell(foo, "bar", quux))
+        .inOrder();
+  }
 }
diff --git a/src/test/py/bazel/bzlmod/bazel_lockfile_test.py b/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
index 67cbacf1126ec1..d06a5451150daf 100644
--- a/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
+++ b/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
@@ -1518,6 +1518,86 @@ def testInnateModuleExtension(self):
           'general']['generatedRepoSpecs']['hello']['attributes']
       self.assertEqual(hello_attrs['value'], '@@//:hello.txt')
 
+  def testExtensionRepoMappingChange(self):
+    # Regression test for #20721
+    self.main_registry.createCcModule('foo', '1.0')
+    self.main_registry.createCcModule('foo', '2.0')
+    self.main_registry.createCcModule('bar', '1.0')
+    self.main_registry.createCcModule('bar', '2.0')
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name="foo",version="1.0")',
+            'bazel_dep(name="bar",version="1.0")',
+            'ext = use_extension(":ext.bzl", "ext")',
+            'use_repo(ext, "repo")',
+        ],
+    )
+    self.ScratchFile(
+        'BUILD.bazel',
+        [
+            'load("@repo//:defs.bzl", "STR")',
+            'print("STR="+STR)',
+            'filegroup(name="lol")',
+        ],
+    )
+    self.ScratchFile(
+        'ext.bzl',
+        [
+            'def _repo_impl(rctx):',
+            '  rctx.file("BUILD")',
+            '  rctx.file("defs.bzl", "STR = " + repr(str(rctx.attr.value)))',
+            'repo = repository_rule(_repo_impl,attrs={"value":attr.label()})',
+            'def _ext_impl(mctx):',
+            '  print("ran the extension!")',
+            '  repo(name = "repo", value = Label("@foo//:lib_foo"))',
+            'ext = module_extension(_ext_impl)',
+        ],
+    )
+
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    self.assertIn('STR=@@foo~1.0//:lib_foo', '\n'.join(stderr))
+
+    # Shutdown bazel to make sure we rely on the lockfile and not skyframe
+    self.RunBazel(['shutdown'])
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    self.assertNotIn('ran the extension!', '\n'.join(stderr))
+
+    # Shutdown bazel to make sure we rely on the lockfile and not skyframe
+    self.RunBazel(['shutdown'])
+    # Now, for something spicy: upgrade foo to 2.0 and change nothing else.
+    # The extension should rerun despite the lockfile being present, and no
+    # usages or .bzl files having changed.
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name="foo",version="2.0")',
+            'bazel_dep(name="bar",version="1.0")',
+            'ext = use_extension(":ext.bzl", "ext")',
+            'use_repo(ext, "repo")',
+        ],
+    )
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    self.assertIn('STR=@@foo~2.0//:lib_foo', '\n'.join(stderr))
+
+    # Shutdown bazel to make sure we rely on the lockfile and not skyframe
+    self.RunBazel(['shutdown'])
+    # More spicy! upgrade bar to 2.0 and change nothing else.
+    # The extension should NOT rerun, since it never used the @bar repo mapping.
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name="foo",version="2.0")',
+            'bazel_dep(name="bar",version="2.0")',
+            'ext = use_extension(":ext.bzl", "ext")',
+            'use_repo(ext, "repo")',
+        ],
+    )
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    stderr = '\n'.join(stderr)
+    self.assertNotIn('ran the extension!', stderr)
+    self.assertIn('STR=@@foo~2.0//:lib_foo', stderr)
+
 
 if __name__ == '__main__':
   absltest.main()