From ab2e8c161d9447ab36e367a2a5047f9e325dda57 Mon Sep 17 00:00:00 2001 From: Sebastian Mark Date: Wed, 20 Nov 2024 14:07:00 +0100 Subject: [PATCH 1/5] feat: implement stale client removal after timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - add goroutine to remove stale clients from the connection pool - update client struct to track `LastPong` time - set write deadline for websocket connections - move package variables and const to `vars.go` - log additional information when broadcasting errors 🤖 --- cmd/server/main.go | 1 + internal/websocket/broadcast.go | 7 +++-- internal/websocket/client_commands.go | 1 + internal/websocket/handle_connection.go | 17 +++++++------ internal/websocket/staleClients.go | 34 +++++++++++++++++++++++++ internal/websocket/vars.go | 18 +++++++++++++ pkg/models/client.go | 18 +++++-------- 7 files changed, 75 insertions(+), 21 deletions(-) create mode 100644 internal/websocket/staleClients.go create mode 100644 internal/websocket/vars.go diff --git a/cmd/server/main.go b/cmd/server/main.go index d8eb8f6..78c5c99 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -54,6 +54,7 @@ func Start() { r := http.NewServeMux() r.HandleFunc("/", websocket.HandleConnection) go websocket.SendPermanentBroadCastMessage() + go websocket.RemoveStaleClients() helper.Logger.Info("GoTomato started", "version", metadata.GoTomatoVersion) helper.Logger.Info("Websocket listening", "address", listen) diff --git a/internal/websocket/broadcast.go b/internal/websocket/broadcast.go index def8309..dce2f82 100644 --- a/internal/websocket/broadcast.go +++ b/internal/websocket/broadcast.go @@ -12,6 +12,7 @@ import ( // Sends continous messages to all connected WebSocket clients func SendPermanentBroadCastMessage() { tick := time.NewTicker(time.Second) + for { // Marshal the message into JSON format jsonMessage, err := json.Marshal(shared.State) @@ -19,14 +20,16 @@ func SendPermanentBroadCastMessage() { helper.Logger.Error("Error marshalling message:", "msg", err) return } + // Iterate over all connected clients and broadcast the message for _, client := range Clients { + // Send message to client err := client.SendMessage(websocket.TextMessage, jsonMessage) if err != nil { - helper.Logger.Error("Error broadcasting to client:", "msg", err) - // The client is responsible for closing itself on error + helper.Logger.Error("Error broadcasting to client:", "msg", err, "host", client.RealIP, "clients", len(Clients)) } } + <-tick.C } } diff --git a/internal/websocket/client_commands.go b/internal/websocket/client_commands.go index 7149faa..6f85108 100644 --- a/internal/websocket/client_commands.go +++ b/internal/websocket/client_commands.go @@ -21,6 +21,7 @@ func handleClientCommands(c models.WebsocketClient) { _, message, err := ws.ReadMessage() if err != nil { + // remove client on error/disconnect delete(Clients, ws.LocalAddr()) helper.Logger.Info("Client disconnected:", "msg", err, "host", c.RealIP, "clients", len(Clients)) break diff --git a/internal/websocket/handle_connection.go b/internal/websocket/handle_connection.go index 9c4b43f..ca55229 100644 --- a/internal/websocket/handle_connection.go +++ b/internal/websocket/handle_connection.go @@ -2,18 +2,13 @@ package websocket import ( "github.com/gorilla/websocket" - "net" "net/http" - "sync" + "time" "git.smsvc.net/pomodoro/GoTomato/internal/helper" "git.smsvc.net/pomodoro/GoTomato/pkg/models" ) -// Clients is a map of connected WebSocket clients, where each client is represented by the WebsocketClient struct -var Clients = make(map[net.Addr]*models.WebsocketClient) -var mu sync.Mutex // Mutex to protect access to the Clients map - // Upgrade HTTP requests to WebSocket connections var upgrader = websocket.Upgrader{ CheckOrigin: func(r *http.Request) bool { return true }, @@ -31,9 +26,15 @@ func HandleConnection(w http.ResponseWriter, r *http.Request) { // Register the new client client := models.WebsocketClient{ - Conn: ws, - RealIP: r.RemoteAddr, + Conn: ws, + LastPong: time.Now(), + RealIP: r.RemoteAddr, } + client.Conn.SetPongHandler(func(s string) error { + client.LastPong = time.Now() + return nil + }) + mu.Lock() Clients[ws.LocalAddr()] = &client mu.Unlock() diff --git a/internal/websocket/staleClients.go b/internal/websocket/staleClients.go new file mode 100644 index 0000000..b636c49 --- /dev/null +++ b/internal/websocket/staleClients.go @@ -0,0 +1,34 @@ +package websocket + +import ( + "time" + + "git.smsvc.net/pomodoro/GoTomato/internal/helper" + "git.smsvc.net/pomodoro/GoTomato/pkg/models" + "github.com/gorilla/websocket" +) + +// Check and remove stale clients +func RemoveStaleClients() { + ticker := time.NewTicker(STALE_CHECK_INTERVALL * time.Second) + defer ticker.Stop() + + for range ticker.C { + mu.Lock() + for _, client := range Clients { + client.Conn.SetWriteDeadline(time.Now().Add(SEND_TIMEOUT * time.Second)) + client.Conn.WriteMessage(websocket.PingMessage, nil) + + if isStale(client) { + helper.Logger.Info("Removing stale client", "host", client.RealIP, "lastPong", client.LastPong.Format(time.RFC3339)) + client.Conn.Close() + delete(Clients, client.Conn.LocalAddr()) + } + } + mu.Unlock() + } +} + +func isStale(client *models.WebsocketClient) bool { + return time.Since(client.LastPong) > (STALE_CLIENT_TIMEOUT * time.Second) +} diff --git a/internal/websocket/vars.go b/internal/websocket/vars.go new file mode 100644 index 0000000..5aedb53 --- /dev/null +++ b/internal/websocket/vars.go @@ -0,0 +1,18 @@ +package websocket + +import ( + "net" + "sync" + + "git.smsvc.net/pomodoro/GoTomato/pkg/models" +) + +const SEND_TIMEOUT = 10 +const STALE_CLIENT_TIMEOUT = 90 +const STALE_CHECK_INTERVALL = 30 + +// Clients is a map of connected WebSocket clients, where each client is represented by the WebsocketClient struct +var Clients = make(map[net.Addr]*models.WebsocketClient) + +// Mutex to protect access to the Clients map +var mu sync.Mutex diff --git a/pkg/models/client.go b/pkg/models/client.go index eef9663..84979ee 100644 --- a/pkg/models/client.go +++ b/pkg/models/client.go @@ -1,9 +1,9 @@ package models import ( - "github.com/gorilla/websocket" + "time" - "git.smsvc.net/pomodoro/GoTomato/internal/helper" + "github.com/gorilla/websocket" ) // Represents a command from the client (start/stop) @@ -15,17 +15,13 @@ type ClientCommand struct { // Represents a single client type WebsocketClient struct { - Conn *websocket.Conn - RealIP string + Conn *websocket.Conn + LastPong time.Time + RealIP string } // Sends a message to the websocket. -// Automatically locks and unlocks the client mutex, to ensure that only one goroutine can write at a time. func (c *WebsocketClient) SendMessage(messageType int, data []byte) error { - err := c.Conn.WriteMessage(messageType, data) - if err != nil { - helper.Logger.Error("Error writing to WebSocket:", "msg", err) - c.Conn.Close() // Close the connection on error - } - return err + c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second)) + return c.Conn.WriteMessage(messageType, data) } From 064a720436c2cf58a56d18e0666bb7b13045f496 Mon Sep 17 00:00:00 2001 From: Sebastian Mark Date: Wed, 20 Nov 2024 21:51:12 +0100 Subject: [PATCH 2/5] fix: add more mutex locks for `Clients` map --- internal/websocket/broadcast.go | 2 ++ internal/websocket/client_commands.go | 2 ++ 2 files changed, 4 insertions(+) diff --git a/internal/websocket/broadcast.go b/internal/websocket/broadcast.go index dce2f82..f9fcab3 100644 --- a/internal/websocket/broadcast.go +++ b/internal/websocket/broadcast.go @@ -22,6 +22,7 @@ func SendPermanentBroadCastMessage() { } // Iterate over all connected clients and broadcast the message + mu.Lock() for _, client := range Clients { // Send message to client err := client.SendMessage(websocket.TextMessage, jsonMessage) @@ -29,6 +30,7 @@ func SendPermanentBroadCastMessage() { helper.Logger.Error("Error broadcasting to client:", "msg", err, "host", client.RealIP, "clients", len(Clients)) } } + mu.Unlock() <-tick.C } diff --git a/internal/websocket/client_commands.go b/internal/websocket/client_commands.go index 6f85108..7712d08 100644 --- a/internal/websocket/client_commands.go +++ b/internal/websocket/client_commands.go @@ -22,7 +22,9 @@ func handleClientCommands(c models.WebsocketClient) { _, message, err := ws.ReadMessage() if err != nil { // remove client on error/disconnect + mu.Lock() delete(Clients, ws.LocalAddr()) + mu.Unlock() helper.Logger.Info("Client disconnected:", "msg", err, "host", c.RealIP, "clients", len(Clients)) break } From 1ed866bfe54bb595a57e14623ebce26820fc6443 Mon Sep 17 00:00:00 2001 From: Sebastian Mark Date: Thu, 21 Nov 2024 08:43:02 +0100 Subject: [PATCH 3/5] refactor: merge `SendMessage()` into `SendPermanentBroadCastMessage()` --- internal/websocket/broadcast.go | 3 ++- pkg/models/client.go | 6 ------ 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/internal/websocket/broadcast.go b/internal/websocket/broadcast.go index f9fcab3..57f4449 100644 --- a/internal/websocket/broadcast.go +++ b/internal/websocket/broadcast.go @@ -25,7 +25,8 @@ func SendPermanentBroadCastMessage() { mu.Lock() for _, client := range Clients { // Send message to client - err := client.SendMessage(websocket.TextMessage, jsonMessage) + client.Conn.SetWriteDeadline(time.Now().Add(SEND_TIMEOUT * time.Second)) + err := client.Conn.WriteMessage(websocket.TextMessage, jsonMessage) if err != nil { helper.Logger.Error("Error broadcasting to client:", "msg", err, "host", client.RealIP, "clients", len(Clients)) } diff --git a/pkg/models/client.go b/pkg/models/client.go index 84979ee..9d5dbdb 100644 --- a/pkg/models/client.go +++ b/pkg/models/client.go @@ -19,9 +19,3 @@ type WebsocketClient struct { LastPong time.Time RealIP string } - -// Sends a message to the websocket. -func (c *WebsocketClient) SendMessage(messageType int, data []byte) error { - c.Conn.SetWriteDeadline(time.Now().Add(TIMEOUT * time.Second)) - return c.Conn.WriteMessage(messageType, data) -} From a09de61a333863f15efa5485e07069539376e9da Mon Sep 17 00:00:00 2001 From: Sebastian Mark Date: Thu, 21 Nov 2024 18:49:10 +0100 Subject: [PATCH 4/5] refactor: improve broadcast loop ticker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - rename tick to ticker for clarity - change infinite loop to range over ticker channel - add defer statement to stop the ticker properly - remove unnecessary channel receive operation 🤖 --- internal/websocket/broadcast.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/internal/websocket/broadcast.go b/internal/websocket/broadcast.go index 57f4449..e2c175d 100644 --- a/internal/websocket/broadcast.go +++ b/internal/websocket/broadcast.go @@ -11,9 +11,10 @@ import ( // Sends continous messages to all connected WebSocket clients func SendPermanentBroadCastMessage() { - tick := time.NewTicker(time.Second) + ticker := time.NewTicker(time.Second) + defer ticker.Stop() - for { + for range ticker.C { // Marshal the message into JSON format jsonMessage, err := json.Marshal(shared.State) if err != nil { @@ -32,7 +33,5 @@ func SendPermanentBroadCastMessage() { } } mu.Unlock() - - <-tick.C } } From 9a86adaf85b9a86a6a56b486f222697809a1dc56 Mon Sep 17 00:00:00 2001 From: Sebastian Mark Date: Thu, 21 Nov 2024 18:54:08 +0100 Subject: [PATCH 5/5] fix: introduce const for broadcast interval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - change ticker to use BROADCAST_INTERVAL constant - define BROADCAST_INTERVAL as 1 second 🤖 --- internal/websocket/broadcast.go | 2 +- internal/websocket/vars.go | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/internal/websocket/broadcast.go b/internal/websocket/broadcast.go index e2c175d..32a6477 100644 --- a/internal/websocket/broadcast.go +++ b/internal/websocket/broadcast.go @@ -11,7 +11,7 @@ import ( // Sends continous messages to all connected WebSocket clients func SendPermanentBroadCastMessage() { - ticker := time.NewTicker(time.Second) + ticker := time.NewTicker(BROADCAST_INTERVAL * time.Second) defer ticker.Stop() for range ticker.C { diff --git a/internal/websocket/vars.go b/internal/websocket/vars.go index 5aedb53..ad02180 100644 --- a/internal/websocket/vars.go +++ b/internal/websocket/vars.go @@ -7,6 +7,7 @@ import ( "git.smsvc.net/pomodoro/GoTomato/pkg/models" ) +const BROADCAST_INTERVAL = 1 const SEND_TIMEOUT = 10 const STALE_CLIENT_TIMEOUT = 90 const STALE_CHECK_INTERVALL = 30