[WX-495] DRS Parallel Downloads

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

@@ -2,13 +2,12 @@ package drs.localizer
 
 import cats.data.NonEmptyList
 import cats.effect.{ExitCode, IO, IOApp}
-import cats.implicits._
-import cloud.nio.impl.drs.DrsPathResolver.{FatalRetryDisposition, RegularRetryDisposition}
+import cats.implicits.toTraverseOps
 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.toValidatedUriType
 import drs.localizer.downloaders._
 import org.apache.commons.csv.{CSVFormat, CSVParser}
 
@@ -17,7 +16,9 @@ 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] = {
@@ -42,47 +43,91 @@ object DrsLocalizerMain extends IOApp with StrictLogging {
     initialInterval = 10 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]): IO[Downloader] =
       IO.pure(GcsUriDownloader(gcsPath, serviceAccountJsonOption, downloadLoc, requesterPaysProjectOption))
+
+    override def buildBulkAccessUrlDownloader(urlsToDownload: List[ResolvedDrsUrl]): IO[Downloader] = {
+      IO.pure(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 {
+  /**
+    * Helper function to read a CSV file as a map from drs URL to requested 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)
+      val list = csvParser.getRecords.asScala.map(record => UnresolvedDrsUrl(record.get(0), record.get(1))).toList
+      list
+    }
+  }
+
+
+  def runLocalizer(commandLineArguments: CommandLineArguments, drsCredentials: DrsCredentials) : IO[ExitCode] = {
+    val urlList : IO[List[UnresolvedDrsUrl]] = 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))
+        loadCSVManifest(manifestPath)
       case None =>
-        val drsObject = commandLineArguments.drsObject.get
-        val containerPath = commandLineArguments.containerPath.get
-        localizeFile(commandLineArguments, drsCredentials, drsObject, containerPath)
+        IO.pure(List(UnresolvedDrsUrl(commandLineArguments.drsObject.get, commandLineArguments.containerPath.get)))
+      }
+    IO{
+      val main = new DrsLocalizerMain(urlList, defaultDownloaderFactory, drsCredentials, commandLineArguments.googleRequesterPaysProject)
+      main.resolveAndDownload().unsafeRunSync().exitCode
+    }
     }
-  }
 
-  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)
+  /**
+    * 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] = {
+    IO {
+      val downloaders: List[Downloader] = buildDownloaders().unsafeRunSync()
+      val results: List[DownloadResult] = downloaders.map(downloader => downloader.download.unsafeRunSync())
+      results.find(res => res != DownloadSuccess).getOrElse(DownloadSuccess)
+    }
+  }
+
   def getDrsPathResolver: IO[DrsLocalizerDrsPathResolver] = {
     IO {
       val drsConfig = DrsConfig.fromEnv(sys.env)
@@ -91,6 +136,63 @@ class DrsLocalizerMain(drsUrl: String,
     }
   }
 
+  /**
+    * Runs a synchronous HTTP request to resolve the provided DRS URL with the provided resolver.
+    */
+  def resolveSingleUrl(resolverObject: DrsLocalizerDrsPathResolver, drsUrlToResolve: UnresolvedDrsUrl): IO[ResolvedDrsUrl] = {
+    IO {
+      val fields = NonEmptyList.of(DrsResolverField.GsUri, DrsResolverField.GoogleServiceAccount, DrsResolverField.AccessUrl, DrsResolverField.Hashes)
+      //Insert retry logic here.
+      val drsResponse = resolverObject.resolveDrs(drsUrlToResolve.drsUrl, fields).unsafeRunSync()
+      ResolvedDrsUrl(drsResponse, drsUrlToResolve.downloadDestinationPath, toValidatedUriType(drsResponse.accessUrl, drsResponse.gsUri))
+    }
+  }
+
+  /**
+    * 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 =>
+          resolveSingleUrl(resolver, unresolvedUrl)
+        }.traverse(identity)
+      }
+    }
+  }
+
+  /**
+    * 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).flatMap { pendingDownloads =>
+      val accessUrls = pendingDownloads.filter(url => url.uriType == URIType.ACCESS)
+      val googleUrls = pendingDownloads.filter(url => url.uriType == URIType.GCS)
+      val bulkDownloader: Option[List[IO[Downloader]]] = if(accessUrls.isEmpty) None else Option(List(buildBulkAccessUrlDownloader(accessUrls)))
+      val googleDownloaders: Option[List[IO[Downloader]]] = if(googleUrls.isEmpty) None else  Option(buildGoogleDownloaders(googleUrls))
+      val combined: List[IO[Downloader]] = googleDownloaders.map(list => list).getOrElse(List()) ++ bulkDownloader.map(list => list).getOrElse(List())
+      combined.traverse(identity)
+    }
+  }
+
+  def buildGoogleDownloaders(resolvedGoogleUrls: List[ResolvedDrsUrl]) : List[IO[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]) : IO[Downloader] = {
+    downloaderFactory.buildBulkAccessUrlDownloader(resolvedUrls)
+  }
+
+
+  /*
   def resolveAndDownloadWithRetries(downloadRetries: Int,
                                     checksumRetries: Int,
                                     downloaderFactory: DownloaderFactory,
@@ -108,7 +210,8 @@ class DrsLocalizerMain(drsUrl: String,
         IO.raiseError(new RuntimeException(s"Exhausted $checksumRetries checksum retries to resolve, download and checksum $drsUrl", t))
       }
     }
-
+*/
+    /*
     def maybeRetryForDownloadFailure(t: Throwable): IO[DownloadResult] = {
       t match {
         case _: FatalRetryDisposition =>
@@ -121,7 +224,8 @@ class DrsLocalizerMain(drsUrl: String,
           IO.raiseError(new RuntimeException(s"Exhausted $downloadRetries download retries to resolve, download and checksum $drsUrl", t))
       }
     }
-
+*/
+      /*
     resolveAndDownload(downloaderFactory).redeemWith({
       maybeRetryForDownloadFailure
     },
@@ -134,33 +238,11 @@ class DrsLocalizerMain(drsUrl: String,
       case ChecksumFailure =>
         maybeRetryForChecksumFailure(new RuntimeException(s"Checksum failure for $drsUrl on checksum retry attempt $checksumAttempt of $checksumRetries"))
       case o => IO.pure(o)
-    })
-  }
+    }*
 
-  private [localizer] def resolveAndDownload(downloaderFactory: DownloaderFactory): IO[DownloadResult] = {
-    resolve(downloaderFactory) flatMap { _.download }
-  }
 
-  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))
-      }
-    } yield downloader
+    IO.raiseError(new RuntimeException(s"Exhausted $downloadRetries download retries to resolve, download and checksum $drsUrl", t))
   }
+         */
 }
+