status/ws.go

package main

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

// wsClient wraps a websocket connection with synchronized writes.
type wsClient struct {
	conn *websocket.Conn
	mu   sync.Mutex
}

// wsHub tracks active websocket clients and broadcasts updates to them.
type wsHub struct {
	mu      sync.Mutex
	clients map[*wsClient]struct{}
}

// wsEnvelope is the websocket payload shape used for live service updates.
type wsEnvelope struct {
	Type        string         `json:"type"`
	GeneratedAt time.Time      `json:"generated_at"`
	Service     *serviceStatus `json:"service,omitempty"`
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// writeJSON writes a JSON payload to the websocket connection safely.
func (c *wsClient) writeJSON(payload any) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	return c.conn.WriteJSON(payload)
}

// newWSHub creates a websocket hub for tracking connected clients.
func newWSHub() *wsHub {
	return &wsHub{clients: make(map[*wsClient]struct{})}
}

// add registers a websocket client with the hub.
func (h *wsHub) add(client *wsClient) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.clients[client] = struct{}{}
}

// remove unregisters a websocket client and closes its connection.
func (h *wsHub) remove(client *wsClient) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if _, ok := h.clients[client]; ok {
		delete(h.clients, client)
		_ = client.conn.Close()
	}
}

// broadcast sends a payload to all connected websocket clients.
func (h *wsHub) broadcast(payload any) {
	h.mu.Lock()
	clients := make([]*wsClient, 0, len(h.clients))
	for client := range h.clients {
		clients = append(clients, client)
	}
	h.mu.Unlock()

	for _, client := range clients {
		if err := client.writeJSON(payload); err != nil {
			h.remove(client)
		}
	}
}

// serveWebSocket upgrades an HTTP request to a websocket and tracks the client.
func serveWebSocket(c *gin.Context, hub *wsHub) {
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}

	client := &wsClient{conn: conn}
	hub.add(client)

	go func() {
		defer hub.remove(client)
		for {
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()
}