From 44779b80d185e05b7ab0fd037360beca00547f3e Mon Sep 17 00:00:00 2001
From: Xdng Yng <wyverald@gmail.com>
Date: Tue, 23 Jan 2024 11:35:50 -0800
Subject: [PATCH] Disregard WORKSPACE while verifying lockfile repo mapping
 entries in extension eval

See code comment and linked issue for more information.

Fixes #20942.

Closes #20982.

PiperOrigin-RevId: 600856392
Change-Id: I5b8a0ed3a38e37ab51ffb49b19a59f2e161b9a33
---
 .../bzlmod/SingleExtensionEvalFunction.java   | 17 ++++++-
 .../py/bazel/bzlmod/bazel_lockfile_test.py    | 44 +++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

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 753a8cc44a2dde..fcdbcab5be11fd 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
@@ -342,10 +342,19 @@ public String getRecordedDiffMessages() {
   private static boolean didRepoMappingsChange(
       Environment env, ImmutableTable<RepositoryName, String, RepositoryName> recordedRepoMappings)
       throws InterruptedException, NeedsSkyframeRestartException {
+    // Request repo mappings for any 'source repos' in the recorded mapping entries.
+    // Note specially that the main repo needs to be treated differently: if any .bzl file from the
+    // main repo was used for module extension eval, it _has_ to be before WORKSPACE is evaluated
+    // (see relevant code in BzlLoadFunction#getRepositoryMapping), so we only request the main repo
+    // mapping _without_ WORKSPACE repos. See #20942 for more information.
     SkyframeLookupResult result =
         env.getValuesAndExceptions(
             recordedRepoMappings.rowKeySet().stream()
-                .map(RepositoryMappingValue::key)
+                .map(
+                    repoName ->
+                        repoName.isMain()
+                            ? RepositoryMappingValue.KEY_FOR_ROOT_MODULE_WITHOUT_WORKSPACE_REPOS
+                            : RepositoryMappingValue.key(repoName))
                 .collect(toImmutableSet()));
     if (env.valuesMissing()) {
       // This likely means that one of the 'source repos' in the recorded mapping entries is no
@@ -354,7 +363,11 @@ private static boolean didRepoMappingsChange(
     }
     for (Table.Cell<RepositoryName, String, RepositoryName> cell : recordedRepoMappings.cellSet()) {
       RepositoryMappingValue repoMappingValue =
-          (RepositoryMappingValue) result.get(RepositoryMappingValue.key(cell.getRowKey()));
+          (RepositoryMappingValue)
+              result.get(
+                  cell.getRowKey().isMain()
+                      ? RepositoryMappingValue.KEY_FOR_ROOT_MODULE_WITHOUT_WORKSPACE_REPOS
+                      : RepositoryMappingValue.key(cell.getRowKey()));
       if (repoMappingValue == null) {
         throw new NeedsSkyframeRestartException();
       }
diff --git a/src/test/py/bazel/bzlmod/bazel_lockfile_test.py b/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
index f09ca3b38a362b..bbdaaa1e7cb4f2 100644
--- a/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
+++ b/src/test/py/bazel/bzlmod/bazel_lockfile_test.py
@@ -1863,6 +1863,50 @@ def testExtensionRepoMappingChange_sourceRepoNoLongerExistent(self):
     self.assertIn('ran the extension!', '\n'.join(stderr))
     self.assertIn('STR=@@quux~1.0//:quux.h', '\n'.join(stderr))
 
+  def testExtensionRepoMappingChange_mainRepoEvalCycleWithWorkspace(self):
+    # Regression test for #20942
+    self.main_registry.createCcModule('foo', '1.0')
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            'bazel_dep(name="foo",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)',
+        ],
+    )
+    # any `load` in WORKSPACE should trigger the bug
+    self.ScratchFile('WORKSPACE.bzlmod', ['load("@repo//:defs.bzl","STR")'])
+
+    _, _, stderr = self.RunBazel(['build', '--enable_workspace', ':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'])
+    # Build again. This should _NOT_ trigger a failure!
+    _, _, stderr = self.RunBazel(['build', '--enable_workspace', ':lol'])
+    self.assertNotIn('ran the extension!', '\n'.join(stderr))
+
 
 if __name__ == '__main__':
   absltest.main()