refactor: enhance context management by adding drain func and unified channel initialization

Author: yosebyte

internal/client.go

@@ -30,6 +30,7 @@ func NewClient(parsedURL *url.URL, logger *logs.Logger) *Client {
 			logger:     logger,
 			semaphore:  make(chan struct{}, semaphoreLimit),
 			signalChan: make(chan string, semaphoreLimit),
+			errChan:    make(chan error, 3),
 		},
 		tunnelName: parsedURL.Hostname(),
 	}
@@ -73,8 +74,8 @@ func (c *Client) Run() {
 
 // start 启动客户端服务
 func (c *Client) start() error {
-	// 初始化基本信息
-	c.initBackground()
+	// 初始化上下文
+	c.initContext()
 
 	// 通过是否监听成功判断单端转发或双端握手
 	if err := c.initTunnelListener(); err == nil {

internal/common.go

@@ -167,13 +167,12 @@ func (c *Common) getAddress(parsedURL *url.URL) {
 	}
 }
 
-// initBackground 初始化基本信息
-func (c *Common) initBackground() {
+// initContext 初始化上下文
+func (c *Common) initContext() {
 	if c.cancel != nil {
 		c.cancel()
 	}
 	c.ctx, c.cancel = context.WithCancel(context.Background())
-	c.errChan = make(chan error, 3)
 }
 
 // initTargetListener 初始化目标监听器
@@ -226,6 +225,17 @@ func (c *Common) initTunnelListener() error {
 	return nil
 }
 
+// drain 清空通道中的所有元素
+func drain[T any](ch <-chan T) {
+	for {
+		select {
+		case <-ch:
+		default:
+			return
+		}
+	}
+}
+
 // stop 共用停止服务
 func (c *Common) stop() {
 	// 取消上下文
@@ -285,14 +295,10 @@ func (c *Common) stop() {
 		c.logger.Debug("Tunnel listener closed: %v", c.tunnelListener.Addr())
 	}
 
-	// 清空信号通道
-	for {
-		select {
-		case <-c.signalChan:
-		default:
-			return
-		}
-	}
+	// 清空通道
+	drain(c.semaphore)
+	drain(c.signalChan)
+	drain(c.errChan)
 }
 
 // shutdown 共用优雅关闭

internal/server.go

@@ -34,6 +34,7 @@ func NewServer(parsedURL *url.URL, tlsCode string, tlsConfig *tls.Config, logger
 			logger:     logger,
 			semaphore:  make(chan struct{}, semaphoreLimit),
 			signalChan: make(chan string, semaphoreLimit),
+			errChan:    make(chan error, 3),
 		},
 		tlsConfig: tlsConfig,
 	}
@@ -77,8 +78,8 @@ func (s *Server) Run() {
 
 // start 启动服务端
 func (s *Server) start() error {
-	// 初始化基本信息
-	s.initBackground()
+	// 初始化上下文
+	s.initContext()
 
 	// 初始化隧道监听器
 	if err := s.initTunnelListener(); err != nil {