boc/aamos-admin-upgrade/backend/handlers/websocket.go

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