package handlers import ( "encoding/json" "log" "net/http" "time" "github.com/gorilla/websocket" ) // SystemMetric is the payload pushed to every connected WebSocket client. type SystemMetric = MetricsResponse const ( writeWait = 10 * time.Second pongWait = 60 * time.Second pingPeriod = (pongWait * 9) / 10 maxMsgSize = 512 ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 4096, CheckOrigin: func(r *http.Request) bool { return true }, } // Client is a single WebSocket connection managed by the Hub. type Client struct { hub *Hub conn *websocket.Conn send chan []byte } // Hub maintains the set of active clients and owns all map mutations. // All access to clients happens in the single Run() goroutine — no mutex needed. type Hub struct { clients map[*Client]struct{} register chan *Client unregister chan *Client broadcast chan []byte } func NewHub() *Hub { return &Hub{ clients: make(map[*Client]struct{}), register: make(chan *Client), unregister: make(chan *Client), broadcast: make(chan []byte, 8), } } // Run starts the hub event loop. Call once in a dedicated goroutine. func (h *Hub) Run() { ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() for { select { case c := <-h.register: h.clients[c] = struct{}{} case c := <-h.unregister: if _, ok := h.clients[c]; ok { delete(h.clients, c) close(c.send) } case msg := <-h.broadcast: h.fanout(msg) case <-ticker.C: msg, err := collectSystemMetrics() if err != nil { log.Printf("ws: metrics collect error: %v", err) continue } h.fanout(msg) } } } func (h *Hub) fanout(msg []byte) { for c := range h.clients { select { case c.send <- msg: default: // slow / blocked client — disconnect close(c.send) delete(h.clients, c) } } } // ServeWS upgrades an HTTP request to a WebSocket connection and registers // the resulting client with the hub. func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("ws: upgrade error: %v", err) return } c := &Client{ hub: h, conn: conn, send: make(chan []byte, 16), } h.register <- c go c.writePump() go c.readPump() } // readPump drains incoming frames (clients only send pongs) and detects disconnects. func (c *Client) readPump() { defer func() { c.hub.unregister <- c c.conn.Close() }() c.conn.SetReadLimit(maxMsgSize) _ = c.conn.SetReadDeadline(time.Now().Add(pongWait)) c.conn.SetPongHandler(func(string) error { return c.conn.SetReadDeadline(time.Now().Add(pongWait)) }) for { if _, _, err := c.conn.ReadMessage(); err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure, ) { log.Printf("ws: read error: %v", err) } return } } } // writePump serialises all writes to the connection and sends periodic pings. func (c *Client) writePump() { pingTicker := time.NewTicker(pingPeriod) defer func() { pingTicker.Stop() c.conn.Close() }() for { select { case msg, ok := <-c.send: _ = c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if !ok { _ = c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil { return } case <-pingTicker.C: _ = c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } func collectSystemMetrics() ([]byte, error) { cpu, err := readCPU() if err != nil { return nil, err } ram, err := readRAM() if err != nil { return nil, err } disk, err := readDisk("/") if err != nil { return nil, err } procs, err := readProcesses() if err != nil { return nil, err } m := SystemMetric{ Timestamp: time.Now().UTC().Format(time.RFC3339), CPU: cpu, RAM: ram, Disk: disk, Processes: procs, } return json.Marshal(m) }