From e757910f4eb7311a527961aeb0d45cbdcf21add3 Mon Sep 17 00:00:00 2001
From: MuYu <mr.muzea@gmail.com>
Date: Mon, 28 Oct 2019 21:20:19 +0800
Subject: [PATCH] feat: single thread speed limit

---
 config/config.go         |  2 ++
 downloader/downloader.go | 12 +++++++++++-
 go.mod                   |  2 ++
 go.sum                   |  4 ++++
 main.go                  | 18 +++++++++++++++++-
 5 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/config/config.go b/config/config.go
index 0b46f83b9..62f5bd6db 100644
--- a/config/config.go
+++ b/config/config.go
@@ -57,6 +57,8 @@ var (
 	RetryTimes int
 	// YouTubeStream2 will use data in `url_encoded_fmt_stream_map`
 	YouTubeStream2 bool
+	// SingleThreadSpeedLimit will limit the speed of each thread
+	SingleThreadSpeedLimit int64
 )
 
 // FakeHeaders fake http headers
diff --git a/downloader/downloader.go b/downloader/downloader.go
index 914c6ad43..db406e635 100644
--- a/downloader/downloader.go
+++ b/downloader/downloader.go
@@ -15,6 +15,7 @@ import (
 	"github.com/iawia002/annie/config"
 	"github.com/iawia002/annie/request"
 	"github.com/iawia002/annie/utils"
+	"github.com/juju/ratelimit"
 )
 
 func progressBar(size int64) *pb.ProgressBar {
@@ -63,7 +64,16 @@ func writeFile(
 	writer := io.MultiWriter(file, bar)
 	// Note that io.Copy reads 32kb(maximum) from input and writes them to output, then repeats.
 	// So don't worry about memory.
-	written, copyErr := io.Copy(writer, res.Body)
+	var (
+		written int64
+		copyErr error
+	)
+	if config.SingleThreadSpeedLimit <= 0 {
+		written, copyErr = io.Copy(writer, res.Body)
+	} else {
+		bucket := ratelimit.NewBucketWithRate(float64(config.SingleThreadSpeedLimit), 4*1024)
+		written, copyErr = io.Copy(writer, ratelimit.Reader(res.Body, bucket))
+	}
 	if copyErr != nil && copyErr != io.EOF {
 		return written, fmt.Errorf("file copy error: %s", copyErr)
 	}
diff --git a/go.mod b/go.mod
index f71b69a7f..5b361c997 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,9 @@ require (
 	github.com/PuerkitoBio/goquery v1.4.1
 	github.com/andybalholm/cascadia v1.0.0 // indirect
 	github.com/cheggaaa/pb v1.0.25
+	github.com/docker/go-units v0.4.0
 	github.com/fatih/color v1.7.0
+	github.com/juju/ratelimit v1.0.1
 	github.com/kr/pretty v0.1.0
 	github.com/mattn/go-colorable v0.0.9 // indirect
 	github.com/mattn/go-isatty v0.0.3 // indirect
diff --git a/go.sum b/go.sum
index 9476713f0..d7f3a4733 100644
--- a/go.sum
+++ b/go.sum
@@ -8,8 +8,12 @@ github.com/cheggaaa/pb v1.0.25 h1:tFpebHTkI7QZx1q1rWGOKhbunhZ3fMaxTvHDWn1bH/4=
 github.com/cheggaaa/pb v1.0.25/go.mod h1:pQciLPpbU0oxA0h+VJYYLxO+XeDQb5pZijXscXHm81s=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw=
+github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
+github.com/juju/ratelimit v1.0.1 h1:+7AIFJVQ0EQgq/K9+0Krm7m530Du7tIz0METWzN0RgY=
+github.com/juju/ratelimit v1.0.1/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
diff --git a/main.go b/main.go
index f75ac2d65..350b8e569 100644
--- a/main.go
+++ b/main.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/fatih/color"
 
+	goUnits "github.com/docker/go-units"
 	"github.com/iawia002/annie/config"
 	"github.com/iawia002/annie/downloader"
 	"github.com/iawia002/annie/extractors/bcy"
@@ -84,6 +85,22 @@ func init() {
 	flag.StringVar(&config.YoukuPassword, "password", "", "Youku password")
 	// youtube
 	flag.BoolVar(&config.YouTubeStream2, "ytb-stream2", false, "Use data in url_encoded_fmt_stream_map")
+	var singleThreadSpeedLimit string
+	flag.StringVar(&singleThreadSpeedLimit, "singleThreadSpeedLimit", "", "Limit the speed of each thread, 0 is unlimited")
+
+	flag.Parse()
+
+	if singleThreadSpeedLimit == "" || singleThreadSpeedLimit == "0" {
+		config.SingleThreadSpeedLimit = 0
+	} else {
+		size, err := goUnits.FromHumanSize(singleThreadSpeedLimit)
+		if err != nil {
+			printError("", err)
+			config.SingleThreadSpeedLimit = 0
+		} else {
+			config.SingleThreadSpeedLimit = size
+		}
+	}
 }
 
 func printError(url string, err error) {
@@ -192,7 +209,6 @@ func download(videoURL string) bool {
 }
 
 func main() {
-	flag.Parse()
 	args := flag.Args()
 	if config.Version {
 		utils.PrintVersion()