Merge pull request #91 from otiai10/imsunv/main

Imsunv/main

.gitignore

@@ -4,3 +4,6 @@ coverage.txt
 vendor
 .vagrant
 .idea/
+
+# Test Specific
+test/data/case16/large.file

all_test.go

@@ -2,11 +2,13 @@ package copy
 
 import (
 	"errors"
+	"io"
 	"io/ioutil"
 	"os"
 	"runtime"
 	"strings"
 	"testing"
+	"time"
 
 	. "github.com/otiai10/mint"
 )
@@ -326,3 +328,45 @@ 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) {
+
+	file, err := os.Create("test/data/case16/large.file")
+	if err != nil {
+		t.Errorf("failed to create test file: %v", err)
+		return
+	}
+
+	size := int64(100 * 1024) // 100 KB
+	if err := file.Truncate(size); err != nil {
+		t.Errorf("failed to truncate test file: %v", err)
+		t.SkipNow()
+		return
+	}
+
+	opt := Options{WrapReader: func(src *os.File) io.Reader {
+		return &SleepyReader{src, 1}
+	}}
+
+	start := time.Now()
+	err = Copy("test/data/case16", "test/data.copy/case16", opt)
+	elasped := time.Since(start)
+	Expect(t, err).ToBe(nil)
+	Expect(t, elasped > 5*time.Second).ToBe(true)
+}
+
+type SleepyReader struct {
+	src *os.File
+	sec time.Duration
+}
+
+func (r *SleepyReader) Read(p []byte) (int, error) {
+	n, e := r.src.Read(p)
+	if e != nil && e != io.EOF {
+		return n, e
+	}
+	if n > 0 {
+		time.Sleep(time.Second * r.sec)
+	}
+	return n, e
+}

copy.go

@@ -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.WrapReader != nil {
+		r = opt.WrapReader(s)
+	}
+
 	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
 	}
 

options.go

@@ -1,6 +1,9 @@
 package copy
 
-import "os"
+import (
+	"io"
+	"os"
+)
 
 // Options specifies optional actions on copying.
 type Options struct {
@@ -46,6 +49,11 @@ type Options struct {
 	// See https://golang.org/pkg/io/#CopyBuffer for more information.
 	CopyBufferSize uint
 
+	// If you want to add some limitation on reading src file,
+	// you can wrap the src and provide new reader,
+	// such as `RateLimitReader` in the test case.
+	WrapReader func(src *os.File) io.Reader
+
 	intent struct {
 		src  string
 		dest string
@@ -92,6 +100,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)
+		WrapReader:        nil,                // Do not wrap src files, use them as they are.
 		intent: struct {
 			src  string
 			dest string

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.

test_setup.go

@@ -10,6 +10,7 @@ import (
 )
 
 func setup(m *testing.M) {
+	os.RemoveAll("test/data.copy")
 	os.MkdirAll("test/data.copy", os.ModePerm)
 	os.Symlink("test/data/case01", "test/data/case03/case01")
 	os.Chmod("test/data/case07/dir_0555", 0o555)