瀏覽代碼

force close session manage when main connection is finished

Darien Raymond 8 年之前
父節點
當前提交
2cf809021b
共有 2 個文件被更改,包括 19 次插入2 次删除
  1. 3 0
      app/proxyman/mux/mux.go
  2. 16 2
      app/proxyman/mux/session.go

+ 3 - 0
app/proxyman/mux/mux.go

@@ -121,6 +121,7 @@ func (m *Client) monitor() {
 	for {
 		select {
 		case <-m.ctx.Done():
+			m.sessionManager.Close()
 			m.inboundRay.InboundInput().Close()
 			m.inboundRay.InboundOutput().CloseError()
 			return
@@ -302,6 +303,8 @@ func handle(ctx context.Context, s *Session, output buf.Writer) {
 func (w *ServerWorker) run(ctx context.Context) {
 	input := w.outboundRay.OutboundInput()
 	reader := NewReader(input)
+
+	defer w.sessionManager.Close()
 L:
 	for {
 		select {

+ 16 - 2
app/proxyman/mux/session.go

@@ -74,11 +74,27 @@ func (m *SessionManager) CloseIfNoSession() bool {
 	m.RLock()
 	defer m.RUnlock()
 
+	if m.closed {
+		return true
+	}
+
 	if len(m.sessions) == 0 {
 		return false
 	}
 
 	m.closed = true
+	return true
+}
+
+func (m *SessionManager) Close() {
+	m.RLock()
+	defer m.RUnlock()
+
+	if m.closed {
+		return
+	}
+
+	m.closed = true
 
 	for _, s := range m.sessions {
 		s.input.CloseError()
@@ -86,8 +102,6 @@ func (m *SessionManager) CloseIfNoSession() bool {
 	}
 
 	m.sessions = make(map[uint16]*Session)
-
-	return true
 }
 
 type Session struct {