Improve Copy

Author: nekohasekai

common/bufio/copy.go

@@ -16,6 +16,10 @@ import (
 )
 
 func Copy(destination io.Writer, source io.Reader) (n int64, err error) {
+	return CopyWithIncreateBuffer(destination, source, 0)
+}
+
+func CopyWithIncreateBuffer(destination io.Writer, source io.Reader, increaseBufferAfter int64) (n int64, err error) {
 	if source == nil {
 		return 0, E.New("nil reader")
 	} else if destination == nil {
@@ -46,10 +50,10 @@ func Copy(destination io.Writer, source io.Reader) (n int64, err error) {
 		}
 		break
 	}
-	return CopyWithCounters(destination, source, originSource, readCounters, writeCounters)
+	return CopyWithCounters(destination, source, originSource, readCounters, writeCounters, increaseBufferAfter)
 }
 
-func CopyWithCounters(destination io.Writer, source io.Reader, originSource io.Reader, readCounters []N.CountFunc, writeCounters []N.CountFunc) (n int64, err error) {
+func CopyWithCounters(destination io.Writer, source io.Reader, originSource io.Reader, readCounters []N.CountFunc, writeCounters []N.CountFunc, increaseBufferAfter int64) (n int64, err error) {
 	srcSyscallConn, srcIsSyscall := source.(syscall.Conn)
 	dstSyscallConn, dstIsSyscall := destination.(syscall.Conn)
 	if srcIsSyscall && dstIsSyscall {
@@ -59,10 +63,10 @@ func CopyWithCounters(destination io.Writer, source io.Reader, originSource io.R
 			return
 		}
 	}
-	return CopyExtended(originSource, NewExtendedWriter(destination), NewExtendedReader(source), readCounters, writeCounters)
+	return CopyExtended(originSource, NewExtendedWriter(destination), NewExtendedReader(source), readCounters, writeCounters, increaseBufferAfter)
 }
 
-func CopyExtended(originSource io.Reader, destination N.ExtendedWriter, source N.ExtendedReader, readCounters []N.CountFunc, writeCounters []N.CountFunc) (n int64, err error) {
+func CopyExtended(originSource io.Reader, destination N.ExtendedWriter, source N.ExtendedReader, readCounters []N.CountFunc, writeCounters []N.CountFunc, increaseBufferAfter int64) (n int64, err error) {
 	frontHeadroom := N.CalculateFrontHeadroom(destination)
 	rearHeadroom := N.CalculateRearHeadroom(destination)
 	readWaiter, isReadWaiter := CreateReadWaiter(source)
@@ -74,13 +78,13 @@ func CopyExtended(originSource io.Reader, destination N.ExtendedWriter, source N
 		})
 		if !needCopy || common.LowMemory {
 			var handled bool
-			handled, n, err = copyWaitWithPool(originSource, destination, readWaiter, readCounters, writeCounters)
+			handled, n, err = copyWaitWithPool(originSource, destination, source, readWaiter, readCounters, writeCounters, increaseBufferAfter)
 			if handled {
 				return
 			}
 		}
 	}
-	return CopyExtendedWithPool(originSource, destination, source, readCounters, writeCounters)
+	return CopyExtendedWithPool(originSource, destination, source, readCounters, writeCounters, increaseBufferAfter)
 }
 
 // Deprecated: not used
@@ -121,7 +125,7 @@ func CopyExtendedBuffer(originSource io.Writer, destination N.ExtendedWriter, so
 	}
 }
 
-func CopyExtendedWithPool(originSource io.Reader, destination N.ExtendedWriter, source N.ExtendedReader, readCounters []N.CountFunc, writeCounters []N.CountFunc) (n int64, err error) {
+func CopyExtendedWithPool(originSource io.Reader, destination N.ExtendedWriter, source N.ExtendedReader, readCounters []N.CountFunc, writeCounters []N.CountFunc, increaseBufferAfter int64) (n int64, err error) {
 	options := N.NewReadWaitOptions(source, destination)
 	var notFirstTime bool
 	for {
@@ -153,14 +157,66 @@ func CopyExtendedWithPool(originSource io.Reader, destination N.ExtendedWriter,
 			counter(int64(dataLen))
 		}
 		notFirstTime = true
+		if increaseBufferAfter > 0 && n >= increaseBufferAfter {
+			return CopyExtendedChanWithPool(destination, source, readCounters, writeCounters, options, n)
+		}
 	}
 }
 
+func CopyExtendedChanWithPool(destination N.ExtendedWriter, source N.ExtendedReader, readCounters []N.CountFunc, writeCounters []N.CountFunc, options N.ReadWaitOptions, inputN int64) (n int64, err error) {
+	n += inputN
+	sendChan := make(chan *buf.Buffer, 2)
+	errChan := make(chan error, 1)
+	go func() {
+		for {
+			buffer := options.NewBufferMax()
+			readErr := source.ReadBuffer(buffer)
+			if readErr != nil {
+				buffer.Release()
+				if errors.Is(readErr, io.EOF) {
+					errChan <- nil
+				} else {
+					errChan <- readErr
+				}
+				return
+			}
+			dataLen := buffer.Len()
+			options.PostReturn(buffer)
+			sendChan <- buffer
+			n += int64(dataLen)
+			for _, counter := range readCounters {
+				counter(int64(dataLen))
+			}
+		}
+	}()
+	for {
+		select {
+		case buffer := <-sendChan:
+			dataLen := buffer.Len()
+			err = destination.WriteBuffer(buffer)
+			if err != nil {
+				buffer.Leak()
+				return
+			}
+			for _, counter := range writeCounters {
+				counter(int64(dataLen))
+			}
+		case err = <-errChan:
+			return
+		}
+	}
+}
+
+const (
+	DefaultIncreaseUploadBufferAfter   = 16 * 1000       // 16KB
+	DefaultIncreaseDownloadBufferAfter = 1 * 1000 * 1000 // 1MB
+)
+
 func CopyConn(ctx context.Context, source net.Conn, destination net.Conn) error {
 	var group task.Group
 	if _, dstDuplex := common.Cast[N.WriteCloser](destination); dstDuplex {
 		group.Append("upload", func(ctx context.Context) error {
-			err := common.Error(Copy(destination, source))
+			err := common.Error(CopyWithIncreateBuffer(destination, source, DefaultIncreaseUploadBufferAfter))
 			if err == nil {
 				N.CloseWrite(destination)
 			} else {
@@ -171,12 +227,12 @@ func CopyConn(ctx context.Context, source net.Conn, destination net.Conn) error
 	} else {
 		group.Append("upload", func(ctx context.Context) error {
 			defer common.Close(destination)
-			return common.Error(Copy(destination, source))
+			return common.Error(CopyWithIncreateBuffer(destination, source, DefaultIncreaseUploadBufferAfter))
 		})
 	}
 	if _, srcDuplex := common.Cast[N.WriteCloser](source); srcDuplex {
 		group.Append("download", func(ctx context.Context) error {
-			err := common.Error(Copy(source, destination))
+			err := common.Error(CopyWithIncreateBuffer(source, destination, DefaultIncreaseDownloadBufferAfter))
 			if err == nil {
 				N.CloseWrite(source)
 			} else {
@@ -187,7 +243,7 @@ func CopyConn(ctx context.Context, source net.Conn, destination net.Conn) error
 	} else {
 		group.Append("download", func(ctx context.Context) error {
 			defer common.Close(source)
-			return common.Error(Copy(source, destination))
+			return common.Error(CopyWithIncreateBuffer(source, destination, DefaultIncreaseDownloadBufferAfter))
 		})
 	}
 	group.Cleanup(func() {

common/bufio/copy_direct.go

@@ -23,14 +23,14 @@ func copyDirect(source syscall.Conn, destination syscall.Conn, readCounters []N.
 	return
 }
 
-func copyWaitWithPool(originSource io.Reader, destination N.ExtendedWriter, source N.ReadWaiter, readCounters []N.CountFunc, writeCounters []N.CountFunc) (handled bool, n int64, err error) {
+func copyWaitWithPool(originSource io.Reader, destination N.ExtendedWriter, source N.ExtendedReader, readWaiter N.ReadWaiter, readCounters []N.CountFunc, writeCounters []N.CountFunc, increaseBufferAfter int64) (handled bool, n int64, err error) {
 	handled = true
 	var (
 		buffer       *buf.Buffer
 		notFirstTime bool
 	)
 	for {
-		buffer, err = source.WaitReadBuffer()
+		buffer, err = readWaiter.WaitReadBuffer()
 		if err != nil {
 			if errors.Is(err, io.EOF) {
 				err = nil
@@ -55,6 +55,10 @@ func copyWaitWithPool(originSource io.Reader, destination N.ExtendedWriter, sour
 			counter(int64(dataLen))
 		}
 		notFirstTime = true
+		if increaseBufferAfter > 0 && n >= increaseBufferAfter {
+			n, err = CopyExtendedChanWithPool(destination, source, readCounters, writeCounters, N.NewReadWaitOptions(source, destination), n)
+			return
+		}
 	}
 }
 

common/network/direct.go

@@ -43,6 +43,11 @@ func (o ReadWaitOptions) NewBuffer() *buf.Buffer {
 	return o.newBuffer(buf.BufferSize, true)
 }
 
+func (o ReadWaitOptions) NewBufferMax() *buf.Buffer {
+	const maxBufferSize = 64<<10 - 1
+	return o.newBuffer(maxBufferSize, true)
+}
+
 func (o ReadWaitOptions) NewPacketBuffer() *buf.Buffer {
 	return o.newBuffer(buf.UDPBufferSize, true)
 }