Use repository cache for BCR files

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

@@ -77,9 +77,11 @@ java_library(
         "Registry.java",
         "RegistryFactory.java",
         "RegistryFactoryImpl.java",
+        "RegistryFileDownloadEvent.java",
     ],
     deps = [
         ":common",
+        "//src/main/java/com/google/devtools/build/lib/bazel/repository/cache",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/events",
@@ -98,11 +100,13 @@ java_library(
     srcs = ["BazelLockFileModule.java"],
     deps = [
         ":exception",
+        ":registry",
         ":resolution",
         ":resolution_impl",
         "//src/main/java/com/google/devtools/build/lib:runtime",
         "//src/main/java/com/google/devtools/build/lib/bazel/bzlmod:module_extension",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository:repository_options",
+        "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/util:abrupt_exit_exception",
         "//src/main/java/com/google/devtools/build/lib/vfs",
@@ -157,6 +161,7 @@ java_library(
         ":root_module_file_fixup",
         "//src/main/java/com/google/devtools/build/docgen/annot",
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_directories",
+        "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//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",
@@ -222,6 +227,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/analysis:blaze_version_info",
         "//src/main/java/com/google/devtools/build/lib/bazel:bazel_version",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository:repository_options",
+        "//src/main/java/com/google/devtools/build/lib/bazel/repository/cache",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/downloader",
         "//src/main/java/com/google/devtools/build/lib/bazel/repository/starlark",
         "//src/main/java/com/google/devtools/build/lib/cmdline",

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

@@ -33,6 +33,7 @@
 import com.google.common.collect.ImmutableTable;
 import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
+import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.RepositoryMapping;
@@ -47,8 +48,10 @@
 import com.google.devtools.build.skyframe.SkyFunctionException.Transience;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
+import java.net.URI;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nullable;
 

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

@@ -19,6 +19,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
@@ -65,6 +66,7 @@ public class BazelLockFileFunction implements SkyFunction {
           .setLocalOverrideHashes(ImmutableMap.of())
           .setModuleDepGraph(ImmutableMap.of())
           .setModuleExtensions(ImmutableMap.of())
+          .setFileHashes(ImmutableMap.of())
           .build();
 
   public BazelLockFileFunction(Path rootDirectory) {

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

@@ -16,13 +16,15 @@
 
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Comparator.comparing;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.eventbus.Subscribe;
 import com.google.common.flogger.GoogleLogger;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions;
 import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
+import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.runtime.BlazeModule;
 import com.google.devtools.build.lib.runtime.CommandEnvironment;
@@ -33,9 +35,14 @@
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.MemoizingEvaluator;
 import java.io.IOException;
+import java.net.URI;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Comparator;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import javax.annotation.Nullable;
 
 /**
@@ -47,11 +54,15 @@ public class BazelLockFileModule extends BlazeModule {
   private MemoizingEvaluator evaluator;
   private Path workspaceRoot;
   @Nullable private BazelModuleResolutionEvent moduleResolutionEvent;
+  private final List<RegistryFileDownloadEvent> downloadEvents = new ArrayList<>();
 
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
   @Override
   public void beforeCommand(CommandEnvironment env) {
+    this.downloadEvents.clear();
+    this.moduleResolutionEvent = null;
+
     evaluator = env.getSkyframeExecutor().getEvaluator();
     workspaceRoot = env.getWorkspace();
     RepositoryOptions options = env.getOptions().getOptions(RepositoryOptions.class);
@@ -88,6 +99,15 @@ public void afterCommand() throws AbruptExitException {
       newLockfile =
           moduleResolutionEvent.getResolutionOnlyLockfileValue().toBuilder()
               .setModuleExtensions(combineModuleExtensions(oldLockfile, extensionLockFileInfos))
+              .setFileHashes(
+                  // We ensure that the same file is not downloaded multiple times and can thus
+                  // collect without worrying about duplicates.
+                  downloadEvents.stream()
+                      .sorted(comparing(RegistryFileDownloadEvent::getUri))
+                      .collect(
+                          toImmutableMap(
+                              RegistryFileDownloadEvent::getUri,
+                              RegistryFileDownloadEvent::getChecksum)))
               .build();
     } catch (ExternalDepsException e) {
       logger.atSevere().withCause(e).log(
@@ -104,7 +124,6 @@ public void afterCommand() throws AbruptExitException {
     if (!newLockfile.equals(oldLockfile)) {
       updateLockfile(workspaceRoot, newLockfile);
     }
-    this.moduleResolutionEvent = null;
   }
 
   /**
@@ -248,4 +267,10 @@ public void bazelModuleResolved(BazelModuleResolutionEvent moduleResolutionEvent
     // sent after the command has modified the module file.
     this.moduleResolutionEvent = moduleResolutionEvent;
   }
+
+  @SuppressWarnings("unused")
+  @Subscribe
+  public void registryFileDownloaded(RegistryFileDownloadEvent downloadEvent) {
+    downloadEvents.add(downloadEvent);
+  }
 }

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

@@ -20,13 +20,16 @@
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.events.ExtendedEventHandler.Postable;
 import com.google.devtools.build.lib.skyframe.SkyFunctions;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.SkyValue;
 import com.ryanharter.auto.value.gson.GenerateTypeAdapter;
+import java.net.URI;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * The result of reading the lockfile. Contains the lockfile version, module hash, definitions of
@@ -37,14 +40,15 @@
 @GenerateTypeAdapter
 public abstract class BazelLockFileValue implements SkyValue, Postable {
 
-  public static final int LOCK_FILE_VERSION = 7;
+  public static final int LOCK_FILE_VERSION = 8;
 
   @SerializationConstant public static final SkyKey KEY = () -> SkyFunctions.BAZEL_LOCK_FILE;
 
   static Builder builder() {
     return new AutoValue_BazelLockFileValue.Builder()
         .setLockFileVersion(LOCK_FILE_VERSION)
-        .setModuleExtensions(ImmutableMap.of());
+        .setModuleExtensions(ImmutableMap.of())
+        .setFileHashes(ImmutableMap.of());
   }
 
   /** Current version of the lock file */
@@ -62,6 +66,9 @@ static Builder builder() {
   /** The post-selection dep graph retrieved from the lock file. */
   public abstract ImmutableMap<ModuleKey, Module> getModuleDepGraph();
 
+  /** Hashes of files retrieved from registries. */
+  public abstract ImmutableMap<String, Optional<Checksum>> getFileHashes();
+
   /** Mapping the extension id to the module extension data */
   public abstract ImmutableMap<
           ModuleExtensionId, ImmutableMap<ModuleExtensionEvalFactors, LockFileModuleExtension>>
@@ -74,6 +81,8 @@ static Builder builder() {
   public abstract static class Builder {
     public abstract Builder setLockFileVersion(int value);
 
+    public abstract Builder setFileHashes(ImmutableMap<String, Optional<Checksum>> value);
+
     public abstract Builder setModuleFileHash(String value);
 
     public abstract Builder setFlags(BzlmodFlagsAndEnvVars value);

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

@@ -26,6 +26,8 @@
 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.bazel.repository.cache.RepositoryCache;
+import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
 import com.google.devtools.build.lib.cmdline.RepositoryName;
@@ -475,10 +477,61 @@ public RepoRecordedInput.Dirents read(JsonReader jsonReader) throws IOException
             }
           };
 
+  // This can't reuse the existing type adapter factory for Optional as we need to explicitly
+  // serialize null values but don't want to rely on GSON's serializeNulls.
+  private static final class OptionalChecksumTypeAdapterFactory implements TypeAdapterFactory {
+
+    @Override
+    public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> typeToken) {
+      if (typeToken.getRawType() != Optional.class) {
+        return null;
+      }
+      Type type = typeToken.getType();
+      if (!(type instanceof ParameterizedType)) {
+        return null;
+      }
+      Type elementType = ((ParameterizedType) type).getActualTypeArguments()[0];
+      if (elementType != Checksum.class) {
+        return null;
+      }
+      @SuppressWarnings("unchecked")
+      TypeAdapter<T> typeAdapter = (TypeAdapter<T>) new OptionalChecksumTypeAdapter();
+      return typeAdapter;
+    }
+
+    private static class OptionalChecksumTypeAdapter extends TypeAdapter<Optional<Checksum>> {
+      // This value must not be a valid checksum string.
+      private static final String NOT_FOUND_MARKER = "not found";
+
+      @Override
+      public void write(JsonWriter jsonWriter, Optional<Checksum> checksum) throws IOException {
+        if (checksum.isPresent()) {
+          jsonWriter.value(checksum.get().toString());
+        } else {
+          jsonWriter.value(NOT_FOUND_MARKER);
+        }
+      }
+
+      @Override
+      public Optional<Checksum> read(JsonReader jsonReader) throws IOException {
+        String checksumString = jsonReader.nextString();
+        if (NOT_FOUND_MARKER.equals(checksumString)) {
+          return Optional.empty();
+        }
+        try {
+          return Optional.of(Checksum.fromString(RepositoryCache.KeyType.SHA256, checksumString));
+        } catch (Checksum.InvalidChecksumException e) {
+          throw new JsonParseException(String.format("Invalid checksum: %s", checksumString), e);
+        }
+      }
+    }
+  }
+
   public static Gson createLockFileGson(Path moduleFilePath, Path workspaceRoot) {
     return newGsonBuilder()
         .setPrettyPrinting()
         .registerTypeAdapterFactory(new LocationTypeAdapterFactory(moduleFilePath, workspaceRoot))
+        .registerTypeAdapterFactory(new OptionalChecksumTypeAdapterFactory())
         .create();
   }