From 4c3174d534c536ed969f67e9549f975c04434482 Mon Sep 17 00:00:00 2001
From: MrBounty <adrien.bouvais.pro@gmail.com>
Date: Tue, 6 Aug 2024 10:34:06 +0200
Subject: [PATCH] Moved websocket stuff to a new file

---
 WebSocket.go | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go      | 53 ++++--------------------------------------------
 2 files changed, 61 insertions(+), 49 deletions(-)
 create mode 100644 WebSocket.go

diff --git a/WebSocket.go b/WebSocket.go
new file mode 100644
index 0000000..9a70a44
--- /dev/null
+++ b/WebSocket.go
@@ -0,0 +1,57 @@
+package main
+
+import (
+	"fmt"
+
+	"github.com/gofiber/contrib/websocket"
+)
+
+var userWSChannels = make(map[string]*websocket.Conn)
+
+// Function to send events to all clients
+func sendEvent(userID string, event string, data string) {
+	message := fmt.Sprintf(`{"event": "%s", "data": "%s"}`, event, data)
+	sendWSMessage(userID, message)
+}
+
+func sendWSMessage(userID string, message string) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	conn, ok := userWSChannels[userID]
+	if !ok {
+		return
+	}
+
+	if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
+		fmt.Println("write:", err)
+		conn.Close()
+		delete(userWSChannels, userID)
+	}
+}
+
+func handleWebSocket(c *websocket.Conn) {
+	userID := c.Query("userID")
+	if userID == "" {
+		c.Close()
+		return
+	}
+
+	mu.Lock()
+	userWSChannels[userID] = c
+	mu.Unlock()
+
+	defer func() {
+		mu.Lock()
+		delete(userWSChannels, userID)
+		mu.Unlock()
+		c.Close()
+	}()
+
+	for {
+		if _, _, err := c.ReadMessage(); err != nil {
+			fmt.Println("read:", err)
+			break
+		}
+	}
+}
diff --git a/main.go b/main.go
index 06800ea..ff69c29 100644
--- a/main.go
+++ b/main.go
@@ -10,6 +10,7 @@ import (
 	"github.com/gofiber/contrib/websocket"
 	"github.com/gofiber/fiber/v2"
 	"github.com/gofiber/fiber/v2/middleware/logger"
+	"github.com/gofiber/fiber/v2/middleware/monitor"
 	"github.com/gofiber/fiber/v2/middleware/recover"
 	"github.com/gofiber/template/django/v3"
 	"github.com/stripe/stripe-go"
@@ -29,57 +30,8 @@ var (
 	explainLLMconvChatTmpl  *pongo2.Template
 	mu                      sync.Mutex
 	app                     *fiber.App
-	userWSChannels          = make(map[string]*websocket.Conn)
 )
 
-// Function to send events to all clients
-func sendEvent(userID string, event string, data string) {
-	message := fmt.Sprintf(`{"event": "%s", "data": "%s"}`, event, data)
-	sendWSMessage(userID, message)
-}
-
-func sendWSMessage(userID string, message string) {
-	mu.Lock()
-	defer mu.Unlock()
-
-	conn, ok := userWSChannels[userID]
-	if !ok {
-		return
-	}
-
-	if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
-		fmt.Println("write:", err)
-		conn.Close()
-		delete(userWSChannels, userID)
-	}
-}
-
-func handleWebSocket(c *websocket.Conn) {
-	userID := c.Query("userID")
-	if userID == "" {
-		c.Close()
-		return
-	}
-
-	mu.Lock()
-	userWSChannels[userID] = c
-	mu.Unlock()
-
-	defer func() {
-		mu.Lock()
-		delete(userWSChannels, userID)
-		mu.Unlock()
-		c.Close()
-	}()
-
-	for {
-		if _, _, err := c.ReadMessage(); err != nil {
-			log.Println("read:", err)
-			break
-		}
-	}
-}
-
 func main() {
 	// Use STRIPE_KEY environment variable
 	stripe.Key = os.Getenv("STRIPE_KEY")
@@ -157,6 +109,9 @@ func main() {
 	app.Post("/createLLM", createLLM)
 	app.Post("/updateLLMPositionBatch", updateLLMPositionBatch)
 
+	//Admin monitoring
+	app.Get("/metrics", monitor.New())
+
 	// Add static files
 	app.Static("/", "./static")