MUX: Prevent goroutine leak (#5110)

common/mux/session.go

@@ -12,6 +12,7 @@ import (
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/protocol"
+	"github.com/xtls/xray-core/common/signal/done"
 	"github.com/xtls/xray-core/transport/pipe"
 )
 
@@ -53,7 +54,7 @@ func (m *SessionManager) Count() int {
 func (m *SessionManager) Allocate(Strategy *ClientStrategy) *Session {
 	m.Lock()
 	defer m.Unlock()
-	
+
 	MaxConcurrency := int(Strategy.MaxConcurrency)
 	MaxConnection := uint16(Strategy.MaxConnection)
 
@@ -65,6 +66,7 @@ func (m *SessionManager) Allocate(Strategy *ClientStrategy) *Session {
 	s := &Session{
 		ID:     m.count,
 		parent: m,
+		done:   done.New(),
 	}
 	m.sessions[s.ID] = s
 	return s
@@ -115,7 +117,7 @@ func (m *SessionManager) Get(id uint16) (*Session, bool) {
 	return s, found
 }
 
-func (m *SessionManager) CloseIfNoSession() bool {
+func (m *SessionManager) CloseIfNoSessionAndIdle(checkSize int, checkCount int) bool {
 	m.Lock()
 	defer m.Unlock()
 
@@ -123,11 +125,13 @@ func (m *SessionManager) CloseIfNoSession() bool {
 		return true
 	}
 
-	if len(m.sessions) != 0 {
+	if len(m.sessions) != 0 || checkSize != 0 || checkCount != int(m.count) {
 		return false
 	}
 
 	m.closed = true
+
+	m.sessions = nil
 	return true
 }
 
@@ -157,6 +161,7 @@ type Session struct {
 	ID           uint16
 	transferType protocol.TransferType
 	closed       bool
+	done         *done.Instance
 	XUDP         *XUDP
 }
 
@@ -171,6 +176,9 @@ func (s *Session) Close(locked bool) error {
 		return nil
 	}
 	s.closed = true
+	if s.done != nil {
+		s.done.Close()
+	}
 	if s.XUDP == nil {
 		common.Interrupt(s.input)
 		common.Close(s.output)