diff --git a/src/main/resources/application.conf b/src/main/resources/application.conf index 7a24eab..e5d366e 100644 --- a/src/main/resources/application.conf +++ b/src/main/resources/application.conf @@ -6,4 +6,5 @@ ferload-client { manifest-separator="\t" download-files-pool=10 download-agreement="yes" + size-estimation-timeout=60 } diff --git a/src/main/scala/ca/ferlab/ferload/client/clients/S3Client.scala b/src/main/scala/ca/ferlab/ferload/client/clients/S3Client.scala index f349a69..0ab20ac 100644 --- a/src/main/scala/ca/ferlab/ferload/client/clients/S3Client.scala +++ b/src/main/scala/ca/ferlab/ferload/client/clients/S3Client.scala @@ -11,7 +11,7 @@ import java.io.File import java.net.URL import java.nio.file.{Files, Paths} import java.util.concurrent.Executors -import scala.concurrent.duration.Duration +import scala.concurrent.duration.{Duration, SECONDS} import scala.concurrent.{Await, ExecutionContext, ExecutionContextExecutorService, Future} class S3Client(nThreads: Int = 1) extends IS3 { @@ -23,10 +23,9 @@ class S3Client(nThreads: Int = 1) extends IS3 { private lazy val awsClient = AmazonS3ClientBuilder.standard().withRegion(Regions.DEFAULT_REGION).build() private lazy val trx = TransferManagerBuilder.standard().withS3Client(awsClient).build() - sys.addShutdownHook(trx.shutdownNow()) sys.addShutdownHook(trx.shutdownNow()) - override def getTotalExpectedDownloadSize(links: Map[String, String]): Long = { + override def getTotalExpectedDownloadSize(links: Map[String, String], timeout: Long): Long = { val sizes = Future.traverse(links.keySet)(fileName => { val link = links(fileName) Future({ @@ -34,11 +33,12 @@ class S3Client(nThreads: Int = 1) extends IS3 { Option(header).map(s => s.toLong).getOrElse(0L) }) }) - Await.result(sizes, Duration.Inf).sum + Await.result(sizes, Duration(timeout, SECONDS)).sum } override def download(outputDir: File, links: Map[String, String]): Set[File] = { - val padding = links.keySet.max.length + 1 + // find the file with the biggest name + val padding = links.keysIterator.reduceLeft((l,r) => if (l.length > r.length) l else r).length + 1 val downloads = Future.traverse(links.keySet)(fileName => { val link = links(fileName) Future(download(outputDir, fileName, link, padding)) diff --git a/src/main/scala/ca/ferlab/ferload/client/clients/inf/IS3.scala b/src/main/scala/ca/ferlab/ferload/client/clients/inf/IS3.scala index a417679..7355ae1 100644 --- a/src/main/scala/ca/ferlab/ferload/client/clients/inf/IS3.scala +++ b/src/main/scala/ca/ferlab/ferload/client/clients/inf/IS3.scala @@ -5,7 +5,7 @@ import java.io.File trait IS3 { def getTotalAvailableDiskSpaceAt(manifest: File): Long - def getTotalExpectedDownloadSize(links: Map[String, String]): Long + def getTotalExpectedDownloadSize(links: Map[String, String], timeout: Long): Long def download(outputDir: File, links: Map[String, String]): Set[File] } diff --git a/src/main/scala/ca/ferlab/ferload/client/commands/Download.scala b/src/main/scala/ca/ferlab/ferload/client/commands/Download.scala index 0732bc2..fc1b8b2 100644 --- a/src/main/scala/ca/ferlab/ferload/client/commands/Download.scala +++ b/src/main/scala/ca/ferlab/ferload/client/commands/Download.scala @@ -12,7 +12,7 @@ import picocli.CommandLine.{Command, IExitCodeGenerator, Option} import java.io.{File, FileReader} import java.util.Optional -import scala.util.{Failure, Success, Using} +import scala.util.{Failure, Success, Try, Using} @Command(name = "download", mixinStandardHelpOptions = true, description = Array("Download files based on provided manifest."), version = Array("0.1")) @@ -78,11 +78,19 @@ class Download(userConfig: UserConfig, } val links: Map[String, String] = CommandBlock("Retrieve Ferload download link(s)", successEmoji, padding) { - ferload.getDownloadLinks(token, manifestContent) + Try(ferload.getDownloadLinks(token, manifestContent)) match { + case Success(links) => links + case Failure(e) => { + // always refresh token if failed + userConfig.remove(Token) + userConfig.save() + throw e + } + } } val totalExpectedDownloadSize = CommandBlock("Compute total average expected download size", successEmoji, padding) { - s3.getTotalExpectedDownloadSize(links) + s3.getTotalExpectedDownloadSize(links, appConfig.getLong("size-estimation-timeout")) } val downloadAgreement = appConfig.getString("download-agreement") diff --git a/src/main/scala/ca/ferlab/ferload/client/configurations/UserConfig.scala b/src/main/scala/ca/ferlab/ferload/client/configurations/UserConfig.scala index e10fe99..8a89473 100644 --- a/src/main/scala/ca/ferlab/ferload/client/configurations/UserConfig.scala +++ b/src/main/scala/ca/ferlab/ferload/client/configurations/UserConfig.scala @@ -17,6 +17,10 @@ class UserConfig(val path: String) { def set(configName: UserConfigName, value: String): Unit = { Option(value).map(props.put(configName.name, _)) } + + def remove(configName: UserConfigName): Unit = { + props.remove(configName.name) + } def clear(): Unit = { props.clear() diff --git a/src/test/scala/ca/ferlab/ferload/client/commands/DownloadTest.scala b/src/test/scala/ca/ferlab/ferload/client/commands/DownloadTest.scala index 83fea9f..2b370be 100644 --- a/src/test/scala/ca/ferlab/ferload/client/commands/DownloadTest.scala +++ b/src/test/scala/ca/ferlab/ferload/client/commands/DownloadTest.scala @@ -118,7 +118,7 @@ class DownloadTest extends AnyFunSuite with BeforeAndAfter { Set(new File("f1"), new File("f2")) } - override def getTotalExpectedDownloadSize(links: Map[String, String]): Long = 0L + override def getTotalExpectedDownloadSize(links: Map[String, String], timeout: Long): Long = 0L override def getTotalAvailableDiskSpaceAt(manifest: File): Long = 1L } @@ -135,7 +135,7 @@ class DownloadTest extends AnyFunSuite with BeforeAndAfter { Set(new File("f1"), new File("f2")) } - override def getTotalExpectedDownloadSize(links: Map[String, String]): Long = 0L + override def getTotalExpectedDownloadSize(links: Map[String, String], timeout: Long): Long = 0L override def getTotalAvailableDiskSpaceAt(manifest: File): Long = 1L } @@ -152,7 +152,7 @@ class DownloadTest extends AnyFunSuite with BeforeAndAfter { Set(new File("f1"), new File("f2")) } - override def getTotalExpectedDownloadSize(links: Map[String, String]): Long = 2L + override def getTotalExpectedDownloadSize(links: Map[String, String], timeout: Long): Long = 2L override def getTotalAvailableDiskSpaceAt(manifest: File): Long = 1L } @@ -169,7 +169,7 @@ class DownloadTest extends AnyFunSuite with BeforeAndAfter { Set(new File("f1"), new File("f2")) } - override def getTotalExpectedDownloadSize(links: Map[String, String]): Long = 1L + override def getTotalExpectedDownloadSize(links: Map[String, String], timeout: Long): Long = 1L override def getTotalAvailableDiskSpaceAt(manifest: File): Long = 2L } diff --git a/src/test/scala/ca/ferlab/ferload/client/configurations/UserConfigTest.scala b/src/test/scala/ca/ferlab/ferload/client/configurations/UserConfigTest.scala index cb711b7..12d304c 100644 --- a/src/test/scala/ca/ferlab/ferload/client/configurations/UserConfigTest.scala +++ b/src/test/scala/ca/ferlab/ferload/client/configurations/UserConfigTest.scala @@ -15,9 +15,11 @@ class UserConfigTest extends AnyFunSuite with BeforeAndAfter { userConfig.clear() } - test("set / get") { + test("set / get / remove") { userConfig.set(Username, "foo") assert(userConfig.get(Username).equals("foo")) + userConfig.remove(Username) + assert(userConfig.get(Username) == null) } test("clear") {