From 2f18b472f6825fd66f839d882cbb5bf9a63dce27 Mon Sep 17 00:00:00 2001
From: Sebastian Mark <smark@posteo.net>
Date: Thu, 21 Nov 2024 08:27:15 +0100
Subject: [PATCH] move ping from `SendMessage()` to `RemoveStaleClients()`

---
 internal/websocket/staleClients.go | 26 +++++++++++++++++++++++++-
 pkg/models/client.go               | 18 +-----------------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/internal/websocket/staleClients.go b/internal/websocket/staleClients.go
index 456ef11..247af04 100644
--- a/internal/websocket/staleClients.go
+++ b/internal/websocket/staleClients.go
@@ -2,8 +2,26 @@ package websocket
 
 import (
 	"time"
+
+	"git.smsvc.net/pomodoro/GoTomato/pkg/models"
+	"github.com/gorilla/websocket"
 )
 
+func sendPing(client *models.WebsocketClient) bool {
+	client.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+	err := client.Conn.WriteMessage(websocket.PingMessage, nil)
+
+	if err != nil {
+		return false
+	}
+
+	return true
+}
+
+func isStale(client *models.WebsocketClient) bool {
+	return time.Since(client.LastPong) > (90 * time.Second)
+}
+
 // Check and remove stale clients
 func RemoveStaleClients() {
 	ticker := time.NewTicker(30 * time.Second)
@@ -12,7 +30,13 @@ func RemoveStaleClients() {
 	for range ticker.C {
 		mu.Lock()
 		for _, client := range Clients {
-			if client.IsStale() {
+			if !sendPing(client) {
+				client.Conn.Close()
+				delete(Clients, client.Conn.LocalAddr())
+
+			}
+
+			if isStale(client) {
 				client.Conn.Close()
 				delete(Clients, client.Conn.LocalAddr())
 			}
diff --git a/pkg/models/client.go b/pkg/models/client.go
index a99fe2d..ae85a94 100644
--- a/pkg/models/client.go
+++ b/pkg/models/client.go
@@ -25,21 +25,5 @@ type WebsocketClient struct {
 // Sends a message to the websocket.
 func (c *WebsocketClient) SendMessage(messageType int, data []byte) error {
 	c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
-	err := c.Conn.WriteMessage(websocket.PingMessage, nil)
-	if err != nil {
-		return err
-	}
-
-	c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second))
-	err = c.Conn.WriteMessage(messageType, data)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
-// Checks if the websockets last Pong is recent
-func (c *WebsocketClient) IsStale() bool {
-	return time.Since(c.LastPong) > (TIMEOUT * time.Second)
+	return c.Conn.WriteMessage(messageType, data)
 }