Improve packet reactor

Author: nekohasekai

common/bufio/channel_poller.go

@@ -24,20 +24,18 @@ const (
 )
 
 type PacketReactor struct {
-	ctx           context.Context
-	cancel        context.CancelFunc
-	channelPoller *ChannelPoller
-	fdPoller      *FDPoller
-	fdPollerOnce  sync.Once
-	fdPollerErr   error
+	ctx          context.Context
+	cancel       context.CancelFunc
+	fdPoller     *FDPoller
+	fdPollerOnce sync.Once
+	fdPollerErr  error
 }
 
 func NewPacketReactor(ctx context.Context) *PacketReactor {
 	ctx, cancel := context.WithCancel(ctx)
 	return &PacketReactor{
-		ctx:           ctx,
-		cancel:        cancel,
-		channelPoller: NewChannelPoller(ctx),
+		ctx:    ctx,
+		cancel: cancel,
 	}
 }
 
@@ -50,14 +48,10 @@ func (r *PacketReactor) getFDPoller() (*FDPoller, error) {
 
 func (r *PacketReactor) Close() error {
 	r.cancel()
-	var errs []error
-	if r.channelPoller != nil {
-		errs = append(errs, r.channelPoller.Close())
-	}
 	if r.fdPoller != nil {
-		errs = append(errs, r.fdPoller.Close())
+		return r.fdPoller.Close()
 	}
-	return E.Errors(errs...)
+	return nil
 }
 
 type reactorConnection struct {
@@ -80,6 +74,7 @@ type reactorStream struct {
 	destination  N.PacketWriter
 	originSource N.PacketReader
 
+	pushable      N.PacketPushable
 	pollable      N.PacketPollable
 	options       N.ReadWaitOptions
 	readWaiter    N.PacketReadWaiter
@@ -159,7 +154,9 @@ func (r *PacketReactor) prepareStream(conn *reactorConnection, source N.PacketRe
 		stream.readWaiter.InitializeReadWaiter(stream.options)
 	}
 
-	if pollable, ok := source.(N.PacketPollable); ok {
+	if pushable, ok := source.(N.PacketPushable); ok {
+		stream.pushable = pushable
+	} else if pollable, ok := source.(N.PacketPollable); ok {
 		stream.pollable = pollable
 	} else if creator, ok := source.(N.PacketPollableCreator); ok {
 		stream.pollable, _ = creator.CreatePacketPollable()
@@ -169,25 +166,32 @@ func (r *PacketReactor) prepareStream(conn *reactorConnection, source N.PacketRe
 }
 
 func (r *PacketReactor) registerStream(stream *reactorStream) {
+	if stream.pushable != nil {
+		stream.pushable.SetOnDataReady(func() {
+			if stream.state.CompareAndSwap(stateIdle, stateActive) {
+				go stream.runActiveLoop(nil)
+			}
+		})
+		if stream.pushable.HasPendingData() {
+			if stream.state.CompareAndSwap(stateIdle, stateActive) {
+				go stream.runActiveLoop(nil)
+			}
+		}
+		return
+	}
+
 	if stream.pollable == nil {
 		go stream.runLegacyCopy()
 		return
 	}
 
-	switch stream.pollable.PollMode() {
-	case N.PacketPollModeChannel:
-		r.channelPoller.Add(stream, stream.pollable.PacketChannel())
-	case N.PacketPollModeFD:
-		fdPoller, err := r.getFDPoller()
-		if err != nil {
-			go stream.runLegacyCopy()
-			return
-		}
-		err = fdPoller.Add(stream, stream.pollable.FD())
-		if err != nil {
-			go stream.runLegacyCopy()
-		}
-	default:
+	fdPoller, err := r.getFDPoller()
+	if err != nil {
+		go stream.runLegacyCopy()
+		return
+	}
+	err = fdPoller.Add(stream, stream.pollable.FD())
+	if err != nil {
 		go stream.runLegacyCopy()
 	}
 }
@@ -259,9 +263,18 @@ func (s *reactorStream) runActiveLoop(firstPacket *N.PacketBuffer) {
 				if setter, ok := s.source.(interface{ SetReadDeadline(time.Time) error }); ok {
 					setter.SetReadDeadline(time.Time{})
 				}
-				if s.state.CompareAndSwap(stateActive, stateIdle) {
-					s.returnToPool()
+				if !s.state.CompareAndSwap(stateActive, stateIdle) {
+					return
+				}
+				if s.pushable != nil {
+					if s.pushable.HasPendingData() {
+						if s.state.CompareAndSwap(stateIdle, stateActive) {
+							continue
+						}
+					}
+					return
 				}
+				s.returnToPool()
 				return
 			}
 			if !notFirstTime {
@@ -310,30 +323,18 @@ func (s *reactorStream) returnToPool() {
 		return
 	}
 
-	if s.pollable == nil {
+	if s.pollable == nil || s.connection.reactor.fdPoller == nil {
 		return
 	}
 
-	switch s.pollable.PollMode() {
-	case N.PacketPollModeChannel:
-		channel := s.pollable.PacketChannel()
-		s.connection.reactor.channelPoller.Add(s, channel)
-		if s.state.Load() != stateIdle {
-			s.connection.reactor.channelPoller.Remove(channel)
-		}
-	case N.PacketPollModeFD:
-		if s.connection.reactor.fdPoller == nil {
-			return
-		}
-		fd := s.pollable.FD()
-		err := s.connection.reactor.fdPoller.Add(s, fd)
-		if err != nil {
-			s.closeWithError(err)
-			return
-		}
-		if s.state.Load() != stateIdle {
-			s.connection.reactor.fdPoller.Remove(fd)
-		}
+	fd := s.pollable.FD()
+	err := s.connection.reactor.fdPoller.Add(s, fd)
+	if err != nil {
+		s.closeWithError(err)
+		return
+	}
+	if s.state.Load() != stateIdle {
+		s.connection.reactor.fdPoller.Remove(fd)
 	}
 }
 
@@ -385,15 +386,8 @@ func (c *reactorConnection) removeFromPollers() {
 }
 
 func (c *reactorConnection) removeStreamFromPoller(stream *reactorStream) {
-	if stream == nil || stream.pollable == nil {
+	if stream == nil || stream.pollable == nil || c.reactor.fdPoller == nil {
 		return
 	}
-	switch stream.pollable.PollMode() {
-	case N.PacketPollModeChannel:
-		c.reactor.channelPoller.Remove(stream.pollable.PacketChannel())
-	case N.PacketPollModeFD:
-		if c.reactor.fdPoller != nil {
-			c.reactor.fdPoller.Remove(stream.pollable.FD())
-		}
-	}
+	c.reactor.fdPoller.Remove(stream.pollable.FD())
 }

common/bufio/packet_reactor.go

@@ -103,10 +103,6 @@ func (p *testPacketPipe) SetWriteDeadline(t time.Time) error {
 	return nil
 }
 
-func (p *testPacketPipe) PacketChannel() <-chan *N.PacketBuffer {
-	return p.inChan
-}
-
 func (p *testPacketPipe) send(data []byte, destination M.Socksaddr) {
 	packet := N.NewPacketBuffer()
 	newBuf := buf.NewSize(len(data))
@@ -255,10 +251,6 @@ func (c *channelPacketConn) SetReadDeadline(t time.Time) error {
 	return nil
 }
 
-func (c *channelPacketConn) PacketChannel() <-chan *N.PacketBuffer {
-	return c.packetChan
-}
-
 func (c *channelPacketConn) Close() error {
 	c.closeOnce.Do(func() {
 		close(c.done)
@@ -551,7 +543,7 @@ func TestBatchCopy_FDPoller_DataIntegrity(t *testing.T) {
 	assert.Equal(t, serverPair.sendHash, clientPair.recvHash, "server->client mismatch")
 }
 
-func TestBatchCopy_ChannelPoller_DataIntegrity(t *testing.T) {
+func TestBatchCopy_LegacyChannel_DataIntegrity(t *testing.T) {
 	t.Parallel()
 
 	clientConn, err := net.ListenPacket("udp", "127.0.0.1:0")

common/bufio/packet_reactor_test.go

@@ -1,20 +1,19 @@
 package network
 
-type PacketPollMode int
-
-const (
-	PacketPollModeChannel PacketPollMode = iota
-	PacketPollModeFD
-)
+// PacketPushable represents a packet source that receives pushed data
+// from external code and notifies reactor via callback.
+type PacketPushable interface {
+	SetOnDataReady(callback func())
+	HasPendingData() bool
+}
 
-// PacketPollable provides polling support for packet connections
+// PacketPollable provides FD-based polling for packet connections.
+// Mirrors StreamPollable for consistency.
 type PacketPollable interface {
-	PollMode() PacketPollMode
-	PacketChannel() <-chan *PacketBuffer
 	FD() int
 }
 
-// PacketPollableCreator creates a PacketPollable dynamically
+// PacketPollableCreator creates a PacketPollable dynamically.
 type PacketPollableCreator interface {
 	CreatePacketPollable() (PacketPollable, bool)
 }