diff --git a/internal/types/websocket.go b/internal/types/websocket.go index f876b0fc..70ce0cb9 100644 --- a/internal/types/websocket.go +++ b/internal/types/websocket.go @@ -16,6 +16,7 @@ type CheckOrigin func(r *http.Request) bool type WebSocketPeer interface { Send(event string, payload interface{}) + Ping() error Destroy(reason string) } diff --git a/internal/websocket/manager.go b/internal/websocket/manager.go index ef9aed37..11cc3392 100644 --- a/internal/websocket/manager.go +++ b/internal/websocket/manager.go @@ -215,10 +215,10 @@ func (manager *WebSocketManagerCtx) connect(connection *websocket.Conn, r *http. session.SetWebSocketConnected(peer, false) }() - manager.handle(connection, session) + manager.handle(connection, peer, session) } -func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, session types.Session) { +func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, peer types.WebSocketPeer, session types.Session) { // add session id to logger context logger := manager.logger.With().Str("session_id", session.ID()).Logger() @@ -281,11 +281,10 @@ func (manager *WebSocketManagerCtx) handle(connection *websocket.Conn, session t case <-cancel: return case <-manager.shutdown: - err := connection.Close() - manager.logger.Err(err).Msg("connection shutdown") + peer.Destroy("connection shutdown") return case <-ticker.C: - if err := connection.WriteMessage(websocket.PingMessage, nil); err != nil { + if err := peer.Ping(); err != nil { logger.Err(err).Msg("ping message has failed") return } diff --git a/internal/websocket/peer.go b/internal/websocket/peer.go index d25367b4..3d76a13a 100644 --- a/internal/websocket/peer.go +++ b/internal/websocket/peer.go @@ -2,6 +2,7 @@ package websocket import ( "encoding/json" + "errors" "sync" "github.com/gorilla/websocket" @@ -66,6 +67,17 @@ func (peer *WebSocketPeerCtx) Send(event string, payload interface{}) { Msg("sending message to client") } +func (peer *WebSocketPeerCtx) Ping() error { + peer.mu.Lock() + defer peer.mu.Unlock() + + if peer.connection == nil { + return errors.New("peer connection not found") + } + + return peer.connection.WriteMessage(websocket.PingMessage, nil) +} + func (peer *WebSocketPeerCtx) Destroy(reason string) { peer.Send( event.SYSTEM_DISCONNECT,