mirror of
https://github.com/andwati/telecom-platform.git
synced 2026-10-18 11:02:54 +00:00
feat: Add WebSocket handler with hub management, client connections, and real-time broadcasting
- Add websocket.Upgrader with CheckOrigin allowing all origins in development - Add Message struct with type/data/timestamp fields for WebSocket messages - Add Client struct with conn/send/mutex/closed fields and readPump/writePump/Close methods - Add Hub struct with clients map, broadcast/register/unregister channels, and sync.RWMutex - Add NewHub constructor and Run method with client registration/
This commit is contained in:
nutcas3
@@ -0,0 +1,264 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true // Allow all origins in development
|
||||
},
|
||||
}
|
||||
|
||||
// Message represents a WebSocket message
|
||||
type Message struct {
|
||||
Type string `json:"type"`
|
||||
Data any `json:"data"`
|
||||
Timestamp string `json:"timestamp"`
|
||||
}
|
||||
|
||||
// Client represents a WebSocket client
|
||||
type Client struct {
|
||||
conn *websocket.Conn
|
||||
send chan Message
|
||||
mutex sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
// Hub manages WebSocket clients and broadcasts
|
||||
type Hub struct {
|
||||
clients map[*Client]bool
|
||||
broadcast chan Message
|
||||
register chan *Client
|
||||
unregister chan *Client
|
||||
mutex sync.RWMutex
|
||||
}
|
||||
|
||||
// NewHub creates a new WebSocket hub
|
||||
func NewHub() *Hub {
|
||||
return &Hub{
|
||||
clients: make(map[*Client]bool),
|
||||
broadcast: make(chan Message),
|
||||
register: make(chan *Client),
|
||||
unregister: make(chan *Client),
|
||||
}
|
||||
}
|
||||
|
||||
// Run starts the WebSocket hub
|
||||
func (h *Hub) Run() {
|
||||
for {
|
||||
select {
|
||||
case client := <-h.register:
|
||||
h.mutex.Lock()
|
||||
h.clients[client] = true
|
||||
h.mutex.Unlock()
|
||||
log.Printf("Client connected. Total clients: %d", len(h.clients))
|
||||
|
||||
case client := <-h.unregister:
|
||||
h.mutex.Lock()
|
||||
if _, ok := h.clients[client]; ok {
|
||||
delete(h.clients, client)
|
||||
close(client.send)
|
||||
}
|
||||
h.mutex.Unlock()
|
||||
log.Printf("Client disconnected. Total clients: %d", len(h.clients))
|
||||
|
||||
case message := <-h.broadcast:
|
||||
h.mutex.RLock()
|
||||
for client := range h.clients {
|
||||
select {
|
||||
case client.send <- message:
|
||||
default:
|
||||
close(client.send)
|
||||
delete(h.clients, client)
|
||||
}
|
||||
}
|
||||
h.mutex.RUnlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Broadcast sends a message to all connected clients
|
||||
func (h *Hub) Broadcast(messageType string, data any) {
|
||||
message := Message{
|
||||
Type: messageType,
|
||||
Data: data,
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339),
|
||||
}
|
||||
select {
|
||||
case h.broadcast <- message:
|
||||
default:
|
||||
log.Println("Broadcast channel is full, dropping message")
|
||||
}
|
||||
}
|
||||
|
||||
// GetClientCount returns the number of connected clients
|
||||
func (h *Hub) GetClientCount() int {
|
||||
h.mutex.RLock()
|
||||
defer h.mutex.RUnlock()
|
||||
return len(h.clients)
|
||||
}
|
||||
|
||||
// readPump handles messages from the WebSocket connection
|
||||
func (c *Client) readPump(h *Hub) {
|
||||
defer func() {
|
||||
h.unregister <- c
|
||||
c.conn.Close()
|
||||
}()
|
||||
|
||||
c.conn.SetReadLimit(512)
|
||||
c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
c.conn.SetPongHandler(func(string) error {
|
||||
c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
_, messageBytes, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||
log.Printf("WebSocket error: %v", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
var message Message
|
||||
if err := json.Unmarshal(messageBytes, &message); err != nil {
|
||||
log.Printf("Invalid message format: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle incoming messages if needed
|
||||
log.Printf("Received message: %s", message.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// writePump handles sending messages to the WebSocket connection
|
||||
func (c *Client) writePump() {
|
||||
ticker := time.NewTicker(54 * time.Second)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
c.conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case message, ok := <-c.send:
|
||||
c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||||
if !ok {
|
||||
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
||||
return
|
||||
}
|
||||
|
||||
if err := c.conn.WriteJSON(message); err != nil {
|
||||
log.Printf("WebSocket write error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
case <-ticker.C:
|
||||
c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
|
||||
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close safely closes the client connection
|
||||
func (c *Client) Close() {
|
||||
c.mutex.Lock()
|
||||
defer c.mutex.Unlock()
|
||||
|
||||
if !c.closed {
|
||||
c.closed = true
|
||||
close(c.send)
|
||||
c.conn.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// NewClient creates a new WebSocket client
|
||||
func NewClient(conn *websocket.Conn) *Client {
|
||||
return &Client{
|
||||
conn: conn,
|
||||
send: make(chan Message, 256),
|
||||
}
|
||||
}
|
||||
|
||||
// HandleWebSocket handles WebSocket connections
|
||||
func HandleWebSocket(hub *Hub) gin.HandlerFunc {
|
||||
return gin.HandlerFunc(func(c *gin.Context) {
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
log.Printf("WebSocket upgrade error: %v", err)
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "Failed to upgrade to WebSocket"})
|
||||
return
|
||||
}
|
||||
|
||||
client := NewClient(conn)
|
||||
hub.register <- client
|
||||
|
||||
// Start goroutines for reading and writing
|
||||
go client.readPump(hub)
|
||||
go client.writePump()
|
||||
})
|
||||
}
|
||||
|
||||
// Global hub instance
|
||||
var globalHub *Hub
|
||||
|
||||
// InitializeWebSocket initializes the WebSocket hub
|
||||
func InitializeWebSocket() {
|
||||
globalHub = NewHub()
|
||||
go globalHub.Run()
|
||||
}
|
||||
|
||||
// GetHub returns the global WebSocket hub
|
||||
func GetHub() *Hub {
|
||||
return globalHub
|
||||
}
|
||||
|
||||
// BroadcastServiceUpdate broadcasts a service update to all clients
|
||||
func BroadcastServiceUpdate(serviceName, status string, replicas, readyReplicas int) {
|
||||
if globalHub != nil {
|
||||
data := map[string]any{
|
||||
"service": serviceName,
|
||||
"status": status,
|
||||
"replicas": replicas,
|
||||
"readyReplicas": readyReplicas,
|
||||
}
|
||||
globalHub.Broadcast("service_update", data)
|
||||
}
|
||||
}
|
||||
|
||||
// BroadcastSystemMetrics broadcasts system metrics to all clients
|
||||
func BroadcastSystemMetrics(cpu, memory, disk, network float64) {
|
||||
if globalHub != nil {
|
||||
data := map[string]any{
|
||||
"cpu": cpu,
|
||||
"memory": memory,
|
||||
"disk": disk,
|
||||
"network": network,
|
||||
}
|
||||
globalHub.Broadcast("system_metrics", data)
|
||||
}
|
||||
}
|
||||
|
||||
// BroadcastAlertUpdate broadcasts an alert update to all clients
|
||||
func BroadcastAlertUpdate(id, severity, message, status string) {
|
||||
if globalHub != nil {
|
||||
data := map[string]any{
|
||||
"id": id,
|
||||
"severity": severity,
|
||||
"message": message,
|
||||
"status": status,
|
||||
}
|
||||
globalHub.Broadcast("alert_update", data)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user