Moved websocket stuff to a new file

This commit is contained in:
Adrien Bouvais 2024-08-06 10:34:06 +02:00
parent eed7bb36f2
commit 4c3174d534
2 changed files with 61 additions and 49 deletions

57
WebSocket.go Normal file
View File

@ -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
}
}
}

53
main.go
View File

@ -10,6 +10,7 @@ import (
"github.com/gofiber/contrib/websocket" "github.com/gofiber/contrib/websocket"
"github.com/gofiber/fiber/v2" "github.com/gofiber/fiber/v2"
"github.com/gofiber/fiber/v2/middleware/logger" "github.com/gofiber/fiber/v2/middleware/logger"
"github.com/gofiber/fiber/v2/middleware/monitor"
"github.com/gofiber/fiber/v2/middleware/recover" "github.com/gofiber/fiber/v2/middleware/recover"
"github.com/gofiber/template/django/v3" "github.com/gofiber/template/django/v3"
"github.com/stripe/stripe-go" "github.com/stripe/stripe-go"
@ -29,57 +30,8 @@ var (
explainLLMconvChatTmpl *pongo2.Template explainLLMconvChatTmpl *pongo2.Template
mu sync.Mutex mu sync.Mutex
app *fiber.App 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() { func main() {
// Use STRIPE_KEY environment variable // Use STRIPE_KEY environment variable
stripe.Key = os.Getenv("STRIPE_KEY") stripe.Key = os.Getenv("STRIPE_KEY")
@ -157,6 +109,9 @@ func main() {
app.Post("/createLLM", createLLM) app.Post("/createLLM", createLLM)
app.Post("/updateLLMPositionBatch", updateLLMPositionBatch) app.Post("/updateLLMPositionBatch", updateLLMPositionBatch)
//Admin monitoring
app.Get("/metrics", monitor.New())
// Add static files // Add static files
app.Static("/", "./static") app.Static("/", "./static")