Record repo mapping entries for labels in module extension tags

The mapping of an apparent name in a module extension tag's `attr.label` attribute to the corresponding canonical name has to be recorded in the lockfile so that instantiated repo rules don't reference the stale repos.

Fixes bazelbuild#25063

Closes bazelbuild#25067.

PiperOrigin-RevId: 720592796
Change-Id: Ia202ca4a8482a81da8085ee18ecaca5fe233bddb

Author: fmeum

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

@@ -36,7 +36,7 @@ public abstract class BazelLockFileValue implements SkyValue {
 
   // NOTE: See "HACK" note in BazelLockFileModule. While this hack exists, normal increments of the
   // lockfile version need to be done by 2 at a time (i.e. keep LOCK_FILE_VERSION an odd number).
-  public static final int LOCK_FILE_VERSION = 11;
+  public static final int LOCK_FILE_VERSION = 18;
 
   @SerializationConstant public static final SkyKey KEY = () -> SkyFunctions.BAZEL_LOCK_FILE;
 

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

@@ -272,7 +272,7 @@ private RunModuleExtensionResult runInternal(
     try (Mutability mu =
             Mutability.create("module extension", usagesValue.getExtensionUniqueName());
         ModuleExtensionContext moduleContext =
-            createContext(env, usagesValue, starlarkSemantics, extensionId)) {
+            createContext(env, usagesValue, starlarkSemantics, extensionId, repoMappingRecorder)) {
       StarlarkThread thread = new StarlarkThread(mu, starlarkSemantics);
       thread.setPrintHandler(Event.makeDebugPrintHandler(env.getListener()));
       threadContext.storeInThread(thread);
@@ -339,7 +339,8 @@ private ModuleExtensionContext createContext(
       Environment env,
       SingleExtensionUsagesValue usagesValue,
       StarlarkSemantics starlarkSemantics,
-      ModuleExtensionId extensionId)
+      ModuleExtensionId extensionId,
+      Label.RepoMappingRecorder repoMappingRecorder)
       throws ExternalDepsException {
     Path workingDirectory =
         directories
@@ -354,7 +355,8 @@ private ModuleExtensionContext createContext(
               abridgedModule,
               extension,
               usagesValue.getRepoMappings().get(moduleKey),
-              usagesValue.getExtensionUsages().get(moduleKey)));
+              usagesValue.getExtensionUsages().get(moduleKey),
+              repoMappingRecorder));
     }
     ModuleExtensionUsage rootUsage = usagesValue.getExtensionUsages().get(ModuleKey.ROOT);
     boolean rootModuleHasNonDevDependency =

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

@@ -19,6 +19,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.devtools.build.docgen.annot.DocCategory;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import com.google.devtools.build.lib.packages.LabelConverter;
@@ -107,12 +108,14 @@ public static StarlarkBazelModule create(
       AbridgedModule module,
       ModuleExtension extension,
       RepositoryMapping repoMapping,
-      @Nullable ModuleExtensionUsage usage)
+      @Nullable ModuleExtensionUsage usage,
+      Label.RepoMappingRecorder repoMappingRecorder)
       throws ExternalDepsException {
     LabelConverter labelConverter =
         new LabelConverter(
             PackageIdentifier.create(repoMapping.ownerRepo(), PathFragment.EMPTY_FRAGMENT),
-            repoMapping);
+            repoMapping,
+            repoMappingRecorder);
     ImmutableList<Tag> tags = usage == null ? ImmutableList.of() : usage.getTags();
     HashMap<String, ArrayList<TypeCheckedTag>> typeCheckedTags = new HashMap<>();
     for (String tagClassName : extension.getTagClasses().keySet()) {

src/main/java/com/google/devtools/build/lib/packages/LabelConverter.java

@@ -21,6 +21,7 @@
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
 import java.util.HashMap;
 import java.util.Map;
+import javax.annotation.Nullable;
 import net.starlark.java.eval.StarlarkThread;
 
 /**
@@ -42,16 +43,35 @@ public static LabelConverter forBzlEvaluatingThread(StarlarkThread thread) {
 
   private final Label.PackageContext packageContext;
   private final Map<String, Label> labelCache = new HashMap<>();
+  @Nullable private final Label.RepoMappingRecorder repoMappingRecorder;
 
-  public LabelConverter(Label.PackageContext packageContext) {
+  private LabelConverter(
+      Label.PackageContext packageContext,
+      @Nullable Label.RepoMappingRecorder repoMappingRecorder) {
     this.packageContext = packageContext;
+    this.repoMappingRecorder = repoMappingRecorder;
+  }
+
+  public LabelConverter(Label.PackageContext packageContext) {
+    this(packageContext, null);
   }
 
   /** Creates a label converter using the given base package and repo mapping. */
   public LabelConverter(PackageIdentifier base, RepositoryMapping repositoryMapping) {
     this(Label.PackageContext.of(base, repositoryMapping));
   }
 
+  /**
+   * Creates a label converter using the given base package and repo mapping, recording all repo
+   * mapping lookups in the given recorder.
+   */
+  public LabelConverter(
+      PackageIdentifier base,
+      RepositoryMapping repositoryMapping,
+      Label.RepoMappingRecorder repoMappingRecorder) {
+    this(Label.PackageContext.of(base, repositoryMapping), repoMappingRecorder);
+  }
+
   /** Returns the base package identifier that relative labels will be resolved against. */
   PackageIdentifier getBasePackage() {
     return packageContext.packageIdentifier();
@@ -65,7 +85,7 @@ public Label convert(String input) throws LabelSyntaxException {
     // label-strings across all their attribute values.
     Label converted = labelCache.get(input);
     if (converted == null) {
-      converted = Label.parseWithPackageContext(input, packageContext);
+      converted = Label.parseWithPackageContext(input, packageContext, repoMappingRecorder);
       labelCache.put(input, converted);
     }
     return converted;

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleExtensionResolutionTest.java

@@ -894,6 +894,24 @@ public void labels_passedOnToRepoRule() throws Exception {
     }
     assertThat(result.get(skyKey).getModule().getGlobal("data"))
         .isEqualTo("get up at 6am. go to bed at 11pm.");
+
+    SkyKey extensionSkyKey =
+        SingleExtensionValue.key(
+            ModuleExtensionId.create(
+                Label.parseCanonicalUnchecked("@@ext+//:defs.bzl"), "ext", Optional.empty()));
+    EvaluationResult<SingleExtensionValue> extensionResult =
+        evaluator.evaluate(ImmutableList.of(extensionSkyKey), evaluationContext);
+    if (extensionResult.hasError()) {
+      throw extensionResult.getError().getException();
+    }
+    assertThat(
+            extensionResult
+                .get(extensionSkyKey)
+                .lockFileInfo()
+                .get()
+                .moduleExtension()
+                .getRecordedRepoMappingEntries())
+        .containsCell(RepositoryName.create("foo+"), "bar", RepositoryName.create("bar+"));
   }
 
   @Test

src/test/java/com/google/devtools/build/lib/bazel/bzlmod/StarlarkBazelModuleTest.java

@@ -26,6 +26,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.cmdline.RepositoryName;
 import com.google.devtools.build.lib.packages.BuildType;
 import com.google.devtools.build.lib.packages.Type;
 import com.google.devtools.build.lib.util.FileTypeSet;
@@ -89,6 +90,7 @@ public void basic() throws Exception {
     Module module = buildModule("foo", "1.0").setKey(fooKey).addDep("bar", barKey).build();
     AbridgedModule abridgedModule = AbridgedModule.from(module);
 
+    Label.RepoMappingRecorder repoMappingRecorder = new Label.RepoMappingRecorder();
     StarlarkBazelModule moduleProxy =
         StarlarkBazelModule.create(
             abridgedModule,
@@ -97,7 +99,8 @@ public void basic() throws Exception {
                 ImmutableMap.of(
                     fooKey, fooKey.getCanonicalRepoNameWithoutVersionForTesting(),
                     barKey, barKey.getCanonicalRepoNameWithoutVersionForTesting())),
-            usage);
+            usage,
+            repoMappingRecorder);
 
     assertThat(moduleProxy.getName()).isEqualTo("foo");
     assertThat(moduleProxy.getVersion()).isEqualTo("1.0");
@@ -124,6 +127,9 @@ public void basic() throws Exception {
             StarlarkList.immutableOf(
                 Label.parseCanonical("@@foo~//:pom.xml"),
                 Label.parseCanonical("@@bar~//:pom.xml")));
+
+    assertThat(repoMappingRecorder.recordedEntries())
+        .containsCell(RepositoryName.create("foo+"), "bar", RepositoryName.create("bar+"));
   }
 
   @Test
@@ -145,7 +151,8 @@ public void unknownTagClass() throws Exception {
                     module.getRepoMappingWithBazelDepsOnly(
                         ImmutableMap.of(
                             fooKey, fooKey.getCanonicalRepoNameWithoutVersionForTesting())),
-                    usage));
+                    usage,
+                    new Label.RepoMappingRecorder()));
     assertThat(e).hasMessageThat().contains("does not have a tag class named blep");
   }
 }

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

@@ -1883,6 +1883,95 @@ def testExtensionRepoMappingChange_BzlInit(self):
     self.assertNotIn('ran the extension!', stderr)
     self.assertIn('STR=@@bar~//:lib_foo', stderr)
 
+  def testExtensionRepoMappingChange_tag(self):
+    # Regression test for #20721
+    self.main_registry.createCcModule('foo', '1.0')
+    self.main_registry.createCcModule('bar', '1.0')
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name="foo",version="1.0",repo_name="repo_name")',
+            'bazel_dep(name="bar",version="1.0")',
+            'ext = use_extension(":ext.bzl", "ext")',
+            'ext.tag(label = "@repo_name//:lib_foo")',
+            '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 = mctx.modules[0].tags.tag[0].label)',
+            'tag = tag_class(',
+            '  attrs = {',
+            '    "label": attr.label(),',
+            '  },',
+            ')',
+            'ext = module_extension(_ext_impl, tag_classes={"tag": tag})',
+        ],
+    )
+
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    self.assertIn('STR=@@foo+//: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: let repo_name point to bar 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="1.0")',
+            'bazel_dep(name="bar",version="1.0",repo_name="repo_name")',
+            'ext = use_extension(":ext.bzl", "ext")',
+            'ext.tag(label = "@repo_name//:lib_foo")',
+            'use_repo(ext, "repo")',
+        ],
+    )
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    stderr = '\n'.join(stderr)
+    self.assertIn('ran the extension!', stderr)
+    self.assertIn('STR=@@bar+//:lib_foo', stderr)
+
+    # Shutdown bazel to make sure we rely on the lockfile and not skyframe
+    self.RunBazel(['shutdown'])
+    # More spicy! change the repo_name of foo, but nothing else.
+    # The extension should NOT rerun, since it never used the @other_name repo
+    # mapping.
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name="foo",version="1.0",repo_name="other_name")',
+            'bazel_dep(name="bar",version="1.0",repo_name="repo_name")',
+            'ext = use_extension(":ext.bzl", "ext")',
+            'ext.tag(label = "@repo_name//:lib_foo")',
+            'use_repo(ext, "repo")',
+        ],
+    )
+    _, _, stderr = self.RunBazel(['build', ':lol'])
+    stderr = '\n'.join(stderr)
+    self.assertNotIn('ran the extension!', stderr)
+    self.assertIn('STR=@@bar+//:lib_foo', stderr)
+
   def testExtensionRepoMappingChange_loadsAndRepoRelativeLabels(self):
     # Regression test for #20721; same test as above, except that the call to
     # Label() in ext.bzl is now moved to @{foo,bar}//:defs.bzl and doesn't