fix(network): clean up resources on panic in callback to prevent memory leaks

Author: timandy

cluster/acceptor.go

@@ -3,6 +3,7 @@ package cluster
 import (
 	"context"
 	"net"
+	"sync/atomic"
 
 	"github.com/lonng/nano/cluster/clusterpb"
 	"github.com/lonng/nano/internal/env"
@@ -22,17 +23,20 @@ type acceptor struct {
 	lastMid    uint64
 	rpcHandler rpcHandler
 	gateAddr   string
+	state      atomic.Int32 // current acceptor state
 }
 
 // newAcceptor 构造函数
 func newAcceptor(sid int64, node *Node, gateClient clusterpb.GateClient, rpcHandler rpcHandler, gateAddr string) *acceptor {
-	return &acceptor{
+	a := &acceptor{
 		sid:        sid,
 		node:       node,
 		gateClient: gateClient,
 		rpcHandler: rpcHandler,
 		gateAddr:   gateAddr,
 	}
+	a.state.Store(statusWorking)
+	return a
 }
 
 // RemoteAddr 返回一个假的地址
@@ -47,6 +51,9 @@ func (a *acceptor) LastMid() uint64 {
 
 // RPC 调用集群内的服务
 func (a *acceptor) RPC(route string, v any) error {
+	if a.status() == statusClosed {
+		return ErrBrokenPipe
+	}
 	// TODO: buffer
 	data, err := env.Marshal(v)
 	if err != nil {
@@ -63,6 +70,9 @@ func (a *acceptor) RPC(route string, v any) error {
 
 // Push 调用 Gate, 推送数据给客户端
 func (a *acceptor) Push(route string, v any) error {
+	if a.status() == statusClosed {
+		return ErrBrokenPipe
+	}
 	// TODO: buffer
 	data, err := env.Marshal(v)
 	if err != nil {
@@ -84,6 +94,9 @@ func (a *acceptor) Response(v any) error {
 
 // ResponseMid 调用 Gate, 返回响应数据给客户端
 func (a *acceptor) ResponseMid(mid uint64, v any) error {
+	if a.status() == statusClosed {
+		return ErrBrokenPipe
+	}
 	// TODO: buffer
 	data, err := env.Marshal(v)
 	if err != nil {
@@ -100,10 +113,13 @@ func (a *acceptor) ResponseMid(mid uint64, v any) error {
 
 // Close 集群模式下, Worker 节点关闭会话, 通知 Gate 也关闭(主动关闭)
 func (a *acceptor) Close() error {
+	if !a.state.CompareAndSwap(statusWorking, statusClosed) {
+		return ErrCloseClosedSession
+	}
 	// TODO: buffer
 	// 先删除
 	s, found := a.node.delSession(a.sid)
-	// 通知 gate
+	// 通知 Gate 关闭连接
 	request := &clusterpb.CloseSessionRequest{
 		SessionId: a.sid,
 	}
@@ -114,3 +130,8 @@ func (a *acceptor) Close() error {
 	}
 	return err
 }
+
+// status 获取当前状态
+func (a *acceptor) status() int32 {
+	return a.state.Load()
+}

cluster/agent.go

@@ -4,7 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"reflect"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -33,17 +33,19 @@ var _ session.NetworkEntity = (*agent)(nil)
 // agent 与客户端直接通信的网络对象(单点模式中的节点, 或集群模式的网关)
 type agent struct {
 	// regular agent member
-	session    *session.Session    // session
-	conn       net.Conn            // low-level conn fd
-	lastMid    uint64              // last message id
-	state      atomic.Int32        // current agent state
-	chDie      chan struct{}       // wait for close
-	chSend     chan pendingMessage // push message queue
-	lastAt     int64               // last heartbeat unix time stamp
-	decoder    *packet.Decoder     // binary decoder
-	pipeline   pipeline.Pipeline   //
-	rpcHandler rpcHandler          //
-	srv        reflect.Value       // cached session reflect.Value
+	session       *session.Session    // session
+	conn          net.Conn            // low-level conn fd
+	lastMid       uint64              // last message id
+	chDie         chan struct{}       // wait for close
+	chSend        chan pendingMessage // push message queue
+	lastAt        atomic.Int64        // last heartbeat unix time stamp
+	decoder       *packet.Decoder     // binary decoder
+	pipeline      pipeline.Pipeline   //
+	rpcHandler    rpcHandler          //
+	writeReady    atomic.Bool         // write 协程是否已经启动
+	connCloseOnce sync.Once           // 确保 conn 只关闭一次
+	chanCloseOnce sync.Once           // 确保 chDie 和 chSend 只关闭一次
+	state         atomic.Int32        // current agent state
 }
 
 // pendingMessage 待发送的消息
@@ -59,18 +61,17 @@ func newAgent(conn net.Conn, pipeline pipeline.Pipeline, rpcHandler rpcHandler)
 	a := &agent{
 		conn:       conn,
 		chDie:      make(chan struct{}),
-		lastAt:     time.Now().Unix(),
 		chSend:     make(chan pendingMessage, agentWriteBacklog),
 		decoder:    packet.NewDecoder(),
 		pipeline:   pipeline,
 		rpcHandler: rpcHandler,
 	}
+	a.lastAt.Store(time.Now().Unix())
 	a.state.Store(statusStart)
 
 	// binding session
 	s := session.New(a)
 	a.session = s
-	a.srv = reflect.ValueOf(s)
 	return a
 }
 
@@ -157,7 +158,7 @@ func (a *agent) ResponseMid(mid uint64, v any) error {
 	return a.send(pendingMessage{typ: message.Response, mid: mid, payload: v})
 }
 
-// Close 关闭低级连接, 任何被阻塞读取或写入作都将被取消阻塞并返回错误
+// Close 设置关闭状态, 发送关闭信号; 如果 write 协程未就绪, 直接关闭底层连接; 否则由 write 协程 flush 完数据后关闭
 func (a *agent) Close() error {
 	if a.status() == statusClosed {
 		return ErrCloseClosedSession
@@ -168,19 +169,38 @@ func (a *agent) Close() error {
 		log.Info("Session closing, ID=%d, UID=%d, IP=%s", a.session.ID(), a.session.UID(), a.conn.RemoteAddr())
 	}
 
-	// 防止关闭已经是关闭状态的 chan
-	select {
-	case <-a.chDie:
-		// expect
-	default:
-		close(a.chDie)
+	// 关闭 chan, 发出停止信号
+	a.closeChanOnce()
+
+	// 如果 write 协程已经启动, 则不需要关闭底层连接, 因为 write 协程 flush 完会自动处理
+	if a.writeReady.Load() {
+		return nil
 	}
-	return nil
+
+	// 如果 write 协程还未启动, 则直接关闭底层连接
+	return a.closeConnOnce()
 }
 
 // String 返回描述信息
 func (a *agent) String() string {
-	return fmt.Sprintf("Remote=%s, LastTime=%d", a.conn.RemoteAddr().String(), atomic.LoadInt64(&a.lastAt))
+	return fmt.Sprintf("Remote=%s, LastTime=%d", a.conn.RemoteAddr().String(), a.lastAt.Load())
+}
+
+// closeChanOnce 确保 chan 只关闭一次
+func (a *agent) closeChanOnce() {
+	a.chanCloseOnce.Do(func() {
+		close(a.chDie)
+		close(a.chSend)
+	})
+}
+
+// closeConnOnce 确保 conn 只关闭一次
+func (a *agent) closeConnOnce() error {
+	var err error
+	a.connCloseOnce.Do(func() {
+		err = a.conn.Close()
+	})
+	return err
 }
 
 // status 获取当前状态
@@ -202,28 +222,33 @@ func (a *agent) write() {
 	// clean func
 	defer func() {
 		ticker.Stop()
-		close(a.chSend)
 		close(chWrite)
-		//非强制退出, 则将所有待发送的消息写入底层连接
+		// 关闭 chan, 必须关闭 chan 后才能执行 flush, 否则阻塞
+		a.closeChanOnce()
+		// 非强制退出, 则将所有待发送的消息写入底层连接
 		if !forceQuit {
 			a.flush(chWrite)
 		}
-		//更改 agent 状态, 必须先更改状态再关闭底层连接
+		// 更改 agent 状态, 必须先更改状态再关闭底层连接
 		_ = a.Close()
-		//关闭底层连接, 此时 conn.Read() 将返回错误, 因上一步已经把状态关闭, 所以读协程会跳过日志退出
-		_ = a.conn.Close()
+		// 关闭底层连接, 此时 conn.Read() 将返回错误, 因上一步已经把状态关闭, 所以读协程会跳过日志退出
+		_ = a.closeConnOnce()
 		if env.Debug {
 			log.Info("Session write goroutine exit, SessionID=%d, UID=%d", a.session.ID(), a.session.UID())
 		}
 	}()
 
+	// 标记 write 协程已经就绪
+	a.writeReady.Store(true)
+
 	for {
 		select {
 		// 心跳检测
 		case <-ticker.C:
 			deadline := time.Now().Add(-2 * heartbeat).Unix()
-			if atomic.LoadInt64(&a.lastAt) < deadline {
-				log.Info("Session heartbeat timeout, LastTime=%d, Deadline=%d", atomic.LoadInt64(&a.lastAt), deadline)
+			lastAt := a.lastAt.Load()
+			if lastAt < deadline {
+				log.Info("Session heartbeat timeout, LastTime=%d, Deadline=%d", lastAt, deadline)
 				return
 			}
 			chWrite <- getHbd()

cluster/handler.go

@@ -299,7 +299,7 @@ func (h *LocalHandler) processPacket(agent *agent, p *packet.Packet) error {
 		return errors.New("invalid packet type")
 	}
 
-	agent.lastAt = time.Now().Unix()
+	agent.lastAt.Store(time.Now().Unix())
 	return nil
 }
 

session/event.go

@@ -39,6 +39,14 @@ func (lt *event) FireSessionCreated(s *Session) {
 		return
 	}
 
+	// 回调过程中如果发生 panic, 需要关闭底层连接以防止连接泄露
+	defer func() {
+		if r := recover(); r != nil {
+			s.Close()
+			panic(r) // 重新抛出 panic
+		}
+	}()
+
 	for _, fn := range lt.onSessionCreated {
 		fn(s)
 	}
@@ -55,12 +63,12 @@ func (lt *event) FireSessionClosed(s *Session) {
 		return
 	}
 
+	// 执行完回调, 清除会话关联的数据
+	defer s.Clear()
+
 	for _, fn := range lt.onSessionClosed {
 		fn(s)
 	}
-
-	// 执行完回调, 清除会话关联的数据
-	s.Clear()
 }
 
 // MessagePushing 设置消息推送前的回调, 可以修改消息内容和路由