109 lines
2.4 KiB
Go
109 lines
2.4 KiB
Go
package hub
|
|
|
|
import (
|
|
"encoding/json"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// WSMessage is sent to browsers (matches frontend ServerStatus shape).
|
|
type WSMessage struct {
|
|
Type string `json:"type"`
|
|
Data map[string]json.RawMessage `json:"data"`
|
|
}
|
|
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
// serverId -> latest metrics JSON object (not wrapped)
|
|
metrics map[string]json.RawMessage
|
|
|
|
subsMu sync.Mutex
|
|
// subscribers for broadcast
|
|
conns map[*websocket.Conn]struct{}
|
|
}
|
|
|
|
func New() *Hub {
|
|
return &Hub{
|
|
metrics: make(map[string]json.RawMessage),
|
|
conns: make(map[*websocket.Conn]struct{}),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) SetMetrics(serverID string, metricsJSON []byte) {
|
|
h.mu.Lock()
|
|
if len(metricsJSON) == 0 {
|
|
delete(h.metrics, serverID)
|
|
} else {
|
|
h.metrics[serverID] = json.RawMessage(append([]byte(nil), metricsJSON...))
|
|
}
|
|
h.mu.Unlock()
|
|
}
|
|
|
|
func (h *Hub) snapshotData() map[string]json.RawMessage {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
out := make(map[string]json.RawMessage, len(h.metrics))
|
|
now := time.Now().UnixMilli()
|
|
for id, raw := range h.metrics {
|
|
// Wrap as ServerStatus for frontend compatibility
|
|
wrap := map[string]any{
|
|
"serverId": id,
|
|
"online": true,
|
|
"metrics": json.RawMessage(raw),
|
|
"lastUpdate": now,
|
|
}
|
|
b, err := json.Marshal(wrap)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
out[id] = b
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (h *Hub) BuildSnapshotMessage() ([]byte, error) {
|
|
msg := WSMessage{
|
|
Type: "snapshot",
|
|
Data: h.snapshotData(),
|
|
}
|
|
return json.Marshal(msg)
|
|
}
|
|
|
|
// AddSubscriber registers for broadcast fan-out (call after sending the initial snapshot).
|
|
func (h *Hub) AddSubscriber(c *websocket.Conn) {
|
|
h.subsMu.Lock()
|
|
h.conns[c] = struct{}{}
|
|
h.subsMu.Unlock()
|
|
}
|
|
|
|
func (h *Hub) Unregister(c *websocket.Conn) {
|
|
h.subsMu.Lock()
|
|
delete(h.conns, c)
|
|
h.subsMu.Unlock()
|
|
}
|
|
|
|
// BroadcastSnapshot pushes full snapshot to all subscribers.
|
|
func (h *Hub) BroadcastSnapshot() {
|
|
body, err := h.BuildSnapshotMessage()
|
|
if err != nil {
|
|
return
|
|
}
|
|
h.subsMu.Lock()
|
|
defer h.subsMu.Unlock()
|
|
for c := range h.conns {
|
|
_ = c.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
|
if err := c.WriteMessage(websocket.TextMessage, body); err != nil {
|
|
_ = c.Close()
|
|
delete(h.conns, c)
|
|
}
|
|
}
|
|
}
|
|
|
|
// NotifyMetricsUpdate should be called after SetMetrics to push to browsers.
|
|
func (h *Hub) NotifyMetricsUpdate(serverID string, metricsJSON []byte) {
|
|
h.SetMetrics(serverID, metricsJSON)
|
|
h.BroadcastSnapshot()
|
|
}
|