160 lines
3.4 KiB
Go
160 lines
3.4 KiB
Go
package handler
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
var upgrader = websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true // Allow all origins for development
|
|
},
|
|
}
|
|
|
|
// Client represents a WebSocket client
|
|
type Client struct {
|
|
hub *Hub
|
|
conn *websocket.Conn
|
|
send chan []byte
|
|
}
|
|
|
|
// Hub maintains the set of active clients and broadcasts messages to clients
|
|
type Hub struct {
|
|
clients map[*Client]bool
|
|
broadcast chan []byte
|
|
register chan *Client
|
|
unregister chan *Client
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// StockUpdateEvent represents a stock update notification
|
|
type StockUpdateEvent struct {
|
|
Type string `json:"type"` // "stock_update"
|
|
Action string `json:"action"` // "inbound" or "outbound"
|
|
OrderID int64 `json:"order_id"`
|
|
Data interface{} `json:"data,omitempty"`
|
|
}
|
|
|
|
var hubInstance *Hub
|
|
var hubOnce sync.Once
|
|
|
|
// GetHub returns the singleton Hub instance
|
|
func GetHub() *Hub {
|
|
hubOnce.Do(func() {
|
|
hubInstance = &Hub{
|
|
clients: make(map[*Client]bool),
|
|
broadcast: make(chan []byte, 256),
|
|
register: make(chan *Client),
|
|
unregister: make(chan *Client),
|
|
}
|
|
go hubInstance.run()
|
|
})
|
|
return hubInstance
|
|
}
|
|
|
|
func (h *Hub) run() {
|
|
for {
|
|
select {
|
|
case client := <-h.register:
|
|
h.mu.Lock()
|
|
h.clients[client] = true
|
|
h.mu.Unlock()
|
|
log.Printf("WebSocket client connected. Total clients: %d", len(h.clients))
|
|
|
|
case client := <-h.unregister:
|
|
h.mu.Lock()
|
|
if _, ok := h.clients[client]; ok {
|
|
delete(h.clients, client)
|
|
close(client.send)
|
|
}
|
|
h.mu.Unlock()
|
|
log.Printf("WebSocket client disconnected. Total clients: %d", len(h.clients))
|
|
|
|
case message := <-h.broadcast:
|
|
h.mu.RLock()
|
|
for client := range h.clients {
|
|
select {
|
|
case client.send <- message:
|
|
default:
|
|
close(client.send)
|
|
delete(h.clients, client)
|
|
}
|
|
}
|
|
h.mu.RUnlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
// BroadcastStockUpdate sends a stock update event to all connected clients
|
|
func (h *Hub) BroadcastStockUpdate(action string, orderID int64, data interface{}) {
|
|
event := StockUpdateEvent{
|
|
Type: "stock_update",
|
|
Action: action,
|
|
OrderID: orderID,
|
|
Data: data,
|
|
}
|
|
msg, err := json.Marshal(event)
|
|
if err != nil {
|
|
log.Printf("Error marshaling stock update event: %v", err)
|
|
return
|
|
}
|
|
h.broadcast <- msg
|
|
}
|
|
|
|
func (c *Client) readPump() {
|
|
defer func() {
|
|
c.hub.unregister <- c
|
|
c.conn.Close()
|
|
}()
|
|
for {
|
|
_, _, err := c.conn.ReadMessage()
|
|
if err != nil {
|
|
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
|
log.Printf("WebSocket error: %v", err)
|
|
}
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Client) writePump() {
|
|
defer func() {
|
|
c.conn.Close()
|
|
}()
|
|
for {
|
|
message, ok := <-c.send
|
|
if !ok {
|
|
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
|
return
|
|
}
|
|
if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// HandleWebSocket handles WebSocket connections at /ws
|
|
func HandleWebSocket(c *gin.Context) {
|
|
hub := GetHub()
|
|
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
|
if err != nil {
|
|
log.Printf("WebSocket upgrade failed: %v", err)
|
|
return
|
|
}
|
|
client := &Client{
|
|
hub: hub,
|
|
conn: conn,
|
|
send: make(chan []byte, 256),
|
|
}
|
|
hub.register <- client
|
|
go client.writePump()
|
|
go client.readPump()
|
|
}
|