diff --git a/internal/websocket/staleClients.go b/internal/websocket/staleClients.go index 456ef11..d679950 100644 --- a/internal/websocket/staleClients.go +++ b/internal/websocket/staleClients.go @@ -2,8 +2,28 @@ package websocket import ( "time" + + "github.com/gorilla/websocket" + + "git.smsvc.net/pomodoro/GoTomato/pkg/models" ) +const TIMEOUT = 10 + +// Checks if the websockets last Pong is recent +func isStale(client *models.WebsocketClient) bool { + return time.Since(client.LastPong) > (TIMEOUT * time.Second) +} + +func sendPing(client *models.WebsocketClient) bool { + client.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second)) + err := client.Conn.WriteMessage(websocket.PingMessage, nil) + if err != nil { + return false + } + return true +} + // Check and remove stale clients func RemoveStaleClients() { ticker := time.NewTicker(30 * time.Second) @@ -12,7 +32,8 @@ func RemoveStaleClients() { for range ticker.C { mu.Lock() for _, client := range Clients { - if client.IsStale() { + // send ping and check for last pong + if !sendPing(client) || 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) }