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/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..207f01853 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" @@ -40,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") @@ -83,6 +86,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") + + flag.StringVar(&singleThreadSpeedLimit, "singleThreadSpeedLimit", "", "Limit the speed of each thread, 0 is unlimited") + } func printError(url string, err error) { @@ -230,6 +236,17 @@ func main() { config.Cookie = string(data) } } + 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)) {