From d91c0a9f7fb0f899b4792e681fbaf08412dd0618 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Fri, 24 Sep 2021 15:15:59 +0200
Subject: [PATCH] move websocket Ping to peer.

---
 internal/types/websocket.go   |  1 +
 internal/websocket/manager.go |  9 ++++-----
 internal/websocket/peer.go    | 12 ++++++++++++
 3 files changed, 17 insertions(+), 5 deletions(-)

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,