From afcdb2a7fa8a66def57b435787fe61f2c7ceec6e Mon Sep 17 00:00:00 2001
From: Hiromu OCHIAI <otiai10@gmail.com>
Date: Tue, 25 Oct 2022 14:35:22 +0900
Subject: [PATCH] Add `WrapReader`, Fix #89

---
 .gitignore            |  3 +++
 all_test.go           | 31 ++++++++++++++++++++++++-------
 copy.go               |  6 +++---
 options.go            | 14 +++++++++-----
 ratelimited_reader.go | 38 --------------------------------------
 test_setup.go         |  1 +
 6 files changed, 40 insertions(+), 53 deletions(-)
 delete mode 100644 ratelimited_reader.go

diff --git a/.gitignore b/.gitignore
index 3b65243..a793485 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,6 @@ coverage.txt
 vendor
 .vagrant
 .idea/
+
+# Test Specific
+test/data/case16/large.file
diff --git a/all_test.go b/all_test.go
index b1d7fb9..5584711 100644
--- a/all_test.go
+++ b/all_test.go
@@ -2,8 +2,8 @@ package copy
 
 import (
 	"errors"
+	"io"
 	"io/ioutil"
-	"math"
 	"os"
 	"runtime"
 	"strings"
@@ -330,9 +330,6 @@ func TestOptions_PreserveOwner(t *testing.T) {
 }
 
 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 {
@@ -340,16 +337,36 @@ func TestOptions_CopyRateLimit(t *testing.T) {
 		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)
-	copyCost := time.Since(start)
+	elasped := 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)
+	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
 }
diff --git a/copy.go b/copy.go
index 2671931..4bdb397 100644
--- a/copy.go
+++ b/copy.go
@@ -83,10 +83,10 @@ 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
-	if opt.CopyRateLimit > 0 {
-		r = NewRateLimitedReader(r, opt.CopyRateLimit)
+
+	if opt.WrapReader != nil {
+		r = opt.WrapReader(s)
 	}
 
 	if opt.CopyBufferSize != 0 {
diff --git a/options.go b/options.go
index 42eea8f..da08377 100644
--- a/options.go
+++ b/options.go
@@ -1,6 +1,9 @@
 package copy
 
-import "os"
+import (
+	"io"
+	"os"
+)
 
 // Options specifies optional actions on copying.
 type Options struct {
@@ -46,9 +49,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
+	// 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
@@ -96,7 +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)
-		CopyRateLimit:     0,                  // Do not specify, use default rate (unlimited)
+		WrapReader:        nil,                // Do not wrap src files, use them as they are.
 		intent: struct {
 			src  string
 			dest string
diff --git a/ratelimited_reader.go b/ratelimited_reader.go
deleted file mode 100644
index f99b690..0000000
--- a/ratelimited_reader.go
+++ /dev/null
@@ -1,38 +0,0 @@
-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
-}
diff --git a/test_setup.go b/test_setup.go
index 0db62fa..f3c2d83 100644
--- a/test_setup.go
+++ b/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)