Merge ThreadSafeReader into ReadWaiter interface

Author: nekohasekai

common/buf/buffer.go

@@ -8,14 +8,16 @@ import (
 	"sync/atomic"
 
 	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/debug"
 	E "github.com/sagernet/sing/common/exceptions"
+	F "github.com/sagernet/sing/common/format"
 )
 
 type Buffer struct {
 	data    []byte
 	start   int
 	end     int
-	refs    int32
+	refs    atomic.Int32
 	managed bool
 	closed  bool
 }
@@ -281,24 +283,40 @@ func (b *Buffer) FullReset() {
 }
 
 func (b *Buffer) IncRef() {
-	atomic.AddInt32(&b.refs, 1)
+	b.refs.Add(1)
 }
 
 func (b *Buffer) DecRef() {
-	atomic.AddInt32(&b.refs, -1)
+	b.refs.Add(-1)
 }
 
 func (b *Buffer) Release() {
 	if b == nil || b.closed || !b.managed {
 		return
 	}
-	if atomic.LoadInt32(&b.refs) > 0 {
+	if b.refs.Load() > 0 {
 		return
 	}
 	common.Must(Put(b.data))
 	*b = Buffer{closed: true}
 }
 
+func (b *Buffer) Leak() {
+	if debug.Enabled {
+		if b == nil || b.closed || !b.managed {
+			return
+		}
+		refs := b.refs.Load()
+		if refs == 0 {
+			panic("leaking buffer")
+		} else {
+			panic(F.ToString("leaking buffer with ", refs, " references"))
+		}
+	} else {
+		b.Release()
+	}
+}
+
 func (b *Buffer) Cut(start int, end int) *Buffer {
 	b.start += start
 	b.end = len(b.data) - end

common/bufio/bind_wait.go

@@ -12,8 +12,8 @@ type BindPacketReadWaiter struct {
 	readWaiter N.PacketReadWaiter
 }
 
-func (w *BindPacketReadWaiter) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
-	w.readWaiter.InitializeReadWaiter(newBuffer)
+func (w *BindPacketReadWaiter) InitializeReadWaiter(options *N.ReadWaitOptions) (needCopy bool) {
+	return w.readWaiter.InitializeReadWaiter(options)
 }
 
 func (w *BindPacketReadWaiter) WaitReadBuffer() (buffer *buf.Buffer, err error) {
@@ -28,8 +28,8 @@ type UnbindPacketReadWaiter struct {
 	addr       M.Socksaddr
 }
 
-func (w *UnbindPacketReadWaiter) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
-	w.readWaiter.InitializeReadWaiter(newBuffer)
+func (w *UnbindPacketReadWaiter) InitializeReadWaiter(options *N.ReadWaitOptions) (needCopy bool) {
+	return w.readWaiter.InitializeReadWaiter(options)
 }
 
 func (w *UnbindPacketReadWaiter) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {

common/bufio/copy.go

@@ -57,19 +57,21 @@ func Copy(destination io.Writer, source io.Reader) (n int64, err error) {
 }
 
 func CopyExtended(originSource io.Reader, destination N.ExtendedWriter, source N.ExtendedReader, readCounters []N.CountFunc, writeCounters []N.CountFunc) (n int64, err error) {
-	safeSrc := N.IsSafeReader(source)
-	headroom := N.CalculateFrontHeadroom(destination) + N.CalculateRearHeadroom(destination)
-	if safeSrc != nil {
-		if headroom == 0 {
-			return CopyExtendedWithSrcBuffer(originSource, destination, safeSrc, readCounters, writeCounters)
-		}
-	}
+	frontHeadroom := N.CalculateFrontHeadroom(destination)
+	rearHeadroom := N.CalculateRearHeadroom(destination)
 	readWaiter, isReadWaiter := CreateReadWaiter(source)
 	if isReadWaiter {
-		var handled bool
-		handled, n, err = copyWaitWithPool(originSource, destination, readWaiter, readCounters, writeCounters)
-		if handled {
-			return
+		needCopy := readWaiter.InitializeReadWaiter(&N.ReadWaitOptions{
+			FrontHeadroom: frontHeadroom,
+			RearHeadroom:  rearHeadroom,
+			MTU:           N.CalculateMTU(source, destination),
+		})
+		if !needCopy || common.LowMemory {
+			var handled bool
+			handled, n, err = copyWaitWithPool(originSource, destination, readWaiter, readCounters, writeCounters)
+			if handled {
+				return
+			}
 		}
 	}
 	return CopyExtendedWithPool(originSource, destination, source, readCounters, writeCounters)
@@ -113,6 +115,7 @@ func CopyExtendedBuffer(originSource io.Writer, destination N.ExtendedWriter, so
 	}
 }
 
+// Deprecated: Use ReadWaiter interface instead.
 func CopyExtendedWithSrcBuffer(originSource io.Reader, destination N.ExtendedWriter, source N.ThreadSafeReader, readCounters []N.CountFunc, writeCounters []N.CountFunc) (n int64, err error) {
 	var notFirstTime bool
 	for {
@@ -128,7 +131,7 @@ func CopyExtendedWithSrcBuffer(originSource io.Reader, destination N.ExtendedWri
 		dataLen := buffer.Len()
 		err = destination.WriteBuffer(buffer)
 		if err != nil {
-			buffer.Release()
+			buffer.Leak()
 			if !notFirstTime {
 				err = N.ReportHandshakeFailure(originSource, err)
 			}
@@ -173,7 +176,7 @@ func CopyExtendedWithPool(originSource io.Reader, destination N.ExtendedWriter,
 		buffer.Resize(readBuffer.Start(), dataLen)
 		err = destination.WriteBuffer(buffer)
 		if err != nil {
-			buffer.Release()
+			buffer.Leak()
 			if !notFirstTime {
 				err = N.ReportHandshakeFailure(originSource, err)
 			}
@@ -256,35 +259,33 @@ func CopyPacket(destinationConn N.PacketWriter, source N.PacketReader) (n int64,
 			return
 		}
 	}
-	safeSrc := N.IsSafePacketReader(source)
 	frontHeadroom := N.CalculateFrontHeadroom(destinationConn)
 	rearHeadroom := N.CalculateRearHeadroom(destinationConn)
-	headroom := frontHeadroom + rearHeadroom
-	if safeSrc != nil {
-		if headroom == 0 {
-			var copyN int64
-			copyN, err = CopyPacketWithSrcBuffer(originSource, destinationConn, safeSrc, readCounters, writeCounters, n > 0)
-			n += copyN
-			return
-		}
-	}
 	var (
 		handled bool
 		copeN   int64
 	)
 	readWaiter, isReadWaiter := CreatePacketReadWaiter(source)
 	if isReadWaiter {
-		handled, copeN, err = copyPacketWaitWithPool(originSource, destinationConn, readWaiter, readCounters, writeCounters, n > 0)
-		if handled {
-			n += copeN
-			return
+		needCopy := readWaiter.InitializeReadWaiter(&N.ReadWaitOptions{
+			FrontHeadroom: frontHeadroom,
+			RearHeadroom:  rearHeadroom,
+			MTU:           N.CalculateMTU(source, destinationConn),
+		})
+		if !needCopy || common.LowMemory {
+			handled, copeN, err = copyPacketWaitWithPool(originSource, destinationConn, readWaiter, readCounters, writeCounters, n > 0)
+			if handled {
+				n += copeN
+				return
+			}
 		}
 	}
 	copeN, err = CopyPacketWithPool(originSource, destinationConn, source, readCounters, writeCounters, n > 0)
 	n += copeN
 	return
 }
 
+// Deprecated: Use PacketReadWaiter interface instead.
 func CopyPacketWithSrcBuffer(originSource N.PacketReader, destinationConn N.PacketWriter, source N.ThreadSafePacketReader, readCounters []N.CountFunc, writeCounters []N.CountFunc, notFirstTime bool) (n int64, err error) {
 	var buffer *buf.Buffer
 	var destination M.Socksaddr
@@ -302,7 +303,7 @@ func CopyPacketWithSrcBuffer(originSource N.PacketReader, destinationConn N.Pack
 		}
 		err = destinationConn.WritePacket(buffer, destination)
 		if err != nil {
-			buffer.Release()
+			buffer.Leak()
 			if !notFirstTime {
 				err = N.ReportHandshakeFailure(originSource, err)
 			}
@@ -343,7 +344,7 @@ func CopyPacketWithPool(originSource N.PacketReader, destinationConn N.PacketWri
 		buffer.Resize(readBuffer.Start(), dataLen)
 		err = destinationConn.WritePacket(buffer, destination)
 		if err != nil {
-			buffer.Release()
+			buffer.Leak()
 			if !notFirstTime {
 				err = N.ReportHandshakeFailure(originSource, err)
 			}
@@ -379,7 +380,7 @@ func WritePacketWithPool(originSource N.PacketReader, destinationConn N.PacketWr
 		buffer.Resize(readBuffer.Start(), dataLen)
 		err = destinationConn.WritePacket(buffer, packetBuffer.Destination)
 		if err != nil {
-			buffer.Release()
+			buffer.Leak()
 			if !notFirstTime {
 				err = N.ReportHandshakeFailure(originSource, err)
 			}