Skip to content

Commit

Permalink
Add copy rete limiter
Browse files Browse the repository at this point in the history
  • Loading branch information
taoyang committed Oct 13, 2022
1 parent d69252f commit 432adff
Show file tree
Hide file tree
Showing 5 changed files with 71 additions and 2 deletions.
10 changes: 8 additions & 2 deletions copy.go
Original file line number Diff line number Diff line change
Expand Up @@ -83,15 +83,21 @@ 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.
// See https://github.com/otiai10/copy/pull/60#discussion_r627320811 for more details.
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
}

Expand Down
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
19 changes: 19 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -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=
5 changes: 5 additions & 0 deletions options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
38 changes: 38 additions & 0 deletions ratelimited_reader.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 432adff

Please sign in to comment.