io: eliminate full copy of copy loop in CopyN

CL 60630 claimed to and did “improve performance of CopyN”
but in doing so introduced a second copy of the I/O copying loop.
This code is subtle and easy to get wrong and the last thing we
need is of two copies that can drift out of sync. Even the newly
introduced copy contains various subtle changes that are not
obviously semantically equivalent to the original. (They probably
are, but it's not obvious.)

Although the CL description does not explain further what the
important optimization was, it appears that the most critical
one was not allocating a 32kB buffer for CopyN(w, r, 512).

This CL deletes the forked copy of copy and instead applies
the buffer size restriction optimization directly to copy itself.

CL 60630 reported:

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)

Starting with that CL as the baseline, this CL does not change a ton:

name          old time/op  new time/op  delta
CopyNSmall-8   370ns ± 1%   411ns ± 1%  +11.18%  (p=0.000 n=16+14)
CopyNLarge-8  18.2µs ± 1%  18.3µs ± 1%   +0.63%  (p=0.000 n=19+20)

It does give up a small amount of the win of 60630 but preserves
the bulk of it, with the benefit that we will not need to debug these
two copies drifting out of sync in the future.

Change-Id: I05b1a5a7115390c5867847cba606b75d513eb2e2
Reviewed-on: https://go-review.googlesource.com/78122
Run-TryBot: Russ Cox <rsc@golang.org>
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>

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 = copyN(dst, src, n)
+	written, err = Copy(dst, LimitReader(src, n))
 	if written == n {
 		return n, nil
 	}
@@ -346,55 +346,6 @@ 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.
@@ -434,8 +385,16 @@ func copyBuffer(dst Writer, src Reader, buf []byte) (written int64, err error) {
 	if rt, ok := dst.(ReaderFrom); ok {
 		return rt.ReadFrom(src)
 	}
+	size := 32 * 1024
+	if l, ok := src.(*LimitedReader); ok && int64(size) > l.N {
+		if l.N < 1 {
+			size = 1
+		} else {
+			size = int(l.N)
+		}
+	}
 	if buf == nil {
-		buf = make([]byte, 32*1024)
+		buf = make([]byte, size)
 	}
 	for {
 		nr, er := src.Read(buf)

src/io/io_test.go

@@ -32,6 +32,21 @@ func TestCopy(t *testing.T) {
 	}
 }
 
+func TestCopyNegative(t *testing.T) {
+	rb := new(Buffer)
+	wb := new(Buffer)
+	rb.WriteString("hello")
+	Copy(wb, &LimitedReader{R: rb, N: -1})
+	if wb.String() != "" {
+		t.Errorf("Copy on LimitedReader with N<0 copied data")
+	}
+
+	CopyN(wb, rb, -1)
+	if wb.String() != "" {
+		t.Errorf("CopyN with N<0 copied data")
+	}
+}
+
 func TestCopyBuffer(t *testing.T) {
 	rb := new(Buffer)
 	wb := new(Buffer)