Parallel downloading

KotatsuApp · Oct 15, 2024 · 10ec4a5 · 10ec4a5
1 parent 0c67a1c
commit 10ec4a5
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 27 deletions.
diff --git a/src/main/kotlin/org/koitharu/kotatsu/dl/Main.kt b/src/main/kotlin/org/koitharu/kotatsu/dl/Main.kt
@@ -3,12 +3,12 @@ package org.koitharu.kotatsu.dl
 import com.github.ajalt.clikt.command.main
 import com.github.ajalt.clikt.core.ProgramResult
 import com.github.ajalt.clikt.parameters.arguments.argument
-import com.github.ajalt.clikt.parameters.options.convert
-import com.github.ajalt.clikt.parameters.options.flag
-import com.github.ajalt.clikt.parameters.options.option
-import com.github.ajalt.clikt.parameters.options.validate
+import com.github.ajalt.clikt.parameters.options.*
 import com.github.ajalt.clikt.parameters.types.enum
 import com.github.ajalt.clikt.parameters.types.file
+import com.github.ajalt.clikt.parameters.types.int
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
 import org.koitharu.kotatsu.dl.download.DownloadFormat
 import org.koitharu.kotatsu.dl.download.MangaDownloader
 import org.koitharu.kotatsu.dl.parsers.MangaLoaderContextImpl
@@ -41,6 +41,12 @@ class Main : AppCommand(name = "kotatsu-dl") {
         ignoreCase = true,
         key = { it.name.lowercase() },
     )
+    private val parallelism: Int by option(
+        names = arrayOf("-j", "--jobs"),
+        help = "Number of parallel jobs for downloading",
+    ).int().default(1).check("Jobs count should be between 1 and 10") {
+        it in 1..10
+    }
     private val throttle: Boolean by option(
         names = arrayOf("--throttle"),
         help = "Slow down downloading to avoid blocking your IP address by server",
@@ -60,7 +66,7 @@ class Main : AppCommand(name = "kotatsu-dl") {
     override suspend fun invoke(): Int {
         val context = MangaLoaderContextImpl()
         val linkResolver = context.newLinkResolver(link)
-        print("Resolving link...")
+        print("Resolving link…")
         val source = linkResolver.getSource()
         if (source == null) {
             println()
@@ -112,8 +118,11 @@ class Main : AppCommand(name = "kotatsu-dl") {
             format = format,
             throttle = throttle,
             verbose = verbose,
+            parallelism = parallelism,
         )
-        val file = downloader.download()
+        val file = withContext(Dispatchers.Default) {
+            downloader.download()
+        }
         colored {
             print("Done.".green.bold)
             print(" Saved to ")

diff --git a/src/main/kotlin/org/koitharu/kotatsu/dl/download/MangaDownloader.kt b/src/main/kotlin/org/koitharu/kotatsu/dl/download/MangaDownloader.kt
@@ -1,9 +1,9 @@
 package org.koitharu.kotatsu.dl.download
 
 import androidx.collection.MutableIntList
-import kotlinx.coroutines.NonCancellable
-import kotlinx.coroutines.delay
-import kotlinx.coroutines.withContext
+import kotlinx.coroutines.*
+import kotlinx.coroutines.sync.Semaphore
+import kotlinx.coroutines.sync.withPermit
 import me.tongfei.progressbar.ProgressBar
 import me.tongfei.progressbar.ProgressBarBuilder
 import me.tongfei.progressbar.ProgressBarStyle
@@ -35,12 +35,13 @@ class MangaDownloader(
     private val format: DownloadFormat?,
     private val throttle: Boolean,
     private val verbose: Boolean,
+    private val parallelism: Int,
 ) {
 
     private val progressBarStyle = ProgressBarStyle.builder()
         .rightBracket("]")
         .leftBracket("[")
-        .colorCode(ColoredConsole.BRIGHT_YELLOW.toByte())
+        .colorCode(if (ColoredConsole.isSupported()) ColoredConsole.BRIGHT_YELLOW.toByte() else 0)
         .block('#')
         .build()
 
@@ -57,7 +58,7 @@ class MangaDownloader(
             .setTaskName("Downloading")
             .clearDisplayOnFinish()
             .build()
-        progressBar.setExtraMessage("Preparing...")
+        progressBar.setExtraMessage("Preparing…")
         val tempDir = Files.createTempDirectory("kdl_").toFile()
         val counters = MutableIntList()
         val totalChapters = chaptersRange.size(chapters)
@@ -70,6 +71,7 @@ class MangaDownloader(
                     file.delete()
                 }
             }
+            val semaphore = Semaphore(parallelism)
             for (chapter in chapters.withIndex()) {
                 progressBar.setExtraMessage(chapter.value.name)
                 if (chapter.index !in chaptersRange) {
@@ -78,25 +80,31 @@ class MangaDownloader(
                 val pages = runFailsafe(progressBar) { parser.getPages(chapter.value) }
                 counters.add(pages.size)
                 progressBar.maxHint(counters.sum().toLong() + (totalChapters - counters.size) * pages.size)
-                for ((pageIndex, page) in pages.withIndex()) {
-                    runFailsafe(progressBar) {
-                        val url = parser.getPageUrl(page)
-                        val file = downloadFile(url, tempDir, parser.source)
-                        output.addPage(
-                            chapter = chapter,
-                            file = file,
-                            pageNumber = pageIndex,
-                            ext = getFileExtensionFromUrl(url).orEmpty(),
-                        )
-                        progressBar.step()
-                        if (file.extension == "tmp") {
-                            file.delete()
+                coroutineScope {
+                    pages.mapIndexed { pageIndex, page ->
+                        launch {
+                            semaphore.withPermit {
+                                runFailsafe(progressBar) {
+                                    val url = parser.getPageUrl(page)
+                                    val file = downloadFile(url, tempDir, parser.source)
+                                    output.addPage(
+                                        chapter = chapter,
+                                        file = file,
+                                        pageNumber = pageIndex,
+                                        ext = getFileExtensionFromUrl(url).orEmpty(),
+                                    )
+                                    progressBar.step()
+                                    if (file.extension == "tmp") {
+                                        file.delete()
+                                    }
+                                }
+                            }
                         }
-                    }
+                    }.joinAll()
                 }
                 output.flushChapter(chapter.value)
             }
-            progressBar.setExtraMessage("Finalizing...")
+            progressBar.setExtraMessage("Finalizing…")
             output.mergeWithExisting()
             output.finish()
             progressBar.close()

diff --git a/src/main/kotlin/org/koitharu/kotatsu/dl/util/AppCommand.kt b/src/main/kotlin/org/koitharu/kotatsu/dl/util/AppCommand.kt
@@ -39,7 +39,7 @@ abstract class AppCommand(name: String) : CoreSuspendingCliktCommand(name) {
         val exitCode = runCatchingCancellable {
             invoke()
         }.onFailure { e ->
-            System.err.println(e.message)
+            e.printStackTrace()
         }.getOrDefault(1)
         throw ProgramResult(exitCode)
     }