Skip to content

Instantly share code, notes, and snippets.

@LanderlYoung

LanderlYoung/ktd.kts

Last active Sep 28, 2018
Embed
What would you like to do?
Multi Thread Http Downloader implemented in kotlin script
#! /usr/bin/env kscript
@file:DependsOn("com.xenomachina:kotlin-argparser:2.0.7")
@file:DependsOn("me.tongfei:progressbar:0.7.1")
import com.xenomachina.argparser.ArgParser
import com.xenomachina.argparser.InvalidArgumentException
import com.xenomachina.argparser.SystemExitException
import com.xenomachina.argparser.default
import me.tongfei.progressbar.ProgressBar
import me.tongfei.progressbar.ProgressBarStyle
import java.io.RandomAccessFile
import java.net.HttpURLConnection
import java.net.URL
import java.util.*
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicLong
import kotlin.system.exitProcess
val param = parseArgument()
// curl -I http://i.imgur.com/z4d4kWk.jpg
val totalReadCount = AtomicLong()
val speedMeter = AtomicLong()
val url = URL(param.url)
val conn = url.openConnection() as HttpURLConnection
val threadPool = Executors.newCachedThreadPool()
val outputFile = RandomAccessFile(param.file, "rw")
conn.doInput = true
conn.doOutput = true
conn.connect()
if (conn.responseCode !in 200..299) {
conn.disconnect()
println("connection to url failed with code ${conn.responseCode}, url:${param.url}")
exitProcess(-1)
}
println("url: ${param.url}")
println("file: ${param.file}")
println("threads: ${param.threads}")
/*
Content-Length: 146515
Accept-Ranges: bytes
*/
val contentLength = try {
conn.getHeaderField("Content-Length").toLong()
} catch (e: NumberFormatException) {
-1L
}
val acceptRange = conn.getHeaderField("Accept-Ranges") == "bytes"
conn.disconnect()
if (acceptRange && contentLength != -1L) {
val segmentSize = contentLength / param.threads
for (job in 0 until param.threads) {
threadPool.submit {
downloadSegment(
segmentSize * job, if (job != param.threads - 1) {
segmentSize
} else {
contentLength - segmentSize * (param.threads - 1)
}
)
}
}
println("size: ${formatSize(contentLength)}")
println("segment: ${formatSize(segmentSize)}")
} else {
println("server doesn't support http-range, download with single thread!")
threadPool.submit {
downloadSegment(0, Long.MAX_VALUE, false)
}
}
threadPool.shutdown()
println()
val bar = ProgressBar("ktd", 100)
while (!threadPool.isTerminated) {
val speed = speedMeter.getAndSet(0L)
if (contentLength != -1L) {
val progress = totalReadCount.get() * 100L / contentLength
bar.stepTo(progress)
} else {
bar.maxHint(-1)
}
bar.extraMessage = "${formatSize(speed)}/s"
Thread.sleep(1000L)
}
bar.stepTo(bar.max)
bar.close()
// impl
fun downloadSegment(startByte: Long, size: Long, range: Boolean = true) {
val conn = url.openConnection() as HttpURLConnection
// Range: bytes=200-1000, 2000-6576, 19000-
if (range) {
conn.setRequestProperty("Range", "bytes=${startByte}-${startByte + size}")
}
conn.doInput = true
conn.doOutput = true
conn.connect()
val buffer = ByteArray(64 * 1024)
val input = conn.inputStream
var allBytes = 0L
while (allBytes < size) {
val count = input.read(buffer)
if (count == -1) {
break
}
synchronized(outputFile) {
outputFile.seek(startByte + allBytes)
outputFile.write(buffer, 0, count)
}
totalReadCount.addAndGet(count.toLong())
speedMeter.addAndGet(count.toLong())
allBytes += count
}
conn.disconnect()
}
fun formatSize(size: Long): String {
val K = 1024
val M = K * K
val G = M * K
return when {
size >= G -> String.format("%.2f GB", size.toDouble() / G)
size >= M -> String.format("%.2f MB", size.toDouble() / M)
size >= K -> String.format("%.2f KB", size.toDouble() / K)
else -> String.format("%d B", size)
}
}
class CLIArguments(parser: ArgParser) {
val threads by parser.storing(
"-t",
"--threads",
argName = "THREADS_COUNT",
help = "threads used to download") {
try {
toInt()
} catch (e: NumberFormatException) {
throw InvalidArgumentException("invalid thread count [$this]")
}
}
.default(4)
.addValidator {
if (this.value <= 0) {
throw InvalidArgumentException("thread count must be positive")
}
}
val url by parser.positional("URL", "url to download")
val file by parser.positional("FILE", "destination filename").default {
val u = URL(url)
val result = when {
u.path.isNotBlank() -> {
var path = u.path
var idx = path.lastIndexOf("/")
if (idx != -1 && idx != path.length - 1) {
path = path.substring(idx + 1)
}
idx = path.indexOf('?')
if (idx != -1) {
path.substring(0, idx)
} else {
path
}
}
u.host.isNotBlank() -> u.host
else -> ""
}
if (result.isBlank()) "ktd-file" else result
}
}
fun parseArgument(): CLIArguments {
try {
return ArgParser(args).parseInto(::CLIArguments)
} catch (e: SystemExitException) {
e.printAndExit("ktd")
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment