From 51a94b6636166d5ee5510bf36822207d11a7be8e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Majdand=C5=BEi=C4=87?=
 <david.majdandzic@hotmail.com>
Date: Thu, 27 Jun 2024 11:16:25 +0200
Subject: [PATCH] Fix memory leak in server

Remove server instrumentation
---
 ws-server/connection.go | 44 +++++++++++++++++++++--------------------
 ws-server/main.go       | 24 +++++++++++++++++++++-
 ws-server/server.go     |  1 +
 3 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/ws-server/connection.go b/ws-server/connection.go
index b366815..542cdf7 100644
--- a/ws-server/connection.go
+++ b/ws-server/connection.go
@@ -53,39 +53,40 @@ func (ws *WSConnection) messageReader() {
 	for {
 		_, message, err := ws.conn.ReadMessage()
 		if !ws.alive {
-			return
+			break
 		}
 		ws.conn.SetReadDeadline(time.Now().Add(ws.IdleTimeout))
 		if err != nil {
 			ws.ErrorChan <- err
-			return
+			break
 		}
 		log.Printf("Client %d: Received: %s, %d in output channel", ws.id, message, len(ws.ReadChan))
 		ws.ReadChan <- string(message)
 	}
+	log.Printf("Client %d: Stopped reading messages", ws.id)
 }
 
 func (ws *WSConnection) messageSender() {
 	log.Printf("Client %d: Sending messages", ws.id)
 	for {
-		func() {
-			msg, ok := <-ws.WriteChan
-			if !ok || !ws.alive {
-				return
-			}
-			ws.writeLock.Lock()
-			defer ws.writeLock.Unlock()
-	
-			ws.conn.SetWriteDeadline(time.Now().Add(ws.IdleTimeout))
-			log.Printf("Client %d: Sending: %s, %d in input channel", ws.id, msg, len(ws.WriteChan))
-			err := ws.conn.WriteMessage(websocket.TextMessage, []byte(msg))
-			if err != nil {
-				log.Printf("Client %d: Error during message writing: %v", ws.id, err)
-				ws.ErrorChan <- err
-				return
-			}
-		}()
+		msg, ok := <-ws.WriteChan
+		if !ok || !ws.alive {
+			break
+		}
+		ws.writeLock.Lock()
+
+		ws.conn.SetWriteDeadline(time.Now().Add(ws.IdleTimeout))
+		log.Printf("Client %d: Sending: %s, %d in input channel", ws.id, msg, len(ws.WriteChan))
+		err := ws.conn.WriteMessage(websocket.TextMessage, []byte(msg))
+		if err != nil {
+			log.Printf("Client %d: Error during message writing: %v", ws.id, err)
+			ws.ErrorChan <- err
+			ws.writeLock.Unlock()
+			break
+		}
+		ws.writeLock.Unlock()
 	}
+	log.Printf("Client %d: Stopped sending messages", ws.id)
 }
 
 func (ws *WSConnection) pinger() {
@@ -93,7 +94,7 @@ func (ws *WSConnection) pinger() {
 	for {
 		time.Sleep(ws.PingInterval)
 		if !ws.alive {
-			return
+			break
 		}
 
 		// log.Printf("Client %d: Ping", ws.id)
@@ -102,10 +103,11 @@ func (ws *WSConnection) pinger() {
 		if err != nil {
 			log.Printf("Client %d: Error during ping: %+v", ws.id, err)
 			ws.ErrorChan <- err
-			return
+			break
 		}
 		ws.conn.SetWriteDeadline(time.Now().Add(ws.IdleTimeout))
 		ws.writeLock.Unlock()
 		// log.Printf("Client %d: Ping OK", ws.id)
 	}
+	log.Printf("Client %d: Stopped pinger", ws.id)
 }
diff --git a/ws-server/main.go b/ws-server/main.go
index af888ef..2b8992a 100644
--- a/ws-server/main.go
+++ b/ws-server/main.go
@@ -44,10 +44,32 @@ func handleDownload(responseWriter http.ResponseWriter, request *http.Request) {
 }
 
 func init() {
-	log.SetFlags(log.Lmicroseconds | log.Lshortfile)
+	// log.SetFlags(log.Lmicroseconds | log.Lshortfile)
+	log.SetFlags(log.Lmicroseconds)
 }
 
+// Mainly used to detect memory leaks
+// func instrument() {
+//     numGoroutines := runtime.NumGoroutine()
+
+//     var m runtime.MemStats
+//     runtime.ReadMemStats(&m)
+
+// 	malloc := float64(m.Alloc)
+//     ramUsedMB := malloc / 1024 / 1024
+//     kbPerGoroutine := malloc / 1024 / float64(numGoroutines)
+
+//     log.Printf("Number of active goroutines: %d; RAM used: %.2f MB; KB per goroutine: %.2f", numGoroutines, ramUsedMB, kbPerGoroutine)
+// }
+
 func main() {
+	// go func() {
+	// 	for {
+	// 		time.Sleep(1 * time.Second)
+	// 		instrument()
+	// 	}
+	// }()
+	
 	http.HandleFunc("/ws", wsHandler)
 	http.HandleFunc("/download", handleDownload)
 	log.Println("Server starting on :8080")
diff --git a/ws-server/server.go b/ws-server/server.go
index f6deaea..508e5b5 100644
--- a/ws-server/server.go
+++ b/ws-server/server.go
@@ -49,6 +49,7 @@ func (server *WSServer) HandleNew(conn *websocket.Conn) {
 		close(wsconn.WriteChan)
 		close(wsconn.ErrorChan)
 		log.Printf("Client %d: disconnected due to %+v, now %d clients", wsconn.id, err, len(server.connections))
+		wsconn.conn.Close()
 		delete(server.connections, wsconn)
 	}()
 }