From 71cbf3d996ae9f2cc020e3f6c05c56e1bdf6b384 Mon Sep 17 00:00:00 2001 From: MuYu Date: Mon, 28 Oct 2019 21:20:19 +0800 Subject: [PATCH 1/3] 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 244b55539..d80fc2860 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 0e3f18a80..76f1f0a2a 100644 --- a/main.go +++ b/main.go @@ -10,6 +10,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" @@ -83,6 +84,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) { @@ -191,7 +208,6 @@ func download(videoURL string) bool { } func main() { - flag.Parse() args := flag.Args() if config.Version { utils.PrintVersion() From c706711e5fe2bd871d8dbc7883a660e6f91df9aa Mon Sep 17 00:00:00 2001 From: Stegosawr <38107067+Stegosawr@users.noreply.github.com> Date: Tue, 29 Oct 2019 16:20:24 +0100 Subject: [PATCH 2/3] Put the logic into the main method, so it's back in order again --- main.go | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/main.go b/main.go index 76f1f0a2a..0c659be5c 100644 --- a/main.go +++ b/main.go @@ -84,22 +84,9 @@ 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) { @@ -208,6 +195,7 @@ func download(videoURL string) bool { } func main() { + flag.Parse() args := flag.Args() if config.Version { utils.PrintVersion() @@ -246,6 +234,18 @@ func main() { config.Cookie = string(data) } } + var singleThreadSpeedLimit string + 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 + } + } var isErr bool for _, videoURL := range args { if !download(strings.TrimSpace(videoURL)) { From 201b590c2c32b635b7c47b2cc5f92f9fb4d097ed Mon Sep 17 00:00:00 2001 From: MuYu Date: Sat, 9 Nov 2019 14:29:36 +0800 Subject: [PATCH 3/3] doc: single thread speed limit --- README.md | 2 ++ main.go | 5 +++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 5169a32e1..bdddfa180 100644 --- a/README.md +++ b/README.md @@ -514,6 +514,8 @@ $ annie -j https://www.bilibili.com/video/av20203945 Use specified Referrer -cs int HTTP chunk size for downloading (in MB) (default 0) + -singleThreadSpeedLimit string + Limit the speed of each thread, 0 is unlimited, acceptable input "10MiB", "200KB" (default 0) ``` #### Network: diff --git a/main.go b/main.go index 0c659be5c..207f01853 100644 --- a/main.go +++ b/main.go @@ -41,6 +41,8 @@ import ( "github.com/iawia002/annie/utils" ) +var singleThreadSpeedLimit string + func init() { flag.BoolVar(&config.Debug, "d", false, "Debug mode") flag.BoolVar(&config.Version, "v", false, "Show version") @@ -84,7 +86,7 @@ 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") - + flag.StringVar(&singleThreadSpeedLimit, "singleThreadSpeedLimit", "", "Limit the speed of each thread, 0 is unlimited") } @@ -234,7 +236,6 @@ func main() { config.Cookie = string(data) } } - var singleThreadSpeedLimit string if singleThreadSpeedLimit == "" || singleThreadSpeedLimit == "0" { config.SingleThreadSpeedLimit = 0 } else {