fix: improve multiplexing for mysql protocol

protocol/mysql/outbound.go

@@ -4,8 +4,8 @@ import (
 	"context"
 	"crypto/tls"
 	"net"
-	"os"
 	"sync"
+	"sync/atomic"
 
 	"github.com/sagernet/sing-box/adapter"
 	"github.com/sagernet/sing-box/adapter/outbound"
@@ -39,10 +39,16 @@ type Outbound struct {
 	username       string
 	password       string
 	tlsConfig      *tls.Config
+	maxConnections int
+	nextSession    uint32
 
 	sessionAccess sync.Mutex
+	sessions      []*muxSession
+}
+
+type muxSession struct {
 	session *smux.Session
-	sessionConn   net.Conn
+	conn    net.Conn
 }
 
 func NewOutbound(ctx context.Context, router adapter.Router, logger log.ContextLogger, tag string, options option.MySQLOutboundOptions) (adapter.Outbound, error) {
@@ -59,8 +65,14 @@ func NewOutbound(ctx context.Context, router adapter.Router, logger log.ContextL
 		serverAddr:     options.ServerOptions.Build(),
 		username:       options.Username,
 		password:       options.Password,
+		maxConnections: 1,
 	}
 
+	if options.Multiplex != nil && options.Multiplex.Enabled && options.Multiplex.MaxConnections > 1 {
+		outbound.maxConnections = options.Multiplex.MaxConnections
+	}
+	outbound.sessions = make([]*muxSession, outbound.maxConnections)
+
 	if outbound.serverAddr.Port == 0 {
 		outbound.serverAddr.Port = 3306
 	}
@@ -88,14 +100,8 @@ func NewOutbound(ctx context.Context, router adapter.Router, logger log.ContextL
 	return outbound, nil
 }
 
-func (h *Outbound) getSession() (*smux.Session, error) {
-	h.sessionAccess.Lock()
-	defer h.sessionAccess.Unlock()
-
-	if h.session != nil && !h.session.IsClosed() {
-		return h.session, nil
-	}
-
+func (h *Outbound) createSession() (*muxSession, error) {
+	h.logger.InfoContext(h.ctx, "creating smux session")
 	// Dial TCP connection to server
 	conn, err := h.dialer.DialContext(h.ctx, N.NetworkTCP, h.serverAddr)
 	if err != nil {
@@ -135,67 +141,92 @@ func (h *Outbound) getSession() (*smux.Session, error) {
 		return nil, E.Cause(err, "create mux session")
 	}
 
-	h.session = session
-	h.sessionConn = tlsConn
+	return &muxSession{session: session, conn: tlsConn}, nil
+}
 
-	go func() {
+func (h *Outbound) getSession(index int) (*smux.Session, error) {
+	h.sessionAccess.Lock()
+	defer h.sessionAccess.Unlock()
+
+	entry := h.sessions[index]
+	if entry != nil && !entry.session.IsClosed() {
+		return entry.session, nil
+	}
+	if entry != nil {
+		entry.conn.Close()
+		h.sessions[index] = nil
+	}
+
+	entry, err := h.createSession()
+	if err != nil {
+		return nil, err
+	}
+	h.sessions[index] = entry
+
+	go func(index int, session *smux.Session, conn net.Conn) {
 		// When session is closed, clean up
 		<-session.CloseChan()
 		h.sessionAccess.Lock()
-		if h.session == session {
-			h.session = nil
-			h.sessionConn = nil
+		if current := h.sessions[index]; current != nil && current.session == session {
+			h.sessions[index] = nil
 		}
 		h.sessionAccess.Unlock()
-		tlsConn.Close()
-	}()
+		conn.Close()
+	}(index, entry.session, entry.conn)
 
-	return session, nil
+	return entry.session, nil
+}
+
+func (h *Outbound) invalidateSession(index int, session *smux.Session) {
+	h.sessionAccess.Lock()
+	defer h.sessionAccess.Unlock()
+
+	if current := h.sessions[index]; current != nil && current.session == session {
+		h.sessions[index] = nil
+		current.conn.Close()
+	}
 }
 
 func (h *Outbound) openStream(ctx context.Context, command byte, destination M.Socksaddr) (net.Conn, error) {
-	session, err := h.getSession()
+	_ = ctx
+	start := int(atomic.AddUint32(&h.nextSession, 1)-1) % h.maxConnections
+	var lastErr error
+	for i := 0; i < h.maxConnections; i++ {
+		index := (start + i) % h.maxConnections
+		session, err := h.getSession(index)
 		if err != nil {
-		return nil, err
+			lastErr = err
+			continue
 		}
 
 		stream, err := session.OpenStream()
 		if err != nil {
-		// Session might be stale, try once more with a new session
-		h.sessionAccess.Lock()
-		if h.session == session {
-			h.session = nil
-			if h.sessionConn != nil {
-				h.sessionConn.Close()
-				h.sessionConn = nil
-			}
-		}
-		h.sessionAccess.Unlock()
-
-		session, err = h.getSession()
-		if err != nil {
-			return nil, err
-		}
-		stream, err = session.OpenStream()
-		if err != nil {
-			return nil, E.Cause(err, "open mux stream")
-		}
+			h.invalidateSession(index, session)
+			lastErr = err
+			continue
 		}
 
 		// Write stream header: command + destination
 		_, err = stream.Write([]byte{command})
 		if err != nil {
 			stream.Close()
-		return nil, E.Cause(err, "write stream header command")
+			lastErr = E.Cause(err, "write stream header command")
+			continue
 		}
 		err = M.SocksaddrSerializer.WriteAddrPort(stream, destination)
 		if err != nil {
 			stream.Close()
-		return nil, E.Cause(err, "write stream header destination")
+			lastErr = E.Cause(err, "write stream header destination")
+			continue
 		}
 
 		return stream, nil
 	}
+	if lastErr == nil {
+		lastErr = E.New("open mux stream")
+	}
+	return nil, E.Cause(lastErr, "open mux stream")
+}
 
 func (h *Outbound) DialContext(ctx context.Context, network string, destination M.Socksaddr) (net.Conn, error) {
 	switch N.NetworkName(network) {
@@ -226,13 +257,13 @@ func (h *Outbound) ListenPacket(ctx context.Context, destination M.Socksaddr) (n
 func (h *Outbound) InterfaceUpdated() {
 	h.sessionAccess.Lock()
 	defer h.sessionAccess.Unlock()
-	if h.session != nil {
-		h.session.Close()
-		h.session = nil
+	for i, session := range h.sessions {
+		if session == nil {
+			continue
 		}
-	if h.sessionConn != nil {
-		h.sessionConn.Close()
-		h.sessionConn = nil
+		session.session.Close()
+		session.conn.Close()
+		h.sessions[i] = nil
 	}
 }
 
@@ -240,13 +271,12 @@ func (h *Outbound) Close() error {
 	h.sessionAccess.Lock()
 	defer h.sessionAccess.Unlock()
 	var err error
-	if h.session != nil {
-		err = h.session.Close()
-		h.session = nil
+	for i, session := range h.sessions {
+		if session == nil {
+			continue
 		}
-	if h.sessionConn != nil {
-		common.Close(h.sessionConn)
-		h.sessionConn = nil
+		err = common.Close(session.session, session.conn)
+		h.sessions[i] = nil
 	}
 	return err
 }
@@ -264,5 +294,3 @@ func (c *packetConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 func (c *packetConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 	return c.Conn.Write(p)
 }
-
-var _ = os.ErrInvalid // keep import