163 lines
3.6 KiB
Go
163 lines
3.6 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"
|
||
// OrganizationID 是本次库存变化所属公司，供按公司投递过滤（Task 5D 按它路由）。
|
||
OrganizationID int64 `json:"organization_id"`
|
||
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, organizationID, orderID int64, data interface{}) {
|
||
event := StockUpdateEvent{
|
||
Type: "stock_update",
|
||
Action: action,
|
||
OrganizationID: organizationID,
|
||
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()
|
||
}
|