Repository: otiai10/copy
Pull request #90 (2022-10-18 02:23:49)

Add copy rete limiter 
solve #89 


Patch:
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..f99b690
--- /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
+}


Test patch:
diff --git a/all_test.go b/all_test.go
index b77a182..b1d7fb9 100644
--- a/all_test.go
+++ b/all_test.go
@@ -3,10 +3,12 @@ package copy
 import (
 	"errors"
 	"io/ioutil"
+	"math"
 	"os"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	. "github.com/otiai10/mint"
 )
@@ -326,3 +328,28 @@ func TestOptions_PreserveOwner(t *testing.T) {
 	err := Copy("test/data/case13", "test/data.copy/case13", opt)
 	Expect(t, err).ToBe(nil)
 }
+
+func TestOptions_CopyRateLimit(t *testing.T) {
+	opt := Options{CopyRateLimit: 50} // 50 KB/s
+	size := int64(100 * 1024)         // 100 KB
+	costSec := int64(2)               // 2 seconds
+
+	file, err := os.Create("test/data/case16/large.file")
+	if err != nil {
+		t.Errorf("failed to create test file: %v", err)
+		return
+	}
+
+	if err := file.Truncate(size); err != nil {
+		t.Errorf("failed to truncate test file: %v", err)
+		t.SkipNow()
+		return
+	}
+
+	start := time.Now()
+	err = Copy("test/data/case16", "test/data.copy/case16", opt)
+	copyCost := time.Since(start)
+	Expect(t, err).ToBe(nil)
+	t.Log("copy cost", copyCost)
+	Expect(t, int64(math.Floor(copyCost.Seconds()+0.5)) == costSec).ToBe(true)
+}
diff --git a/test/data/case16/README.md b/test/data/case16/README.md
new file mode 100644
index 0000000..29064a8
--- /dev/null
+++ b/test/data/case16/README.md
@@ -0,0 +1,5 @@
+So if you wanted to limit copy rate with a KB/second value, you could add options like this:
+```go
+opt.CopyRateLimit = 100
+```
+The default value is 0, which means no limit.
\ No newline at end of file