Improve Copy

Author: nekohasekai

common/bufio/copy.go

@@ -15,7 +15,13 @@ import (
 	"github.com/sagernet/sing/common/task"
 )
 
+const DefaultIncreaseBufferAfter = 512 * 1000
+
 func Copy(destination io.Writer, source io.Reader) (n int64, err error) {
+	return CopyWithIncreateBuffer(destination, source, DefaultIncreaseBufferAfter)
+}
+
+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 +52,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 +65,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 +80,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 +127,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,6 +159,53 @@ 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
+		}
 	}
 }
 

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)
 }