fix: agent ws connect panic
This commit is contained in:
@@ -407,17 +407,28 @@ func (c *AgentController) WSConnect(ctx *gin.Context) {
|
|||||||
// wsReadPump 读取消息
|
// wsReadPump 读取消息
|
||||||
func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models.Agent) {
|
func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models.Agent) {
|
||||||
defer func() {
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
logger.Errorf("[AgentWS] Agent #%d wsReadPump panic: %v", agent.ID, r)
|
||||||
|
}
|
||||||
c.wsManager.Unregister(agent.ID)
|
c.wsManager.Unregister(agent.ID)
|
||||||
}()
|
}()
|
||||||
|
|
||||||
ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
|
// 检查连接是否有效
|
||||||
|
if ac == nil || ac.IsClosed() {
|
||||||
|
logger.Warnf("[AgentWS] Agent #%d 连接无效", agent.ID)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
ac.SetReadDeadline(time.Now().Add(90 * time.Second))
|
||||||
|
// 注意:SetPongHandler 需要直接访问 Conn,但这里我们在连接建立后立即设置
|
||||||
|
// 所以是安全的,因为此时连接还没有被其他 goroutine 关闭
|
||||||
ac.Conn.SetPongHandler(func(string) error {
|
ac.Conn.SetPongHandler(func(string) error {
|
||||||
ac.Conn.SetReadDeadline(time.Now().Add(90 * time.Second))
|
ac.SetReadDeadline(time.Now().Add(90 * time.Second))
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
|
|
||||||
for {
|
for {
|
||||||
_, message, err := ac.Conn.ReadMessage()
|
_, message, err := ac.ReadMessage()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||||
logger.Warnf("[AgentWS] Agent #%d 读取错误: %v", agent.ID, err)
|
logger.Warnf("[AgentWS] Agent #%d 读取错误: %v", agent.ID, err)
|
||||||
@@ -436,6 +447,12 @@ func (c *AgentController) wsReadPump(ac *services.AgentConnection, agent *models
|
|||||||
|
|
||||||
// wsWritePump 写入消息
|
// wsWritePump 写入消息
|
||||||
func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
|
func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
logger.Errorf("[AgentWS] Agent #%d wsWritePump panic: %v", ac.AgentID, r)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
ticker := time.NewTicker(30 * time.Second)
|
ticker := time.NewTicker(30 * time.Second)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
|
|
||||||
@@ -445,12 +462,17 @@ func (c *AgentController) wsWritePump(ac *services.AgentConnection) {
|
|||||||
if !ok {
|
if !ok {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
if ac.IsClosed() {
|
||||||
|
return
|
||||||
|
}
|
||||||
if err := ac.WriteMessage(message); err != nil {
|
if err := ac.WriteMessage(message); err != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
case <-ticker.C:
|
case <-ticker.C:
|
||||||
ac.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
if ac.IsClosed() {
|
||||||
if err := ac.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
return
|
||||||
|
}
|
||||||
|
if err := ac.WritePing(); err != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -35,6 +35,7 @@ type AgentConnection struct {
|
|||||||
Conn *websocket.Conn
|
Conn *websocket.Conn
|
||||||
Send chan []byte
|
Send chan []byte
|
||||||
LastPing time.Time
|
LastPing time.Time
|
||||||
|
closed bool
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -266,27 +267,71 @@ func (m *AgentWSManager) cleanupLoop() {
|
|||||||
func (c *AgentConnection) Close() {
|
func (c *AgentConnection) Close() {
|
||||||
c.mu.Lock()
|
c.mu.Lock()
|
||||||
defer c.mu.Unlock()
|
defer c.mu.Unlock()
|
||||||
|
if c.closed {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
c.closed = true
|
||||||
if c.Conn != nil {
|
if c.Conn != nil {
|
||||||
c.Conn.Close()
|
c.Conn.Close()
|
||||||
c.Conn = nil
|
|
||||||
}
|
}
|
||||||
if c.Send != nil {
|
if c.Send != nil {
|
||||||
close(c.Send)
|
close(c.Send)
|
||||||
c.Send = nil
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// IsClosed 检查连接是否已关闭
|
||||||
|
func (c *AgentConnection) IsClosed() bool {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
return c.closed
|
||||||
|
}
|
||||||
|
|
||||||
// WriteMessage 写入消息
|
// WriteMessage 写入消息
|
||||||
func (c *AgentConnection) WriteMessage(data []byte) error {
|
func (c *AgentConnection) WriteMessage(data []byte) error {
|
||||||
c.mu.Lock()
|
c.mu.Lock()
|
||||||
defer c.mu.Unlock()
|
defer c.mu.Unlock()
|
||||||
if c.Conn == nil {
|
if c.closed || c.Conn == nil {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||||||
return c.Conn.WriteMessage(websocket.TextMessage, data)
|
return c.Conn.WriteMessage(websocket.TextMessage, data)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetReadDeadline 设置读取超时
|
||||||
|
func (c *AgentConnection) SetReadDeadline(t time.Time) error {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
if c.closed || c.Conn == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return c.Conn.SetReadDeadline(t)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReadMessage 读取消息
|
||||||
|
func (c *AgentConnection) ReadMessage() (int, []byte, error) {
|
||||||
|
// 不加锁,因为 ReadMessage 是阻塞的
|
||||||
|
// 但需要先检查连接状态
|
||||||
|
c.mu.Lock()
|
||||||
|
if c.closed || c.Conn == nil {
|
||||||
|
c.mu.Unlock()
|
||||||
|
return 0, nil, websocket.ErrCloseSent
|
||||||
|
}
|
||||||
|
conn := c.Conn
|
||||||
|
c.mu.Unlock()
|
||||||
|
return conn.ReadMessage()
|
||||||
|
}
|
||||||
|
|
||||||
|
// WritePing 发送 ping 消息
|
||||||
|
func (c *AgentConnection) WritePing() error {
|
||||||
|
c.mu.Lock()
|
||||||
|
defer c.mu.Unlock()
|
||||||
|
if c.closed || c.Conn == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||||||
|
return c.Conn.WriteMessage(websocket.PingMessage, nil)
|
||||||
|
}
|
||||||
|
|
||||||
// UpdatePing 更新心跳时间
|
// UpdatePing 更新心跳时间
|
||||||
func (c *AgentConnection) UpdatePing() {
|
func (c *AgentConnection) UpdatePing() {
|
||||||
c.LastPing = time.Now()
|
c.LastPing = time.Now()
|
||||||
|
|||||||
Reference in New Issue
Block a user