Add support for isolated extension usages to the lockfile

Previously, Bazel would crash if a module declares an isolated extension usage and the lockfile is used.

Closes #18991.

PiperOrigin-RevId: 549631385
Change-Id: Id8e706991dc5053b2873847a62f5e0b777347c69

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

@@ -41,7 +41,6 @@
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.util.Base64;
-import java.util.List;
 import java.util.Optional;
 import javax.annotation.Nullable;
 import net.starlark.java.syntax.Location;
@@ -82,52 +81,81 @@ public void write(JsonWriter jsonWriter, ModuleKey moduleKey) throws IOException
         @Override
         public ModuleKey read(JsonReader jsonReader) throws IOException {
           String jsonString = jsonReader.nextString();
-          if (jsonString.equals("<root>")) {
-            return ModuleKey.ROOT;
-          }
-          List<String> parts = Splitter.on('@').splitToList(jsonString);
-          if (parts.get(1).equals("_")) {
-            return ModuleKey.create(parts.get(0), Version.EMPTY);
-          }
-
-          Version version;
           try {
-            version = Version.parse(parts.get(1));
+            return ModuleKey.fromString(jsonString);
           } catch (ParseException e) {
             throw new JsonParseException(
                 String.format("Unable to parse ModuleKey %s version from the lockfile", jsonString),
                 e);
           }
-          return ModuleKey.create(parts.get(0), version);
         }
       };
 
-  // TODO(salmasamy) need to handle "isolated" in module extensions when it is stable
   public static final TypeAdapter<ModuleExtensionId> MODULE_EXTENSION_ID_TYPE_ADAPTER =
       new TypeAdapter<>() {
         @Override
         public void write(JsonWriter jsonWriter, ModuleExtensionId moduleExtId) throws IOException {
-          jsonWriter.value(moduleExtId.getBzlFileLabel() + "%" + moduleExtId.getExtensionName());
+          String isolationKeyPart = moduleExtId.getIsolationKey().map(key -> "%" + key).orElse("");
+          jsonWriter.value(
+              moduleExtId.getBzlFileLabel()
+                  + "%"
+                  + moduleExtId.getExtensionName()
+                  + isolationKeyPart);
         }
 
         @Override
         public ModuleExtensionId read(JsonReader jsonReader) throws IOException {
           String jsonString = jsonReader.nextString();
-          // [0] is labelString, [1] is extensionName
-          List<String> extIdParts = Splitter.on("%").splitToList(jsonString);
+          var extIdParts = Splitter.on('%').splitToList(jsonString);
+          Optional<ModuleExtensionId.IsolationKey> isolationKey;
+          if (extIdParts.size() > 2) {
+            try {
+              isolationKey =
+                  Optional.of(ModuleExtensionId.IsolationKey.fromString(extIdParts.get(2)));
+            } catch (ParseException e) {
+              throw new JsonParseException(
+                  String.format(
+                      "Unable to parse ModuleExtensionID isolation key: '%s' from the lockfile",
+                      extIdParts.get(2)),
+                  e);
+            }
+          } else {
+            isolationKey = Optional.empty();
+          }
           try {
             return ModuleExtensionId.create(
-                Label.parseCanonical(extIdParts.get(0)), extIdParts.get(1), Optional.empty());
+                Label.parseCanonical(extIdParts.get(0)), extIdParts.get(1), isolationKey);
           } catch (LabelSyntaxException e) {
             throw new JsonParseException(
                 String.format(
-                    "Unable to parse ModuleExtensionID bzl file label:  '%s' from the lockfile",
+                    "Unable to parse ModuleExtensionID bzl file label: '%s' from the lockfile",
                     extIdParts.get(0)),
                 e);
           }
         }
       };
 
+  public static final TypeAdapter<ModuleExtensionId.IsolationKey> ISOLATION_KEY_TYPE_ADAPTER =
+      new TypeAdapter<>() {
+        @Override
+        public void write(JsonWriter jsonWriter, ModuleExtensionId.IsolationKey isolationKey)
+            throws IOException {
+          jsonWriter.value(isolationKey.toString());
+        }
+
+        @Override
+        public ModuleExtensionId.IsolationKey read(JsonReader jsonReader) throws IOException {
+          String jsonString = jsonReader.nextString();
+          try {
+            return ModuleExtensionId.IsolationKey.fromString(jsonString);
+          } catch (ParseException e) {
+            throw new JsonParseException(
+                String.format("Unable to parse isolation key: '%s' from the lockfile", jsonString),
+                e);
+          }
+        }
+      };
+
   public static final TypeAdapter<byte[]> BYTE_ARRAY_TYPE_ADAPTER =
       new TypeAdapter<>() {
         @Override
@@ -283,6 +311,7 @@ public static Gson createLockFileGson(Path moduleFilePath) {
         .registerTypeAdapter(Version.class, VERSION_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleKey.class, MODULE_KEY_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleExtensionId.class, MODULE_EXTENSION_ID_TYPE_ADAPTER)
+        .registerTypeAdapter(ModuleExtensionId.IsolationKey.class, ISOLATION_KEY_TYPE_ADAPTER)
         .registerTypeAdapter(AttributeValues.class, new AttributeValuesAdapter())
         .registerTypeAdapter(byte[].class, BYTE_ARRAY_TYPE_ADAPTER)
         .create();

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

@@ -19,8 +19,10 @@
 import static java.util.Comparator.comparing;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.base.Splitter;
 import com.google.devtools.build.lib.cmdline.Label;
 import java.util.Comparator;
+import java.util.List;
 import java.util.Optional;
 
 /** A unique identifier for a {@link ModuleExtension}. */
@@ -49,6 +51,17 @@ abstract static class IsolationKey {
     public static IsolationKey create(ModuleKey module, String usageExportedName) {
       return new AutoValue_ModuleExtensionId_IsolationKey(module, usageExportedName);
     }
+
+    @Override
+    public final String toString() {
+      return getModule() + "~" + getUsageExportedName();
+    }
+
+    public static IsolationKey fromString(String s) throws Version.ParseException {
+      List<String> isolationKeyParts = Splitter.on("~").splitToList(s);
+      return ModuleExtensionId.IsolationKey.create(
+          ModuleKey.fromString(isolationKeyParts.get(0)), isolationKeyParts.get(1));
+    }
   }
 
   public abstract Label getBzlFileLabel();

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

@@ -16,9 +16,11 @@
 package com.google.devtools.build.lib.bazel.bzlmod;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
 import java.util.Comparator;
+import java.util.List;
 
 /** A module name, version pair that identifies a module in the external dependency graph. */
 @AutoValue
@@ -79,4 +81,17 @@ public RepositoryName getCanonicalRepoName() {
     return RepositoryName.createUnvalidated(
         String.format("%s~%s", getName(), getVersion().isEmpty() ? "override" : getVersion()));
   }
+
+  public static ModuleKey fromString(String s) throws Version.ParseException {
+    if (s.equals("<root>")) {
+      return ModuleKey.ROOT;
+    }
+    List<String> parts = Splitter.on('@').splitToList(s);
+    if (parts.get(1).equals("_")) {
+      return ModuleKey.create(parts.get(0), Version.EMPTY);
+    }
+
+    Version version = Version.parse(parts.get(1));
+    return ModuleKey.create(parts.get(0), version);
+  }
 }

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

@@ -299,6 +299,58 @@ def testModuleExtension(self):
     )
     self.assertNotIn('Hello from the other side!', ''.join(stderr))
 
+  def testIsolatedModuleExtension(self):
+    self.ScratchFile(
+        'MODULE.bazel',
+        [
+            (
+                'lockfile_ext = use_extension("extension.bzl", "lockfile_ext",'
+                ' isolate = True)'
+            ),
+            'lockfile_ext.dep(name = "bmbm", versions = ["v1", "v2"])',
+            'use_repo(lockfile_ext, "hello")',
+        ],
+    )
+    self.ScratchFile('BUILD.bazel')
+    self.ScratchFile(
+        'extension.bzl',
+        [
+            'def _repo_rule_impl(ctx):',
+            '    ctx.file("WORKSPACE")',
+            '    ctx.file("BUILD", "filegroup(name=\'lala\')")',
+            '',
+            'repo_rule = repository_rule(implementation=_repo_rule_impl)',
+            '',
+            'def _module_ext_impl(ctx):',
+            '    print("Hello from the other side!")',
+            '    repo_rule(name="hello")',
+            '    for mod in ctx.modules:',
+            '        for dep in mod.tags.dep:',
+            '            print("Name:", dep.name, ", Versions:", dep.versions)',
+            '',
+            '_dep = tag_class(attrs={"name": attr.string(), "versions":',
+            ' attr.string_list()})',
+            'lockfile_ext = module_extension(',
+            '    implementation=_module_ext_impl,',
+            '    tag_classes={"dep": _dep},',
+            '    environ=["GREEN_TREES", "NOT_SET"],',
+            ')',
+        ],
+    )
+
+    # Only set one env var, to make sure null variables don't crash
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'GREEN_TREES': 'In the city'}
+    )
+    self.assertIn('Hello from the other side!', ''.join(stderr))
+    self.assertIn('Name: bmbm , Versions: ["v1", "v2"]', ''.join(stderr))
+
+    self.RunBazel(['shutdown'])
+    _, _, stderr = self.RunBazel(
+        ['build', '@hello//:all'], env_add={'GREEN_TREES': 'In the city'}
+    )
+    self.assertNotIn('Hello from the other side!', ''.join(stderr))
+
   def testModuleExtensionsInDifferentBuilds(self):
     # Test that the module extension stays in the lockfile (as long as it's
     # used in the module) even if it is not in the current build