From dfa23fde3c28ac069ad8b83d0b2ed0a180d82b7d Mon Sep 17 00:00:00 2001 From: nutcas3 Date: Mon, 20 Apr 2026 20:05:40 +0300 Subject: [PATCH] feat: Add WebSocket handler with hub management, client connections, and real-time broadcasting - Add websocket.Upgrader with CheckOrigin allowing all origins in development - Add Message struct with type/data/timestamp fields for WebSocket messages - Add Client struct with conn/send/mutex/closed fields and readPump/writePump/Close methods - Add Hub struct with clients map, broadcast/register/unregister channels, and sync.RWMutex - Add NewHub constructor and Run method with client registration/ --- apps/api-server/internal/websocket/handler.go | 264 ++++++++++++++++++ 1 file changed, 264 insertions(+) create mode 100644 apps/api-server/internal/websocket/handler.go diff --git a/apps/api-server/internal/websocket/handler.go b/apps/api-server/internal/websocket/handler.go new file mode 100644 index 0000000..5218c3e --- /dev/null +++ b/apps/api-server/internal/websocket/handler.go @@ -0,0 +1,264 @@ +package websocket + +import ( + "encoding/json" + "log" + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" +) + +var upgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + return true // Allow all origins in development + }, +} + +// Message represents a WebSocket message +type Message struct { + Type string `json:"type"` + Data any `json:"data"` + Timestamp string `json:"timestamp"` +} + +// Client represents a WebSocket client +type Client struct { + conn *websocket.Conn + send chan Message + mutex sync.Mutex + closed bool +} + +// Hub manages WebSocket clients and broadcasts +type Hub struct { + clients map[*Client]bool + broadcast chan Message + register chan *Client + unregister chan *Client + mutex sync.RWMutex +} + +// NewHub creates a new WebSocket hub +func NewHub() *Hub { + return &Hub{ + clients: make(map[*Client]bool), + broadcast: make(chan Message), + register: make(chan *Client), + unregister: make(chan *Client), + } +} + +// Run starts the WebSocket hub +func (h *Hub) Run() { + for { + select { + case client := <-h.register: + h.mutex.Lock() + h.clients[client] = true + h.mutex.Unlock() + log.Printf("Client connected. Total clients: %d", len(h.clients)) + + case client := <-h.unregister: + h.mutex.Lock() + if _, ok := h.clients[client]; ok { + delete(h.clients, client) + close(client.send) + } + h.mutex.Unlock() + log.Printf("Client disconnected. Total clients: %d", len(h.clients)) + + case message := <-h.broadcast: + h.mutex.RLock() + for client := range h.clients { + select { + case client.send <- message: + default: + close(client.send) + delete(h.clients, client) + } + } + h.mutex.RUnlock() + } + } +} + +// Broadcast sends a message to all connected clients +func (h *Hub) Broadcast(messageType string, data any) { + message := Message{ + Type: messageType, + Data: data, + Timestamp: time.Now().UTC().Format(time.RFC3339), + } + select { + case h.broadcast <- message: + default: + log.Println("Broadcast channel is full, dropping message") + } +} + +// GetClientCount returns the number of connected clients +func (h *Hub) GetClientCount() int { + h.mutex.RLock() + defer h.mutex.RUnlock() + return len(h.clients) +} + +// readPump handles messages from the WebSocket connection +func (c *Client) readPump(h *Hub) { + defer func() { + h.unregister <- c + c.conn.Close() + }() + + c.conn.SetReadLimit(512) + c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + c.conn.SetPongHandler(func(string) error { + c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) + return nil + }) + + for { + _, messageBytes, err := c.conn.ReadMessage() + if err != nil { + if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { + log.Printf("WebSocket error: %v", err) + } + break + } + + var message Message + if err := json.Unmarshal(messageBytes, &message); err != nil { + log.Printf("Invalid message format: %v", err) + continue + } + + // Handle incoming messages if needed + log.Printf("Received message: %s", message.Type) + } +} + +// writePump handles sending messages to the WebSocket connection +func (c *Client) writePump() { + ticker := time.NewTicker(54 * time.Second) + defer func() { + ticker.Stop() + c.conn.Close() + }() + + for { + select { + case message, ok := <-c.send: + c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if !ok { + c.conn.WriteMessage(websocket.CloseMessage, []byte{}) + return + } + + if err := c.conn.WriteJSON(message); err != nil { + log.Printf("WebSocket write error: %v", err) + return + } + + case <-ticker.C: + c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { + return + } + } + } +} + +// Close safely closes the client connection +func (c *Client) Close() { + c.mutex.Lock() + defer c.mutex.Unlock() + + if !c.closed { + c.closed = true + close(c.send) + c.conn.Close() + } +} + +// NewClient creates a new WebSocket client +func NewClient(conn *websocket.Conn) *Client { + return &Client{ + conn: conn, + send: make(chan Message, 256), + } +} + +// HandleWebSocket handles WebSocket connections +func HandleWebSocket(hub *Hub) gin.HandlerFunc { + return gin.HandlerFunc(func(c *gin.Context) { + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + log.Printf("WebSocket upgrade error: %v", err) + c.JSON(http.StatusBadRequest, gin.H{"error": "Failed to upgrade to WebSocket"}) + return + } + + client := NewClient(conn) + hub.register <- client + + // Start goroutines for reading and writing + go client.readPump(hub) + go client.writePump() + }) +} + +// Global hub instance +var globalHub *Hub + +// InitializeWebSocket initializes the WebSocket hub +func InitializeWebSocket() { + globalHub = NewHub() + go globalHub.Run() +} + +// GetHub returns the global WebSocket hub +func GetHub() *Hub { + return globalHub +} + +// BroadcastServiceUpdate broadcasts a service update to all clients +func BroadcastServiceUpdate(serviceName, status string, replicas, readyReplicas int) { + if globalHub != nil { + data := map[string]any{ + "service": serviceName, + "status": status, + "replicas": replicas, + "readyReplicas": readyReplicas, + } + globalHub.Broadcast("service_update", data) + } +} + +// BroadcastSystemMetrics broadcasts system metrics to all clients +func BroadcastSystemMetrics(cpu, memory, disk, network float64) { + if globalHub != nil { + data := map[string]any{ + "cpu": cpu, + "memory": memory, + "disk": disk, + "network": network, + } + globalHub.Broadcast("system_metrics", data) + } +} + +// BroadcastAlertUpdate broadcasts an alert update to all clients +func BroadcastAlertUpdate(id, severity, message, status string) { + if globalHub != nil { + data := map[string]any{ + "id": id, + "severity": severity, + "message": message, + "status": status, + } + globalHub.Broadcast("alert_update", data) + } +}