Fix reactor issues and add context propagation

Author: nekohasekai

common/bufio/fd_poller_windows.go

@@ -120,6 +120,12 @@ func (p *FDPoller) Remove(fd int) {
 	if p.afd != nil {
 		p.afd.Cancel(&entry.ioStatusBlock)
 	}
+
+	if !entry.unpinned {
+		entry.unpinned = true
+		entry.pinner.Unpin()
+	}
+	delete(p.entries, fd)
 }
 
 func (p *FDPoller) wakeup() {

common/bufio/packet_reactor.go

@@ -55,12 +55,13 @@ func (r *PacketReactor) Close() error {
 }
 
 type reactorConnection struct {
-	ctx      context.Context
-	cancel   context.CancelFunc
-	reactor  *PacketReactor
-	onClose  N.CloseHandlerFunc
-	upload   *reactorStream
-	download *reactorStream
+	ctx              context.Context
+	cancel           context.CancelFunc
+	reactor          *PacketReactor
+	onClose          N.CloseHandlerFunc
+	upload           *reactorStream
+	download         *reactorStream
+	stopReactorWatch func() bool
 
 	closeOnce sync.Once
 	done      chan struct{}
@@ -93,6 +94,9 @@ func (r *PacketReactor) Copy(ctx context.Context, source N.PacketConn, destinati
 		onClose: onClose,
 		done:    make(chan struct{}),
 	}
+	conn.stopReactorWatch = common.ContextAfterFunc(r.ctx, func() {
+		conn.closeWithError(r.ctx.Err())
+	})
 
 	conn.upload = r.prepareStream(conn, source, destination)
 	select {
@@ -126,10 +130,12 @@ func (r *PacketReactor) prepareStream(conn *reactorConnection, source N.PacketRe
 		if cachedReader, isCached := source.(N.CachedPacketReader); isCached {
 			packet := cachedReader.ReadCachedPacket()
 			if packet != nil {
-				dataLen := packet.Buffer.Len()
-				err := destination.WritePacket(packet.Buffer, packet.Destination)
+				buffer := packet.Buffer
+				dataLen := buffer.Len()
+				err := destination.WritePacket(buffer, packet.Destination)
 				N.PutPacketBuffer(packet)
 				if err != nil {
+					buffer.Leak()
 					conn.closeWithError(err)
 					return stream
 				}
@@ -151,7 +157,10 @@ func (r *PacketReactor) prepareStream(conn *reactorConnection, source N.PacketRe
 
 	stream.readWaiter, _ = CreatePacketReadWaiter(source)
 	if stream.readWaiter != nil {
-		stream.readWaiter.InitializeReadWaiter(stream.options)
+		needCopy := stream.readWaiter.InitializeReadWaiter(stream.options)
+		if needCopy {
+			stream.readWaiter = nil
+		}
 	}
 
 	if pushable, ok := source.(N.PacketPushable); ok {
@@ -343,7 +352,7 @@ func (s *reactorStream) HandleFDEvent() {
 }
 
 func (s *reactorStream) runLegacyCopy() {
-	_, err := CopyPacket(s.destination, s.source)
+	_, err := CopyPacketWithCounters(s.destination, s.source, s.originSource, s.readCounters, s.writeCounters)
 	s.closeWithError(err)
 }
 
@@ -353,6 +362,12 @@ func (s *reactorStream) closeWithError(err error) {
 
 func (c *reactorConnection) closeWithError(err error) {
 	c.closeOnce.Do(func() {
+		defer close(c.done)
+
+		if c.stopReactorWatch != nil {
+			c.stopReactorWatch()
+		}
+
 		c.err = err
 		c.cancel()
 
@@ -375,8 +390,6 @@ func (c *reactorConnection) closeWithError(err error) {
 		if c.onClose != nil {
 			c.onClose(c.err)
 		}
-
-		close(c.done)
 	})
 }
 

common/bufio/stream_reactor.go

@@ -50,12 +50,13 @@ func (r *StreamReactor) Close() error {
 }
 
 type streamConnection struct {
-	ctx      context.Context
-	cancel   context.CancelFunc
-	reactor  *StreamReactor
-	onClose  N.CloseHandlerFunc
-	upload   *streamDirection
-	download *streamDirection
+	ctx              context.Context
+	cancel           context.CancelFunc
+	reactor          *StreamReactor
+	onClose          N.CloseHandlerFunc
+	upload           *streamDirection
+	download         *streamDirection
+	stopReactorWatch func() bool
 
 	closeOnce sync.Once
 	done      chan struct{}
@@ -96,6 +97,9 @@ func (r *StreamReactor) Copy(ctx context.Context, source net.Conn, destination n
 		onClose: onClose,
 		done:    make(chan struct{}),
 	}
+	conn.stopReactorWatch = common.ContextAfterFunc(r.ctx, func() {
+		conn.closeWithError(r.ctx.Err())
+	})
 
 	conn.upload = r.prepareDirection(conn, source, destination, source, true)
 	select {
@@ -171,7 +175,10 @@ func (r *StreamReactor) prepareDirection(conn *streamConnection, source io.Reade
 
 	direction.readWaiter, _ = CreateReadWaiter(source)
 	if direction.readWaiter != nil {
-		direction.readWaiter.InitializeReadWaiter(direction.options)
+		needCopy := direction.readWaiter.InitializeReadWaiter(direction.options)
+		if needCopy {
+			direction.readWaiter = nil
+		}
 	}
 
 	// Try to get stream pollable for FD-based idle detection
@@ -320,7 +327,7 @@ func (d *streamDirection) HandleFDEvent() {
 }
 
 func (d *streamDirection) runLegacyCopy() {
-	_, err := Copy(d.destination, d.source)
+	_, err := CopyWithCounters(d.destination, d.source, d.originSource, d.readCounters, d.writeCounters, DefaultIncreaseBufferAfter, DefaultBatchSize)
 	d.handleEOFOrError(err)
 }
 
@@ -358,6 +365,12 @@ func (c *streamConnection) checkBothClosed() {
 
 	if uploadClosed && downloadClosed {
 		c.closeOnce.Do(func() {
+			defer close(c.done)
+
+			if c.stopReactorWatch != nil {
+				c.stopReactorWatch()
+			}
+
 			c.cancel()
 			c.removeFromPoller()
 
@@ -367,14 +380,18 @@ func (c *streamConnection) checkBothClosed() {
 			if c.onClose != nil {
 				c.onClose(nil)
 			}
-
-			close(c.done)
 		})
 	}
 }
 
 func (c *streamConnection) closeWithError(err error) {
 	c.closeOnce.Do(func() {
+		defer close(c.done)
+
+		if c.stopReactorWatch != nil {
+			c.stopReactorWatch()
+		}
+
 		c.err = err
 		c.cancel()
 
@@ -397,8 +414,6 @@ func (c *streamConnection) closeWithError(err error) {
 		if c.onClose != nil {
 			c.onClose(c.err)
 		}
-
-		close(c.done)
 	})
 }
 

common/context_afterfunc.go

@@ -0,0 +1,11 @@
+//go:build go1.21
+
+package common
+
+import "context"
+
+// ContextAfterFunc arranges to call f in its own goroutine after ctx is done.
+// Returns a stop function that prevents f from being run.
+func ContextAfterFunc(ctx context.Context, f func()) (stop func() bool) {
+	return context.AfterFunc(ctx, f)
+}

common/context_afterfunc_compat.go

@@ -0,0 +1,41 @@
+//go:build go1.20 && !go1.21
+
+package common
+
+import (
+	"context"
+	"sync"
+)
+
+// ContextAfterFunc arranges to call f in its own goroutine after ctx is done.
+// Returns a stop function that prevents f from being run.
+func ContextAfterFunc(ctx context.Context, f func()) (stop func() bool) {
+	stopCh := make(chan struct{})
+	var once sync.Once
+	stopped := false
+
+	go func() {
+		select {
+		case <-ctx.Done():
+			once.Do(func() {
+				if !stopped {
+					f()
+				}
+			})
+		case <-stopCh:
+		}
+	}()
+
+	return func() bool {
+		select {
+		case <-ctx.Done():
+			return false
+		default:
+			stopped = true
+			once.Do(func() {
+				close(stopCh)
+			})
+			return true
+		}
+	}
+}

common/network/handshake.go

@@ -30,7 +30,7 @@ func ReportHandshakeFailure(reporter any, err error) error {
 			return E.Cause(err, "write handshake failure")
 		})
 	}
-	return nil
+	return err
 }
 
 func CloseOnHandshakeFailure(reporter io.Closer, onClose CloseHandlerFunc, err error) error {

common/udpnat2/conn.go

@@ -161,6 +161,14 @@ func (c *natConn) SetTimeout(timeout time.Duration) bool {
 func (c *natConn) Close() error {
 	c.closeOnce.Do(func() {
 		close(c.doneChan)
+
+		c.queueMutex.Lock()
+		pending := c.dataQueue
+		c.dataQueue = nil
+		c.onDataReady = nil
+		c.queueMutex.Unlock()
+
+		N.ReleaseMultiPacketBuffer(pending)
 		common.Close(c.handler)
 	})
 	return nil