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>

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.