WIP

oss-review-toolkit · Mar 15, 2024 · 8b2b524 · 8b2b524
1 parent b54aa02
commit 8b2b524
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 9 deletions.
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
@@ -150,6 +150,7 @@ maven-resolver-transport-http = { module = "org.apache.maven.resolver:maven-reso
 maven-resolver-transport-wagon = { module = "org.apache.maven.resolver:maven-resolver-transport-wagon", version.ref = "mavenResolver" }
 mockk = { module = "io.mockk:mockk", version.ref = "mockk" }
 mordant = { module = "com.github.ajalt.mordant:mordant", version.ref = "mordant" }
+mordantCoroutines = { module = "com.github.ajalt.mordant:mordant-coroutines", version.ref = "mordant" }
 okhttp = { module = "com.squareup.okhttp3:okhttp", version.ref = "okhttp" }
 postgres = { module = "org.postgresql:postgresql", version.ref = "postgres" }
 postgresEmbedded = { module = "com.opentable.components:otj-pg-embedded", version.ref = "postgresEmbedded" }

diff --git a/plugins/commands/downloader/build.gradle.kts b/plugins/commands/downloader/build.gradle.kts
@@ -33,4 +33,6 @@ dependencies {
 
     implementation(libs.clikt)
     implementation(libs.kotlinx.coroutines)
+    implementation(libs.mordant)
+    implementation(libs.mordantCoroutines)
 }
diff --git a/plugins/commands/downloader/src/main/kotlin/DownloaderCommand.kt b/plugins/commands/downloader/src/main/kotlin/DownloaderCommand.kt
@@ -20,6 +20,7 @@
 package org.ossreviewtoolkit.plugins.commands.downloader
 
 import com.github.ajalt.clikt.core.ProgramResult
+import com.github.ajalt.clikt.core.terminal
 import com.github.ajalt.clikt.parameters.groups.default
 import com.github.ajalt.clikt.parameters.groups.mutuallyExclusiveOptions
 import com.github.ajalt.clikt.parameters.groups.required
@@ -34,12 +35,24 @@ import com.github.ajalt.clikt.parameters.options.split
 import com.github.ajalt.clikt.parameters.options.switch
 import com.github.ajalt.clikt.parameters.types.enum
 import com.github.ajalt.clikt.parameters.types.file
+import com.github.ajalt.mordant.animation.coroutines.animateInCoroutine
+import com.github.ajalt.mordant.animation.progress.MultiProgressBarAnimation
+import com.github.ajalt.mordant.animation.progress.advance
+import com.github.ajalt.mordant.rendering.TextAlign
+import com.github.ajalt.mordant.widgets.progress.percentage
+import com.github.ajalt.mordant.widgets.progress.progressBar
+import com.github.ajalt.mordant.widgets.progress.progressBarContextLayout
+import com.github.ajalt.mordant.widgets.progress.progressBarLayout
+import com.github.ajalt.mordant.widgets.progress.text
+import com.github.ajalt.mordant.widgets.progress.timeElapsed
+import com.github.ajalt.mordant.widgets.progress.timeRemaining
 
 import java.io.File
 
 import kotlin.time.measureTime
 
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.runBlocking
@@ -316,21 +329,39 @@ class DownloaderCommand : OrtCommand(
         }
     }
 
+    @OptIn(ExperimentalCoroutinesApi::class)
     private suspend fun downloadAllPackages(
         packageDownloadDirs: Map<Package, File>,
-        failureMessages: MutableList<String>
+        failureMessages: MutableList<String>,
+        maxParallelDownloads: Int = 8
     ) {
-        withContext(Dispatchers.IO) {
-            packageDownloadDirs.entries.mapIndexed { index, (pkg, dir) ->
-                async {
-                    val progress = "${index + 1} of ${packageDownloadDirs.size}"
+        val overallLayout = progressBarLayout {
+            text(if (dryRun) "Verifying" else "Downloading", align = TextAlign.LEFT)
+            progressBar()
+            percentage()
+            timeRemaining()
+        }
+
+        val taskLayout = progressBarContextLayout<Package> {
+            text(fps = animationFps) { "> Package '${context.id.toCoordinates()}'..." }
+        }
 
-                    val verb = if (dryRun) "Verifying" else "Starting"
-                    echo("$verb download for '${pkg.id.toCoordinates()}' ($progress).")
+        val progress = MultiProgressBarAnimation(terminal).animateInCoroutine()
+        val overall = progress.addTask(overallLayout, context = Unit, total = packageDownloadDirs.size.toLong())
+        val tasks = List(maxParallelDownloads) { progress.addTask(taskLayout, context = Package.EMPTY, total = 1) }
 
-                    downloadPackage(pkg, dir, failureMessages).also {
-                        if (!dryRun) echo("Finished download for ${pkg.id.toCoordinates()} ($progress).")
+        progress.execute()
+
+        withContext(Dispatchers.IO.limitedParallelism(maxParallelDownloads)) {
+            packageDownloadDirs.entries.mapIndexed { index, (pkg, dir) ->
+                async {
+                    with(tasks[index % maxParallelDownloads]) {
+                        update { context = pkg }
+                        downloadPackage(pkg, dir, failureMessages)
+                        advance()
                     }
+
+                    overall.advance()
                 }
             }.awaitAll()
         }