[WX-495] DRS Parallel Downloads

cromwell-drs-localizer/src/main/scala/drs/localizer/DrsLocalizerMain.scala

@@ -2,13 +2,13 @@ package drs.localizer
 
 import cats.data.NonEmptyList
 import cats.effect.{ExitCode, IO, IOApp}
-import cats.implicits._
+import cats.implicits.toTraverseOps
 import cloud.nio.impl.drs.DrsPathResolver.{FatalRetryDisposition, RegularRetryDisposition}
 import cloud.nio.impl.drs._
 import cloud.nio.spi.{CloudNioBackoff, CloudNioSimpleExponentialBackoff}
 import com.typesafe.scalalogging.StrictLogging
 import drs.localizer.CommandLineParser.AccessTokenStrategy.{Azure, Google}
-import drs.localizer.downloaders.AccessUrlDownloader.Hashes
+import drs.localizer.DrsLocalizerMain.{defaultNumRetries, toValidatedUriType}
 import drs.localizer.downloaders._
 import org.apache.commons.csv.{CSVFormat, CSVParser}
 
@@ -17,7 +17,10 @@ import java.nio.charset.Charset
 import scala.concurrent.duration._
 import scala.jdk.CollectionConverters._
 import scala.language.postfixOps
+import drs.localizer.URIType.URIType
 
+case class UnresolvedDrsUrl(drsUrl: String, downloadDestinationPath: String)
+case class ResolvedDrsUrl(drsResponse: DrsResolverResponse, downloadDestinationPath: String, uriType: URIType)
 object DrsLocalizerMain extends IOApp with StrictLogging {
 
   override def run(args: List[String]): IO[ExitCode] = {
@@ -38,51 +41,102 @@ object DrsLocalizerMain extends IOApp with StrictLogging {
 
   def buildParser(): scopt.OptionParser[CommandLineArguments] = new CommandLineParser()
 
+  // Default retry parameters for resolving a DRS url
+  val defaultNumRetries: Int = 5
   val defaultBackoff: CloudNioBackoff = CloudNioSimpleExponentialBackoff(
-    initialInterval = 10 seconds, maxInterval = 60 seconds, multiplier = 2)
+    initialInterval = 1 seconds, maxInterval = 60 seconds, multiplier = 2)
 
   val defaultDownloaderFactory: DownloaderFactory = new DownloaderFactory {
-    override def buildAccessUrlDownloader(accessUrl: AccessUrl, downloadLoc: String, hashes: Hashes): IO[Downloader] =
-      IO.pure(AccessUrlDownloader(accessUrl, downloadLoc, hashes))
+    override def buildGcsUriDownloader(gcsPath: String, serviceAccountJsonOption: Option[String], downloadLoc: String, requesterPaysProjectOption: Option[String]): Downloader =
+      GcsUriDownloader(gcsPath, serviceAccountJsonOption, downloadLoc, requesterPaysProjectOption)
 
-    override def buildGcsUriDownloader(gcsPath: String, serviceAccountJsonOption: Option[String], downloadLoc: String, requesterPaysProjectOption: Option[String]): IO[Downloader] =
-      IO.pure(GcsUriDownloader(gcsPath, serviceAccountJsonOption, downloadLoc, requesterPaysProjectOption))
+    override def buildBulkAccessUrlDownloader(urlsToDownload: List[ResolvedDrsUrl]): Downloader = {
+      BulkAccessUrlDownloader(urlsToDownload)
+    }
   }
 
   private def printUsage: IO[ExitCode] = {
     System.err.println(CommandLineParser.Usage)
     IO.pure(ExitCode.Error)
   }
 
-  def runLocalizer(commandLineArguments: CommandLineArguments, drsCredentials: DrsCredentials): IO[ExitCode] = {
-    commandLineArguments.manifestPath match {
-      case Some(manifestPath) =>
-        val manifestFile = new File(manifestPath)
-        val csvParser = CSVParser.parse(manifestFile, Charset.defaultCharset(), CSVFormat.DEFAULT)
-        val exitCodes: IO[List[ExitCode]] = csvParser.asScala.map(record => {
-          val drsObject = record.get(0)
-          val containerPath = record.get(1)
-          localizeFile(commandLineArguments, drsCredentials, drsObject, containerPath)
-        }).toList.sequence
-        exitCodes.map(_.find(_ != ExitCode.Success).getOrElse(ExitCode.Success))
-      case None =>
-        val drsObject = commandLineArguments.drsObject.get
-        val containerPath = commandLineArguments.containerPath.get
-        localizeFile(commandLineArguments, drsCredentials, drsObject, containerPath)
+  /**
+    * Helper function to read a CSV file as pairs of drsURL -> local download destination.
+    * @param csvManifestPath Path to a CSV file where each row is something like: drs://asdf.ghj, path/to/my/directory
+    */
+  def loadCSVManifest(csvManifestPath: String): IO[List[UnresolvedDrsUrl]] = {
+    IO {
+      val openFile = new File(csvManifestPath)
+      val csvParser = CSVParser.parse(openFile, Charset.defaultCharset(), CSVFormat.DEFAULT)
+      try{
+        csvParser.getRecords.asScala.map(record => UnresolvedDrsUrl(record.get(0), record.get(1))).toList
+      } finally {
+        csvParser.close()
+      }
     }
   }
 
-  private def localizeFile(commandLineArguments: CommandLineArguments, drsCredentials: DrsCredentials, drsObject: String, containerPath: String) = {
-    new DrsLocalizerMain(drsObject, containerPath, drsCredentials, commandLineArguments.googleRequesterPaysProject).
-      resolveAndDownloadWithRetries(downloadRetries = 3, checksumRetries = 1, defaultDownloaderFactory, Option(defaultBackoff)).map(_.exitCode)
+  def runLocalizer(commandLineArguments: CommandLineArguments, drsCredentials: DrsCredentials) : IO[ExitCode] = {
+    val urlList = (commandLineArguments.manifestPath, commandLineArguments.drsObject, commandLineArguments.containerPath) match {
+      case (Some(manifestPath), _, _) => {
+        loadCSVManifest(manifestPath)
+      }
+      case (_, Some(drsObject), Some(containerPath)) => {
+        IO.pure(List(UnresolvedDrsUrl(drsObject, containerPath)))
+      }
+      case(_,_,_) => {
+        throw new RuntimeException("Illegal command line arguments supplied to drs localizer.")
+      }
+    }
+    val main = new DrsLocalizerMain(urlList, defaultDownloaderFactory, drsCredentials, commandLineArguments.googleRequesterPaysProject)
+    main.resolveAndDownload().map(_.exitCode)
+    }
+
+  /**
+    * Helper function to decide which downloader to use based on data from the DRS response.
+    * Throws a runtime exception if the DRS response is invalid.
+    */
+  def toValidatedUriType(accessUrl: Option[AccessUrl], gsUri: Option[String]): URIType = {
+    // if both are provided, prefer using access urls
+    (accessUrl, gsUri) match {
+      case (Some(_), _) =>
+        if(!accessUrl.get.url.startsWith("https://")) { throw new RuntimeException("Resolved Access URL does not start with https://")}
+        URIType.ACCESS
+      case (_, Some(_)) =>
+        if(!gsUri.get.startsWith("gs://")) { throw new RuntimeException("Resolved Google URL does not start with gs://")}
+        URIType.GCS
+      case (_, _) =>
+        throw new RuntimeException("DRS response did not contain any URLs")
+    }
   }
+  }
+
+object URIType extends Enumeration {
+  type URIType = Value
+  val GCS, ACCESS, UNKNOWN = Value
 }
 
-class DrsLocalizerMain(drsUrl: String,
-                       downloadLoc: String,
+class DrsLocalizerMain(toResolveAndDownload: IO[List[UnresolvedDrsUrl]],
+                       downloaderFactory: DownloaderFactory,
                        drsCredentials: DrsCredentials,
                        requesterPaysProjectIdOption: Option[String]) extends StrictLogging {
 
+  /**
+    * This will:
+    *   - resolve all URLS
+    *   - build downloader(s) for them
+    *   - Invoke the downloaders to localize the files.
+    * @return DownloadSuccess if all downloads succeed. An error otherwise.
+    */
+  def resolveAndDownload(): IO[DownloadResult] = {
+    val downloadResults = buildDownloaders().flatMap { downloaderList =>
+      downloaderList.map(downloader => downloader.download).traverse(identity)
+    }
+    downloadResults.map{list =>
+      list.find(result => result != DownloadSuccess).getOrElse(DownloadSuccess)
+    }
+  }
+
   def getDrsPathResolver: IO[DrsLocalizerDrsPathResolver] = {
     IO {
       val drsConfig = DrsConfig.fromEnv(sys.env)
@@ -91,76 +145,86 @@ class DrsLocalizerMain(drsUrl: String,
     }
   }
 
-  def resolveAndDownloadWithRetries(downloadRetries: Int,
-                                    checksumRetries: Int,
-                                    downloaderFactory: DownloaderFactory,
-                                    backoff: Option[CloudNioBackoff],
-                                    downloadAttempt: Int = 0,
-                                    checksumAttempt: Int = 0): IO[DownloadResult] = {
-
-    def maybeRetryForChecksumFailure(t: Throwable): IO[DownloadResult] = {
-      if (checksumAttempt < checksumRetries) {
-        backoff foreach { b => Thread.sleep(b.backoffMillis) }
-        logger.warn(s"Attempting retry $checksumAttempt of $checksumRetries checksum retries to download $drsUrl", t)
-        // In the event of a checksum failure reset the download attempt to zero.
-        resolveAndDownloadWithRetries(downloadRetries, checksumRetries, downloaderFactory, backoff map { _.next }, 0, checksumAttempt + 1)
-      } else {
-        IO.raiseError(new RuntimeException(s"Exhausted $checksumRetries checksum retries to resolve, download and checksum $drsUrl", t))
-      }
+  /**
+    * After resolving all of the URLs, this sorts them into an "Access" or "GCS" bucket.
+    * All access URLS will be downloaded as a batch with a single bulk downloader.
+    * All google URLs will be downloaded individually in their own google downloader.
+    * @return List of all downloaders required to fulfill the request.
+    */
+  def buildDownloaders() : IO[List[Downloader]] = {
+    resolveUrls(toResolveAndDownload).map { pendingDownloads =>
+      val accessUrls = pendingDownloads.filter(url => url.uriType == URIType.ACCESS)
+      val googleUrls = pendingDownloads.filter(url => url.uriType == URIType.GCS)
+      val bulkDownloader: List[Downloader] = if (accessUrls.isEmpty) List() else List(buildBulkAccessUrlDownloader(accessUrls))
+      val googleDownloaders: List[Downloader] = if (googleUrls.isEmpty) List() else buildGoogleDownloaders(googleUrls)
+      bulkDownloader ++ googleDownloaders
     }
+  }
 
-    def maybeRetryForDownloadFailure(t: Throwable): IO[DownloadResult] = {
-      t match {
-        case _: FatalRetryDisposition =>
-          IO.raiseError(t)
-        case _ if downloadAttempt < downloadRetries =>
-          backoff foreach { b => Thread.sleep(b.backoffMillis) }
-          logger.warn(s"Attempting retry $downloadAttempt of $downloadRetries download retries to download $drsUrl", t)
-          resolveAndDownloadWithRetries(downloadRetries, checksumRetries, downloaderFactory, backoff map { _.next }, downloadAttempt + 1, checksumAttempt)
-        case _ =>
-          IO.raiseError(new RuntimeException(s"Exhausted $downloadRetries download retries to resolve, download and checksum $drsUrl", t))
-      }
+  def buildGoogleDownloaders(resolvedGoogleUrls: List[ResolvedDrsUrl]) : List[Downloader] = {
+    resolvedGoogleUrls.map{url=>
+      downloaderFactory.buildGcsUriDownloader(
+        gcsPath = url.drsResponse.gsUri.get,
+        serviceAccountJsonOption = url.drsResponse.googleServiceAccount.map(_.data.spaces2),
+        downloadLoc = url.downloadDestinationPath,
+        requesterPaysProjectOption = requesterPaysProjectIdOption)
     }
+  }
+  def buildBulkAccessUrlDownloader(resolvedUrls: List[ResolvedDrsUrl]) : Downloader = {
+    downloaderFactory.buildBulkAccessUrlDownloader(resolvedUrls)
+  }
 
-    resolveAndDownload(downloaderFactory).redeemWith({
-      maybeRetryForDownloadFailure
-    },
-    {
-      case f: FatalDownloadFailure =>
-        IO.raiseError(new RuntimeException(s"Fatal error downloading DRS object: $f"))
-      case r: RetryableDownloadFailure =>
-        maybeRetryForDownloadFailure(
-          new RuntimeException(s"Retryable download error: $r for $drsUrl on retry attempt $downloadAttempt of $downloadRetries") with RegularRetryDisposition)
-      case ChecksumFailure =>
-        maybeRetryForChecksumFailure(new RuntimeException(s"Checksum failure for $drsUrl on checksum retry attempt $checksumAttempt of $checksumRetries"))
-      case o => IO.pure(o)
-    })
+  /**
+    * Runs a synchronous HTTP request to resolve the provided DRS URL with the provided resolver.
+    */
+  def resolveSingleUrl(resolverObject: DrsLocalizerDrsPathResolver, drsUrlToResolve: UnresolvedDrsUrl): IO[ResolvedDrsUrl] = {
+    val fields = NonEmptyList.of(DrsResolverField.GsUri, DrsResolverField.GoogleServiceAccount, DrsResolverField.AccessUrl, DrsResolverField.Hashes)
+    val drsResponse = resolverObject.resolveDrs(drsUrlToResolve.drsUrl, fields)
+    drsResponse.map(resp => ResolvedDrsUrl(resp, drsUrlToResolve.downloadDestinationPath, toValidatedUriType(resp.accessUrl, resp.gsUri)))
   }
 
-  private [localizer] def resolveAndDownload(downloaderFactory: DownloaderFactory): IO[DownloadResult] = {
-    resolve(downloaderFactory) flatMap { _.download }
+
+  val defaultBackoff: CloudNioBackoff = CloudNioSimpleExponentialBackoff(
+    initialInterval = 10 seconds, maxInterval = 60 seconds, multiplier = 2)
+
+  /**
+    * Runs synchronous HTTP requests to resolve all the DRS urls.
+    */
+  def resolveUrls(unresolvedUrls: IO[List[UnresolvedDrsUrl]]): IO[List[ResolvedDrsUrl]] = {
+    unresolvedUrls.flatMap { unresolvedList =>
+      getDrsPathResolver.flatMap { resolver =>
+        unresolvedList.map { unresolvedUrl =>
+          resolveWithRetries(resolver, unresolvedUrl, defaultNumRetries, Option(defaultBackoff))
+        }.traverse(identity)
+      }
+    }
   }
 
-  private [localizer] def resolve(downloaderFactory: DownloaderFactory): IO[Downloader] = {
-    val fields = NonEmptyList.of(DrsResolverField.GsUri, DrsResolverField.GoogleServiceAccount, DrsResolverField.AccessUrl, DrsResolverField.Hashes)
-    for {
-      resolver <- getDrsPathResolver
-      drsResolverResponse <- resolver.resolveDrs(drsUrl, fields)
-
-      // Currently DRS Resolver only supports resolving DRS paths to access URLs or GCS paths.
-      downloader <- (drsResolverResponse.accessUrl, drsResolverResponse.gsUri) match {
-        case (Some(accessUrl), _) =>
-          downloaderFactory.buildAccessUrlDownloader(accessUrl, downloadLoc, drsResolverResponse.hashes)
-        case (_, Some(gcsPath)) =>
-          val serviceAccountJsonOption = drsResolverResponse.googleServiceAccount.map(_.data.spaces2)
-          downloaderFactory.buildGcsUriDownloader(
-            gcsPath = gcsPath,
-            serviceAccountJsonOption = serviceAccountJsonOption,
-            downloadLoc = downloadLoc,
-            requesterPaysProjectOption = requesterPaysProjectIdOption)
-        case _ =>
-          IO.raiseError(new RuntimeException(DrsPathResolver.ExtractUriErrorMsg))
+  def resolveWithRetries(resolverObject: DrsLocalizerDrsPathResolver,
+                         drsUrlToResolve: UnresolvedDrsUrl,
+                         resolutionRetries: Int,
+                         backoff: Option[CloudNioBackoff],
+                        resolutionAttempt: Int = 0) : IO[ResolvedDrsUrl] = {
+
+    def maybeRetryForResolutionFailure(t: Throwable): IO[ResolvedDrsUrl] = {
+      if (resolutionAttempt < resolutionRetries) {
+        backoff foreach { b => Thread.sleep(b.backoffMillis) }
+        logger.warn(s"Attempting retry $resolutionAttempt of $resolutionRetries drs resolution retries to resolve ${drsUrlToResolve.drsUrl}", t)
+        resolveWithRetries(resolverObject, drsUrlToResolve, resolutionRetries, backoff map { _.next }, resolutionAttempt+1)
+      } else {
+        IO.raiseError(new RuntimeException(s"Exhausted $resolutionRetries resolution retries to resolve $drsUrlToResolve.drsUrl", t))
       }
-    } yield downloader
+    }
+
+    resolveSingleUrl(resolverObject, drsUrlToResolve).redeemWith(
+      recover = maybeRetryForResolutionFailure,
+      bind = {
+      case f: FatalRetryDisposition =>
+        IO.raiseError(new RuntimeException(s"Fatal error resolving DRS URL: $f"))
+      case _: RegularRetryDisposition =>
+        resolveWithRetries(resolverObject, drsUrlToResolve, resolutionRetries, backoff, resolutionAttempt+1)
+      case o => IO.pure(o)
+    })
   }
 }
+