From 432adff346e6d142e623dc0752c9203b138d962f Mon Sep 17 00:00:00 2001 From: taoyang Date: Thu, 13 Oct 2022 16:42:49 +0800 Subject: [PATCH] Add copy rete limiter --- copy.go | 10 ++++++++-- go.mod | 1 + go.sum | 19 +++++++++++++++++++ options.go | 5 +++++ ratelimited_reader.go | 38 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 71 insertions(+), 2 deletions(-) create mode 100644 ratelimited_reader.go diff --git a/copy.go b/copy.go index 77eaf3c..2671931 100644 --- a/copy.go +++ b/copy.go @@ -83,7 +83,12 @@ func fcopy(src, dest string, info os.FileInfo, opt Options) (err error) { var buf []byte = nil var w io.Writer = f - // var r io.Reader = s + + var r io.Reader = s + if opt.CopyRateLimit > 0 { + r = NewRateLimitedReader(r, opt.CopyRateLimit) + } + if opt.CopyBufferSize != 0 { buf = make([]byte, opt.CopyBufferSize) // Disable using `ReadFrom` by io.CopyBuffer. @@ -91,7 +96,8 @@ func fcopy(src, dest string, info os.FileInfo, opt Options) (err error) { w = struct{ io.Writer }{f} // r = struct{ io.Reader }{s} } - if _, err = io.CopyBuffer(w, s, buf); err != nil { + + if _, err = io.CopyBuffer(w, r, buf); err != nil { return err } diff --git a/go.mod b/go.mod index f1c6bc2..eb9fad4 100644 --- a/go.mod +++ b/go.mod @@ -4,5 +4,6 @@ go 1.14 require ( github.com/otiai10/mint v1.4.0 + go.uber.org/ratelimit v0.2.0 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 ) diff --git a/go.sum b/go.sum index ee2753e..4deec0c 100644 --- a/go.sum +++ b/go.sum @@ -1,8 +1,27 @@ +github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 h1:MzBOUgng9orim59UnfUTLRjMpd09C5uEVQ6RPGeCaVI= +github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129/go.mod h1:rFgpPQZYZ8vdbc+48xibu8ALc3yeyd64IhHS+PU6Yyg= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/otiai10/curr v0.0.0-20150429015615-9b4961190c95/go.mod h1:9qAhocn7zKJG+0mI8eUu6xqkFDYS2kb2saOteoSB3cE= github.com/otiai10/curr v1.0.0 h1:TJIWdbX0B+kpNagQrjgq8bCMrbhiuX73M2XwgtDMoOI= github.com/otiai10/curr v1.0.0/go.mod h1:LskTG5wDwr8Rs+nNQ+1LlxRjAtTZZjtJW4rMXl6j4vs= github.com/otiai10/mint v1.3.0/go.mod h1:F5AjcsTsWUqX+Na9fpHb52P8pcRX2CI6A3ctIT91xUo= github.com/otiai10/mint v1.4.0 h1:umwcf7gbpEwf7WFzqmWwSv0CzbeMsae2u9ZvpP8j2q4= github.com/otiai10/mint v1.4.0/go.mod h1:gifjb2MYOoULtKLqUAEILUG/9KONW6f7YsJ6vQLTlFI= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw= +go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= +go.uber.org/ratelimit v0.2.0 h1:UQE2Bgi7p2B85uP5dC2bbRtig0C+OeNRnNEafLjsLPA= +go.uber.org/ratelimit v0.2.0/go.mod h1:YYBV4e4naJvhpitQrWJu1vCpgB7CboMe0qhltKt6mUg= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/options.go b/options.go index 0045813..42eea8f 100644 --- a/options.go +++ b/options.go @@ -46,6 +46,10 @@ type Options struct { // See https://golang.org/pkg/io/#CopyBuffer for more information. CopyBufferSize uint + // Limit the rate of copying files with n KB per second. + // If zero, will not limit the copy operation. + CopyRateLimit int64 + intent struct { src string dest string @@ -92,6 +96,7 @@ func getDefaultOptions(src, dest string) Options { Sync: false, // Do not sync PreserveTimes: false, // Do not preserve the modification time CopyBufferSize: 0, // Do not specify, use default bufsize (32*1024) + CopyRateLimit: 0, // Do not specify, use default rate (unlimited) intent: struct { src string dest string diff --git a/ratelimited_reader.go b/ratelimited_reader.go new file mode 100644 index 0000000..36be9ad --- /dev/null +++ b/ratelimited_reader.go @@ -0,0 +1,38 @@ +package copy + +import ( + "io" + + "go.uber.org/ratelimit" +) + +type RateLimitedReader struct { + src io.Reader + limiter ratelimit.Limiter +} + +// NewRateLimitedReader +// n means the number of Kb to be read per second +func NewRateLimitedReader(src io.Reader, n int64) io.Reader { + return &RateLimitedReader{ + src: src, + limiter: ratelimit.New(int(n)), + } +} + +func (lr *RateLimitedReader) Read(p []byte) (n int, err error) { + n, e := lr.src.Read(p) + if e != nil && e != io.EOF { + return n, e + } + if n > 0 { + nkb := n / 1024 + if nkb == 0 { + nkb = 1 + } + for i := 0; i < nkb; i++ { + lr.limiter.Take() + } + } + return n, e +}