File size: 3,496 Bytes
a8eed5c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | package main
import (
"encoding/json"
"log"
"net/http"
"sync"
"time"
"github.com/gorilla/websocket"
)
// wsMessage is the envelope sent over WebSocket.
type wsMessage struct {
Type string `json:"type"`
Payload interface{} `json:"payload"`
}
// wsClient represents a connected WebSocket client.
type wsClient struct {
conn *websocket.Conn
send chan []byte
}
// Hub manages all connected WebSocket clients.
type Hub struct {
mu sync.RWMutex
clients map[*wsClient]struct{}
register chan *wsClient
unregister chan *wsClient
broadcast chan []byte
}
var hub = &Hub{
clients: make(map[*wsClient]struct{}),
register: make(chan *wsClient, 32),
unregister: make(chan *wsClient, 32),
broadcast: make(chan []byte, 256),
}
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool { return true },
ReadBufferSize: 1024,
WriteBufferSize: 4096,
}
// startHub runs the hub event loop.
func startHub() {
for {
select {
case c := <-hub.register:
hub.mu.Lock()
hub.clients[c] = struct{}{}
hub.mu.Unlock()
log.Printf("[WS] client connected (total=%d)", len(hub.clients))
// Send full state immediately on connect
go func() {
state := getNetworkState()
msg := wsMessage{Type: "NETWORK_UPDATE", Payload: state}
data, _ := json.Marshal(msg)
c.send <- data
}()
case c := <-hub.unregister:
hub.mu.Lock()
if _, ok := hub.clients[c]; ok {
delete(hub.clients, c)
close(c.send)
}
hub.mu.Unlock()
log.Printf("[WS] client disconnected (total=%d)", len(hub.clients))
case data := <-hub.broadcast:
hub.mu.RLock()
for c := range hub.clients {
select {
case c.send <- data:
default:
// Client too slow – drop message
}
}
hub.mu.RUnlock()
}
}
}
// handleWebSocket upgrades an HTTP connection to WebSocket.
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Printf("[WS] upgrade error: %v", err)
return
}
client := &wsClient{
conn: conn,
send: make(chan []byte, 64),
}
hub.register <- client
// Writer goroutine
go func() {
defer func() {
conn.Close()
}()
for data := range client.send {
conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
return
}
}
}()
// Reader goroutine (keep connection alive; ignore client messages)
defer func() {
hub.unregister <- client
}()
conn.SetReadLimit(512)
conn.SetReadDeadline(time.Now().Add(120 * time.Second))
conn.SetPongHandler(func(string) error {
conn.SetReadDeadline(time.Now().Add(120 * time.Second))
return nil
})
for {
_, _, err := conn.ReadMessage()
if err != nil {
break
}
// Reset deadline on any message
conn.SetReadDeadline(time.Now().Add(120 * time.Second))
}
}
// broadcastNetworkUpdate sends a full NETWORK_UPDATE event to all clients.
func broadcastNetworkUpdate() {
state := getNetworkState()
msg := wsMessage{Type: "NETWORK_UPDATE", Payload: state}
data, err := json.Marshal(msg)
if err != nil {
return
}
select {
case hub.broadcast <- data:
default:
}
}
// broadcastEvent sends a typed event to all connected WebSocket clients.
func broadcastEvent(eventType string, payload interface{}) {
msg := wsMessage{Type: eventType, Payload: payload}
data, err := json.Marshal(msg)
if err != nil {
return
}
select {
case hub.broadcast <- data:
default:
}
}
|