[6.4.0] Cherry pick platform dependent lockfile (#19498)

Store different extensions in the lockfile if they depend on os and arch

Includes:
- Old refactor commits for StarlarkRepositoryModule &
RepositoryModuleApi from
- Adding os and arch attributes to module extension 
- Updating the logic to save and fetch module extensions from the
lockfile depending on their os and arch dependency.

fixes #19154

PiperOrigin-RevId: 564707355
Change-Id: Ib48d7590e6d35397471a4ff46c58226eecf339c2

---------

Co-authored-by: Googler <wyv@google.com>
Co-authored-by: Yun Peng <pcloudy@google.com>

src/main/java/com/google/devtools/build/lib/bazel/bzlmod/BUILD

@@ -59,7 +59,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
         "//src/main/java/com/google/devtools/build/lib/packages",
-        "//src/main/java/net/starlark/java/annot",
+        "//src/main/java/com/google/devtools/build/lib/starlarkbuildapi/repository",
         "//src/main/java/net/starlark/java/eval",
         "//src/main/java/net/starlark/java/syntax",
         "//third_party:auto_value",
@@ -125,6 +125,7 @@ java_library(
         "LockFileModuleExtension.java",
         "Module.java",
         "ModuleBase.java",
+        "ModuleExtensionEvalFactors.java",
         "ModuleExtensionEvalStarlarkThreadContext.java",
         "ModuleExtensionResolutionEvent.java",
         "ModuleFileValue.java",
@@ -212,6 +213,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe:precomputed_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",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/skyframe",

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

@@ -32,8 +32,7 @@
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.gson.JsonSyntaxException;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashMap;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -45,7 +44,8 @@ public class BazelLockFileModule extends BlazeModule {
 
   private Path workspaceRoot;
   @Nullable private BazelModuleResolutionEvent moduleResolutionEvent;
-  private final List<ModuleExtensionResolutionEvent> extensionResolutionEvents = new ArrayList<>();
+  private final Map<ModuleExtensionId, ModuleExtensionResolutionEvent>
+      extensionResolutionEventsMap = new HashMap<>();
 
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
@@ -60,7 +60,7 @@ public void beforeCommand(CommandEnvironment env) {
 
   @Override
   public void afterCommand() throws AbruptExitException {
-    if (moduleResolutionEvent == null && extensionResolutionEvents.isEmpty()) {
+    if (moduleResolutionEvent == null && extensionResolutionEventsMap.isEmpty()) {
       return; // nothing changed, do nothing!
     }
 
@@ -91,44 +91,88 @@ public void afterCommand() throws AbruptExitException {
     // Write the new value to the file
     updateLockfile(lockfilePath, lockfile);
     this.moduleResolutionEvent = null;
-    this.extensionResolutionEvents.clear();
+    this.extensionResolutionEventsMap.clear();
   }
 
   /**
-   * Combines the old extensions stored in the lockfile -that are still used- with the new
+   * Combines the old extensions stored in the lockfile -if they are still valid- with the new
    * extensions from the events (if any)
    *
    * @param oldModuleExtensions Module extensions stored in the current lockfile
    */
-  private ImmutableMap<ModuleExtensionId, LockFileModuleExtension> combineModuleExtensions(
-      ImmutableMap<ModuleExtensionId, LockFileModuleExtension> oldModuleExtensions) {
-    ImmutableMap.Builder<ModuleExtensionId, LockFileModuleExtension> updatedExtensionMap =
-        ImmutableMap.builder();
-
-    // This event being null means that no changes occurred to the usages of the stored extensions,
-    // hence no changes to any module resulted in re-running resolution. So we can just add all the
-    // old stored extensions. Otherwise, check the usage of each one.
-    if (moduleResolutionEvent == null) {
-      updatedExtensionMap.putAll(oldModuleExtensions);
-    } else {
-      // Add the old extensions (stored in the lockfile) only if it still has a usage somewhere
-      for (Map.Entry<ModuleExtensionId, LockFileModuleExtension> extensionEntry :
-          oldModuleExtensions.entrySet()) {
-        if (moduleResolutionEvent.getExtensionUsagesById().containsRow(extensionEntry.getKey())) {
-          updatedExtensionMap.put(extensionEntry);
-        }
-      }
-    }
+  private ImmutableMap<
+          ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
+      combineModuleExtensions(
+          ImmutableMap<
+                  ModuleExtensionId,
+                  ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
+              oldModuleExtensions) {
+
+    Map<ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
+        updatedExtensionMap = new HashMap<>();
+
+    // Add old extensions if they are still valid
+    oldModuleExtensions.forEach(
+        (moduleExtensionId, innerMap) -> {
+          ModuleExtensionEvalFactors firstEntryKey = innerMap.keySet().iterator().next();
+          if (shouldKeepExtension(moduleExtensionId, firstEntryKey)) {
+            updatedExtensionMap.put(moduleExtensionId, innerMap);
+          }
+        });
 
     // Add the new resolved extensions
-    for (ModuleExtensionResolutionEvent extensionEvent : extensionResolutionEvents) {
-      updatedExtensionMap.put(extensionEvent.getExtensionId(), extensionEvent.getModuleExtension());
+    for (var event : extensionResolutionEventsMap.values()) {
+      var oldExtensionEntries = updatedExtensionMap.get(event.getExtensionId());
+      ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension> extensionEntries;
+      if (oldExtensionEntries != null) {
+        // extension exists, add the new entry to the existing map
+        extensionEntries =
+            new ImmutableMap.Builder<ModuleExtensionEvalFactors, LockFileModuleExtension>()
+                .putAll(oldExtensionEntries)
+                .put(event.getExtensionFactors(), event.getModuleExtension())
+                .buildKeepingLast();
+      } else {
+        // new extension
+        extensionEntries = ImmutableMap.of(event.getExtensionFactors(), event.getModuleExtension());
+      }
+      updatedExtensionMap.put(event.getExtensionId(), extensionEntries);
     }
+
     // The order in which extensions are added to extensionResolutionEvents depends on the order
     // in which their Skyframe evaluations finish, which is non-deterministic. We ensure a
     // deterministic lockfile by sorting.
     return ImmutableSortedMap.copyOf(
-        updatedExtensionMap.buildKeepingLast(), ModuleExtensionId.LEXICOGRAPHIC_COMPARATOR);
+        updatedExtensionMap, ModuleExtensionId.LEXICOGRAPHIC_COMPARATOR);
+  }
+
+  /**
+   * Decide whether to keep this extension or not depending on both: 1. If its dependency on os &
+   * arch didn't change 2. If it is still has a usage in the module
+   *
+   * @param lockedExtensionKey object holding the old extension id and state of os and arch
+   * @return True if this extension should still be in lockfile, false otherwise
+   */
+  private boolean shouldKeepExtension(
+      ModuleExtensionId extensionId, ModuleExtensionEvalFactors lockedExtensionKey) {
+
+    // If there is a new event for this extension, compare it with the existing ones
+    ModuleExtensionResolutionEvent extEvent = extensionResolutionEventsMap.get(extensionId);
+    if (extEvent != null) {
+      boolean dependencyOnOsChanged =
+          lockedExtensionKey.getOs().isEmpty() != extEvent.getExtensionFactors().getOs().isEmpty();
+      boolean dependencyOnArchChanged =
+          lockedExtensionKey.getArch().isEmpty()
+              != extEvent.getExtensionFactors().getArch().isEmpty();
+      if (dependencyOnOsChanged || dependencyOnArchChanged) {
+        return false;
+      }
+    }
+
+    // If moduleResolutionEvent is null, then no usage has changed. So we don't need this check
+    if (moduleResolutionEvent != null) {
+      return moduleResolutionEvent.getExtensionUsagesById().containsRow(extensionId);
+    }
+    return true;
   }
 
   /**
@@ -161,6 +205,8 @@ public void bazelModuleResolved(BazelModuleResolutionEvent moduleResolutionEvent
 
   @Subscribe
   public void moduleExtensionResolved(ModuleExtensionResolutionEvent extensionResolutionEvent) {
-    this.extensionResolutionEvents.add(extensionResolutionEvent);
+    this.extensionResolutionEventsMap.put(
+        extensionResolutionEvent.getExtensionId(), extensionResolutionEvent);
   }
+
 }

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 = 1;
+  public static final int LOCK_FILE_VERSION = 2;
 
   @SerializationConstant public static final SkyKey KEY = () -> SkyFunctions.BAZEL_LOCK_FILE;
 
@@ -63,7 +63,9 @@ static Builder builder() {
   public abstract ImmutableMap<ModuleKey, Module> getModuleDepGraph();
 
   /** Mapping the extension id to the module extension data */
-  public abstract ImmutableMap<ModuleExtensionId, LockFileModuleExtension> getModuleExtensions();
+  public abstract ImmutableMap<
+          ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
+      getModuleExtensions();
 
   public abstract Builder toBuilder();
 
@@ -81,7 +83,10 @@ public abstract static class Builder {
     public abstract Builder setModuleDepGraph(ImmutableMap<ModuleKey, Module> value);
 
     public abstract Builder setModuleExtensions(
-        ImmutableMap<ModuleExtensionId, LockFileModuleExtension> value);
+        ImmutableMap<
+                ModuleExtensionId,
+                ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
+            value);
 
     public abstract BazelLockFileValue build();
   }
@@ -117,12 +122,12 @@ public ImmutableList<String> getModuleAndFlagsDiff(
   /** 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,
       ImmutableMap<ModuleKey, ModuleExtensionUsage> extensionUsages,
       ImmutableMap<ModuleKey, ModuleExtensionUsage> lockedExtensionUsages) {
-    LockFileModuleExtension lockedExtension = getModuleExtensions().get(extensionId);
 
     ImmutableList.Builder<String> extDiff = new ImmutableList.Builder<>();
     if (!Arrays.equals(transitiveDigest, lockedExtension.getBzlTransitiveDigest())) {

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

@@ -148,6 +148,57 @@ public ModuleExtensionId read(JsonReader jsonReader) throws IOException {
         }
       };
 
+  public static final TypeAdapter<ModuleExtensionEvalFactors>
+      MODULE_EXTENSION_FACTORS_TYPE_ADAPTER =
+          new TypeAdapter<>() {
+
+            private static final String OS_KEY = "os:";
+            private static final String ARCH_KEY = "arch:";
+            // This is used when the module extension doesn't depend on os or arch, to indicate that
+            // its value is "general" and can be used with any platform
+            private static final String GENERAL_EXTENSION = "general";
+
+            @Override
+            public void write(JsonWriter jsonWriter, ModuleExtensionEvalFactors extFactors)
+                throws IOException {
+              if (extFactors.isEmpty()) {
+                jsonWriter.value(GENERAL_EXTENSION);
+              } else {
+                StringBuilder jsonBuilder = new StringBuilder();
+                if (!extFactors.getOs().isEmpty()) {
+                  jsonBuilder.append(OS_KEY).append(extFactors.getOs());
+                }
+                if (!extFactors.getArch().isEmpty()) {
+                  if (jsonBuilder.length() > 0) {
+                    jsonBuilder.append(",");
+                  }
+                  jsonBuilder.append(ARCH_KEY).append(extFactors.getArch());
+                }
+                jsonWriter.value(jsonBuilder.toString());
+              }
+            }
+
+            @Override
+            public ModuleExtensionEvalFactors read(JsonReader jsonReader) throws IOException {
+              String jsonString = jsonReader.nextString();
+              if (jsonString.equals(GENERAL_EXTENSION)) {
+                return ModuleExtensionEvalFactors.create("", "");
+              }
+
+              String os = "";
+              String arch = "";
+              var extParts = Splitter.on(',').splitToList(jsonString);
+              for (String part : extParts) {
+                if (part.startsWith(OS_KEY)) {
+                  os = part.substring(OS_KEY.length());
+                } else if (part.startsWith(ARCH_KEY)) {
+                  arch = part.substring(ARCH_KEY.length());
+                }
+              }
+              return ModuleExtensionEvalFactors.create(os, arch);
+            }
+          };
+
   public static final TypeAdapter<ModuleExtensionId.IsolationKey> ISOLATION_KEY_TYPE_ADAPTER =
       new TypeAdapter<>() {
         @Override
@@ -325,6 +376,8 @@ 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(
+            ModuleExtensionEvalFactors.class, MODULE_EXTENSION_FACTORS_TYPE_ADAPTER)
         .registerTypeAdapter(ModuleExtensionId.IsolationKey.class, ISOLATION_KEY_TYPE_ADAPTER)
         .registerTypeAdapter(AttributeValues.class, new AttributeValuesAdapter())
         .registerTypeAdapter(byte[].class, BYTE_ARRAY_TYPE_ADAPTER)

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

@@ -32,9 +32,6 @@ public abstract class LockFileModuleExtension implements Postable {
 
   public static Builder builder() {
     return new AutoValue_LockFileModuleExtension.Builder()
-        // TODO(salmasamy) can be removed when updating lockfile version
-        .setEnvVariables(ImmutableMap.of())
-        .setAccumulatedFileDigests(ImmutableMap.of())
         .setModuleExtensionMetadata(Optional.empty());
   }