io: Improve performance of CopyN

Benchmarks:
name          old time/op    new time/op    delta
CopyNSmall-4    5.09µs ± 1%    2.25µs ±86%  -55.91%  (p=0.000 n=11+14)
CopyNLarge-4     114µs ±73%     121µs ±72%     ~     (p=0.701 n=14+14)

name          old alloc/op   new alloc/op   delta
CopyNSmall-4    34.6kB ± 0%     1.9kB ±19%  -94.60%  (p=0.000 n=12+14)
CopyNLarge-4     129kB ± 8%     127kB ±18%   -2.00%  (p=0.007 n=14+14)

name          old allocs/op  new allocs/op  delta
CopyNSmall-4      2.00 ± 0%      1.00 ± 0%  -50.00%  (p=0.000 n=14+14)
CopyNLarge-4      2.00 ± 0%      1.00 ± 0%  -50.00%  (p=0.000 n=14+14)

Benchmark code:
type Buffer struct {
	bytes.Buffer
	io.ReaderFrom
}

func BenchmarkCopyNSmall(b *testing.B) {
	bs := bytes.Repeat([]byte{0}, 1024)
	rd := bytes.NewReader(bs)
	buf := new(Buffer)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		io.CopyN(buf, rd, 512)
		rd.Reset(bs)
	}
}

func BenchmarkCopyNLarge(b *testing.B) {
	bs := bytes.Repeat([]byte{0}, 64*1024)
	rd := bytes.NewReader(bs)
	buf := new(Buffer)
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		io.CopyN(buf, rd, (32*1024)+1)
		rd.Reset(bs)
	}
}

Change-Id: Id8d29e55758452c870cf372db640f07baec05849
Reviewed-on: https://go-review.googlesource.com/60630
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>

diff --git a/src/io/io.go b/src/io/io.go
index 28dab08e46cef187706e24fc976d25c5c04c3392..86710ed6f31d7f8c7902c3975abe23076964f5da 100644 (file)
--- a/src/io/io.go
+++ b/src/io/io.go
@@ -335,7 +335,7 @@ func ReadFull(r Reader, buf []byte) (n int, err error) {
 // If dst implements the ReaderFrom interface,
 // the copy is implemented using it.
 func CopyN(dst Writer, src Reader, n int64) (written int64, err error) {
-       written, err = Copy(dst, LimitReader(src, n))
+       written, err = copyN(dst, src, n)
        if written == n {
                return n, nil
        }
@@ -346,6 +346,55 @@ func CopyN(dst Writer, src Reader, n int64) (written int64, err error) {
        return
 }
 
+// copyN copies n bytes (or until an error) from src to dst.
+// It returns the number of bytes copied and the earliest
+// error encountered while copying.
+//
+// If dst implements the ReaderFrom interface,
+// the copy is implemented using it.
+func copyN(dst Writer, src Reader, n int64) (int64, error) {
+       // If the writer has a ReadFrom method, use it to do the copy.
+       if rt, ok := dst.(ReaderFrom); ok {
+               return rt.ReadFrom(LimitReader(src, n))
+       }
+
+       l := 32 * 1024 // same size as in copyBuffer
+       if n < int64(l) {
+               l = int(n)
+       }
+       buf := make([]byte, l)
+
+       var written int64
+       for n > 0 {
+               if n < int64(len(buf)) {
+                       buf = buf[:n]
+               }
+
+               nr, errR := src.Read(buf)
+               if nr > 0 {
+                       n -= int64(nr)
+                       nw, errW := dst.Write(buf[:nr])
+                       if nw > 0 {
+                               written += int64(nw)
+                       }
+                       if errW != nil {
+                               return written, errW
+                       }
+                       if nr != nw {
+                               return written, ErrShortWrite
+                       }
+               }
+
+               if errR != nil {
+                       if errR != EOF {
+                               return written, errR
+                       }
+                       return written, nil
+               }
+       }
+       return written, nil
+}
+
 // Copy copies from src to dst until either EOF is reached
 // on src or an error occurs. It returns the number of bytes
 // copied and the first error encountered while copying, if any.