Allow the use of the gRPC downloader for Bzlmod

src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java

@@ -270,10 +270,7 @@ public void initializeRuleClasses(ConfiguredRuleClassProvider.Builder builder) {
   public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
     DownloadManager downloadManager =
         new DownloadManager(
-            repositoryCache.getDownloadCache(),
-            env.getDownloaderDelegate(),
-            env.getHttpDownloader(),
-            env.getReporter());
+            repositoryCache.getDownloadCache(), env.getDownloaderDelegate(), env.getReporter());
     this.repositoryFetchFunction.setDownloadManager(downloadManager);
     this.moduleFileFunction.setDownloadManager(downloadManager);
     this.repoSpecFunction.setDownloadManager(downloadManager);

src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DelegatingDownloader.java

@@ -16,8 +16,8 @@
 
 import com.google.auth.Credentials;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
-import com.google.devtools.build.lib.vfs.Path;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.net.URL;
 import java.util.List;
 import java.util.Map;
@@ -51,7 +51,8 @@ public void download(
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,
-      Path destination,
+      OutputStream out,
+      @Nullable String destinationPath,
       ExtendedEventHandler eventHandler,
       Map<String, String> clientEnv,
       Optional<String> type,
@@ -67,7 +68,8 @@ public void download(
         credentials,
         checksum,
         canonicalId,
-        destination,
+        out,
+        destinationPath,
         eventHandler,
         clientEnv,
         type,

src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java

@@ -29,7 +29,6 @@
 import com.google.devtools.build.lib.bazel.repository.cache.DownloadCache;
 import com.google.devtools.build.lib.bazel.repository.cache.DownloadCache.KeyType;
 import com.google.devtools.build.lib.bazel.repository.cache.DownloadCacheHitEvent;
-import com.google.devtools.build.lib.bazel.repository.cache.RepositoryCache;
 import com.google.devtools.build.lib.bazel.repository.downloader.UrlRewriter.RewrittenURL;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
@@ -38,6 +37,7 @@
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.net.SocketException;
@@ -66,7 +66,6 @@ public class DownloadManager {
   private ImmutableList<Path> distdir = ImmutableList.of();
   private UrlRewriter rewriter;
   private final Downloader downloader;
-  private final HttpDownloader bzlmodHttpDownloader;
   private final ExtendedEventHandler eventHandler;
   private boolean disableDownload = false;
   private int retries = 0;
@@ -83,17 +82,11 @@ public interface CredentialFactory {
    *
    * @param downloader The (delegating) downloader to use to download files. Is either a
    *     HttpDownloader, or a GrpcRemoteDownloader.
-   * @param bzlmodHttpDownloader The downloader to use for downloading files from the bzlmod
-   *     registry.
    */
   public DownloadManager(
-      DownloadCache downloadCache,
-      Downloader downloader,
-      HttpDownloader bzlmodHttpDownloader,
-      ExtendedEventHandler eventHandler) {
+      DownloadCache downloadCache, Downloader downloader, ExtendedEventHandler eventHandler) {
     this.downloadCache = downloadCache;
     this.downloader = downloader;
-    this.bzlmodHttpDownloader = bzlmodHttpDownloader;
     this.eventHandler = eventHandler;
   }
 
@@ -174,7 +167,7 @@ public Path finalizeDownload(Future<Path> download) throws IOException, Interrup
    * Downloads file to disk and returns path.
    *
    * <p>If the checksum and path to the repository cache is specified, attempt to load the file from
-   * the {@link RepositoryCache}. If it doesn't exist, proceed to download the file and load it into
+   * the {@link DownloadCache}. If it doesn't exist, proceed to download the file and load it into
    * the cache prior to returning the value.
    *
    * @param originalUrls list of mirror URLs with identical content
@@ -394,7 +387,7 @@ private boolean isRetryableException(Throwable e) {
    * <code>--repository_disable_download</code>.
    *
    * <p>If the checksum and path to the repository cache is specified, attempt to load the file from
-   * the {@link RepositoryCache}. If it doesn't exist, proceed to download the file and load it into
+   * the {@link DownloadCache}. If it doesn't exist, proceed to download the file and load it into
    * the cache prior to returning the value.
    *
    * @param originalUrl the original URL of the file
@@ -457,16 +450,21 @@ public byte[] downloadAndReadOneUrlForBzlmod(
       throw new IOException(getRewriterBlockedAllUrlsMessage(ImmutableList.of(originalUrl)));
     }
 
-    byte[] content;
+    var contentOut = new ByteArrayOutputStream();
     for (int attempt = 0; ; ++attempt) {
       try {
-        content =
-            bzlmodHttpDownloader.downloadAndReadOneUrl(
-                rewrittenUrls.get(0),
-                credentialFactory.create(authHeaders),
-                checksum,
-                eventHandler,
-                clientEnv);
+        downloader.download(
+            rewrittenUrls,
+            ImmutableMap.of(),
+            credentialFactory.create(authHeaders),
+            checksum,
+            /* canonicalId= */ null,
+            contentOut,
+            /* destinationPath= */ null,
+            eventHandler,
+            clientEnv,
+            Optional.empty(),
+            "Bazel registry download");
         break;
       } catch (InterruptedIOException e) {
         throw new InterruptedException(e.getMessage());
@@ -476,13 +474,12 @@ public byte[] downloadAndReadOneUrlForBzlmod(
         }
       }
     }
-    if (content == null) {
-      throw new IllegalStateException("Unexpected error: file should have been downloaded.");
-    }
+    byte[] content = contentOut.toByteArray();
 
     if (downloadCache.isEnabled()) {
       if (checksum.isPresent()) {
-        downloadCache.put(checksum.get().toString(), content, checksum.get().getKeyType());
+        downloadCache.put(
+            checksum.get().toString(), content, checksum.get().getKeyType());
       } else {
         downloadCache.put(content, KeyType.SHA256);
       }

src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/Downloader.java

@@ -15,13 +15,16 @@
 package com.google.devtools.build.lib.bazel.repository.downloader;
 
 import com.google.auth.Credentials;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.vfs.Path;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.net.URL;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import javax.annotation.Nullable;
 
 /** Interface for implementing the download of a file. */
 public interface Downloader {
@@ -35,7 +38,9 @@ public interface Downloader {
    * @param urls list of mirror URLs with identical content
    * @param credentials credentials to use when connecting to URLs
    * @param checksum valid checksum which is checked, or absent to disable
-   * @param output path to the destination file to write
+   * @param output the output stream to write the downloaded contents to
+   * @param destinationPath the path to the destination file, if any, or a human-readable and
+   *     path-like description of the desination (used for logging and error messages)
    * @param type extension, e.g. "tar.gz" to force on downloaded filename, or empty to not do this
    * @param context free-form string that describes the origin of the download for logging
    * @throws IOException if download was attempted and ended up failing
@@ -47,10 +52,40 @@ void download(
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,
-      Path output,
+      OutputStream output,
+      @Nullable String destinationPath,
       ExtendedEventHandler eventHandler,
       Map<String, String> clientEnv,
       Optional<String> type,
       String context)
       throws IOException, InterruptedException;
+
+  @VisibleForTesting
+  default void download(
+      List<URL> urls,
+      Map<String, List<String>> headers,
+      Credentials credentials,
+      Optional<Checksum> checksum,
+      String canonicalId,
+      Path destination,
+      ExtendedEventHandler eventHandler,
+      Map<String, String> clientEnv,
+      Optional<String> type,
+      String context)
+      throws IOException, InterruptedException {
+    try (var out = destination.getOutputStream()) {
+      download(
+          urls,
+          headers,
+          credentials,
+          checksum,
+          canonicalId,
+          out,
+          destination.getPathString(),
+          eventHandler,
+          clientEnv,
+          type,
+          context);
+    }
+  }
 }

src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/HttpDownloader.java

@@ -16,7 +16,6 @@
 
 import com.google.auth.Credentials;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.io.ByteStreams;
 import com.google.devtools.build.lib.buildeventstream.BuildEventStreamProtos.BuildEventId.FetchId;
@@ -27,8 +26,7 @@
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.util.JavaSleeper;
 import com.google.devtools.build.lib.util.Sleeper;
-import com.google.devtools.build.lib.vfs.Path;
-import java.io.ByteArrayOutputStream;
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.io.OutputStream;
@@ -41,6 +39,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Semaphore;
+import javax.annotation.Nullable;
 
 /**
  * HTTP implementation of {@link Downloader}.
@@ -79,7 +78,8 @@ public void download(
       Credentials credentials,
       Optional<Checksum> checksum,
       String canonicalId,
-      Path destination,
+      OutputStream destination,
+      @Nullable String destinationPath,
       ExtendedEventHandler eventHandler,
       Map<String, String> clientEnv,
       Optional<String> type,
@@ -97,7 +97,7 @@ public void download(
       semaphore.acquire();
 
       try (HttpStream payload = multiplexer.connect(url, checksum, headers, credentials, type);
-          OutputStream out = destination.getOutputStream()) {
+          OutputStream out = destination) {
         try {
           ByteStreams.copy(payload, out);
         } catch (SocketTimeoutException e) {
@@ -125,53 +125,34 @@ public void download(
     }
 
     if (!success) {
-      final IOException exception =
-          new IOException(
-              "Error downloading "
-                  + urls
-                  + " to "
-                  + destination
-                  + (ioExceptions.isEmpty()
-                      ? ""
-                      : ": " + Iterables.getLast(ioExceptions).getMessage()));
-
-      for (IOException cause : ioExceptions) {
+      // When destinationPath is null (in-memory download), don't add the "Error downloading"
+      // prefix - just use the original exception message to preserve backward compatibility.
+      var message =
+          destinationPath != null
+              ? "Error downloading %s to %s%s"
+                  .formatted(
+                      urls,
+                      destinationPath,
+                      ioExceptions.isEmpty()
+                          ? ""
+                          : ": " + Iterables.getLast(ioExceptions).getMessage())
+              : ioExceptions.isEmpty()
+                  ? "Error downloading " + urls
+                  : Iterables.getLast(ioExceptions).getMessage();
+      // Make it easy for callers to distinguish between not found and other IO errors.
+      var exception =
+          !ioExceptions.isEmpty()
+                  && ioExceptions.stream().allMatch(e -> e instanceof FileNotFoundException)
+              ? new FileNotFoundException(message)
+              : new IOException(message);
+      for (var cause : ioExceptions) {
         exception.addSuppressed(cause);
       }
 
       throw exception;
     }
   }
 
-  /** Downloads the contents of one URL and reads it into a byte array. */
-  public byte[] downloadAndReadOneUrl(
-      URL url,
-      Credentials credentials,
-      Optional<Checksum> checksum,
-      ExtendedEventHandler eventHandler,
-      Map<String, String> clientEnv)
-      throws IOException, InterruptedException {
-    HttpConnectorMultiplexer multiplexer = setUpConnectorMultiplexer(eventHandler, clientEnv);
-
-    ByteArrayOutputStream out = new ByteArrayOutputStream();
-    semaphore.acquire();
-    try (HttpStream payload =
-        multiplexer.connect(url, checksum, ImmutableMap.of(), credentials, Optional.empty())) {
-      ByteStreams.copy(payload, out);
-    } catch (SocketTimeoutException e) {
-      // SocketTimeoutExceptions are InterruptedIOExceptions; however they do not signify
-      // an external interruption, but simply a failed download due to some server timing
-      // out. So rethrow them as ordinary IOExceptions.
-      throw new IOException(e);
-    } catch (InterruptedIOException e) {
-      throw new InterruptedException(e.getMessage());
-    } finally {
-      semaphore.release();
-      // TODO(wyv): Do we need to report any event here?
-    }
-    return out.toByteArray();
-  }
-
   private HttpConnectorMultiplexer setUpConnectorMultiplexer(
       ExtendedEventHandler eventHandler, Map<String, String> clientEnv) {
     ProxyHelper proxyHelper = new ProxyHelper(clientEnv);

src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java

@@ -860,11 +860,11 @@ private void captureCorruptedOutputs(Exception e) {
     if (captureCorruptedOutputsDir != null) {
       if (e instanceof BulkTransferException) {
         for (Throwable suppressed : e.getSuppressed()) {
-          if (suppressed instanceof OutputDigestMismatchException) {
+          if (suppressed instanceof OutputDigestMismatchException outputDigestMismatchException) {
             // Capture corrupted outputs
             try {
-              String outputPath = ((OutputDigestMismatchException) suppressed).getOutputPath();
-              Path localPath = ((OutputDigestMismatchException) suppressed).getLocalPath();
+              String outputPath = outputDigestMismatchException.getOutputPath();
+              Path localPath = outputDigestMismatchException.getLocalPath();
               Path dst = captureCorruptedOutputsDir.getRelative(outputPath);
               dst.createDirectoryAndParents();