Improve half-close handling and queue backpressure

internal/client/sender_workers.go

@@ -151,7 +151,7 @@ func (c *Client) buildPollBatch(connections []*SOCKSConnection) (protocol.Batch,
 	now := time.Now()
 	nowUnixMS := now.UnixMilli()
 	lastUnixMS := c.lastPollUnixMS.Load()
-	minInterval := time.Duration(c.cfg.WorkerPollIntervalMS) * time.Millisecond
+	minInterval := time.Duration(c.cfg.IdlePollIntervalMS) * time.Millisecond
 	if lastUnixMS > 0 && nowUnixMS-lastUnixMS < minInterval.Milliseconds() {
 		return protocol.Batch{}, false
 	}
@@ -304,7 +304,26 @@ func (c *Client) applyResponsePacket(packet protocol.Packet) error {
 		socksConn.LastActivityAt = time.Now()
 		return socksConn.WriteToLocal(packet.Payload)
 
-	case protocol.PacketTypeSOCKSCloseRead, protocol.PacketTypeSOCKSCloseWrite, protocol.PacketTypeSOCKSRST:
+	case protocol.PacketTypeSOCKSCloseRead:
+		_ = socksConn.AckPacket(packet)
+		socksConn.LastActivityAt = time.Now()
+		if err := socksConn.CloseLocalWrite(); err != nil {
+			return err
+		}
+		if socksConn.BothLocalSidesClosed() {
+			return socksConn.CloseLocal()
+		}
+		return nil
+
+	case protocol.PacketTypeSOCKSCloseWrite:
+		_ = socksConn.AckPacket(packet)
+		socksConn.LastActivityAt = time.Now()
+		if socksConn.BothLocalSidesClosed() {
+			return socksConn.CloseLocal()
+		}
+		return nil
+
+	case protocol.PacketTypeSOCKSRST:
 		_ = socksConn.AckPacket(packet)
 		socksConn.LastActivityAt = time.Now()
 		return socksConn.CloseLocal()

internal/client/session.go

@@ -7,6 +7,7 @@
 package client
 
 import (
+	"context"
 	"encoding/hex"
 	"net"
 	"sync"
@@ -38,6 +39,11 @@ type SOCKSConnection struct {
 
 	LocalConn      net.Conn
 	localWriteMu   sync.Mutex
+	localCloseMu   sync.Mutex
+	localReadEOF   bool
+	localWriteEOF  bool
+	closedC        chan struct{}
+	closeOnce      sync.Once
 	connectResultC chan error
 	queueMu        sync.Mutex
 	OutboundQueue  []*SOCKSOutboundQueueItem
@@ -74,6 +80,7 @@ func (s *SOCKSConnectionStore) New(clientSessionKey string, clientAddress string
 		CreatedAt:        now,
 		LastActivityAt:   now,
 		ClientAddress:    clientAddress,
+		closedC:          make(chan struct{}),
 		connectResultC:   make(chan error, 1),
 		InFlight:         make(map[string]*SOCKSOutboundQueueItem),
 	}
@@ -112,13 +119,65 @@ func (s *SOCKSConnection) WriteToLocal(payload []byte) error {
 }
 
 func (s *SOCKSConnection) CloseLocal() error {
-	s.localWriteMu.Lock()
-	defer s.localWriteMu.Unlock()
+	var err error
+	s.closeOnce.Do(func() {
+		s.localWriteMu.Lock()
+		defer s.localWriteMu.Unlock()
+		if s.LocalConn != nil {
+			err = s.LocalConn.Close()
+		}
+		close(s.closedC)
+	})
+	return err
+}
 
-	if s.LocalConn == nil {
+func (s *SOCKSConnection) CloseLocalWrite() error {
+	s.localCloseMu.Lock()
+	defer s.localCloseMu.Unlock()
+
+	if s.localWriteEOF {
 		return nil
 	}
-	return s.LocalConn.Close()
+	s.localWriteEOF = true
+
+	if tcpConn, ok := s.LocalConn.(*net.TCPConn); ok {
+		return tcpConn.CloseWrite()
+	}
+	return s.CloseLocal()
+}
+
+func (s *SOCKSConnection) CloseLocalRead() error {
+	s.localCloseMu.Lock()
+	defer s.localCloseMu.Unlock()
+
+	if s.localReadEOF {
+		return nil
+	}
+	s.localReadEOF = true
+
+	if tcpConn, ok := s.LocalConn.(*net.TCPConn); ok {
+		return tcpConn.CloseRead()
+	}
+	return nil
+}
+
+func (s *SOCKSConnection) MarkLocalReadEOF() {
+	s.localCloseMu.Lock()
+	s.localReadEOF = true
+	s.localCloseMu.Unlock()
+}
+
+func (s *SOCKSConnection) BothLocalSidesClosed() bool {
+	s.localCloseMu.Lock()
+	defer s.localCloseMu.Unlock()
+	return s.localReadEOF && s.localWriteEOF
+}
+
+func (s *SOCKSConnection) WaitUntilClosed(ctx context.Context) {
+	select {
+	case <-ctx.Done():
+	case <-s.closedC:
+	}
 }
 
 func (s *SOCKSConnectionStore) Get(id uint64) *SOCKSConnection {

internal/client/socks5.go

@@ -316,7 +316,9 @@ func (c *Client) captureInitialPayload(ctx context.Context, conn net.Conn, socks
 
 		if err != nil {
 			if errors.Is(err, io.EOF) {
+				socksConn.MarkLocalReadEOF()
 				_ = socksConn.EnqueuePacket(socksConn.BuildSOCKSCloseWritePacket())
+				socksConn.WaitUntilClosed(ctx)
 				return nil
 			}
 			if ne, ok := err.(net.Error); ok && ne.Timeout() {