69 lines
1.4 KiB
Go
69 lines
1.4 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
connections map[uuid.UUID][]*websocket.Conn
|
|
trackers map[uuid.UUID]context.CancelFunc
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
connections: make(map[uuid.UUID][]*websocket.Conn),
|
|
trackers: make(map[uuid.UUID]context.CancelFunc),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Register(orderID uuid.UUID, conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
h.connections[orderID] = append(h.connections[orderID], conn)
|
|
}
|
|
|
|
func (h *Hub) Unregister(orderID uuid.UUID, conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
conns := h.connections[orderID]
|
|
for i, c := range conns {
|
|
if c == conn {
|
|
h.connections[orderID] = append(conns[:i], conns[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Broadcast(orderID uuid.UUID, msg any) {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
for _, conn := range h.connections[orderID] {
|
|
conn.WriteJSON(msg)
|
|
}
|
|
}
|
|
|
|
func (h *Hub) StartTracker(orderID uuid.UUID, fn func(ctx context.Context)) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if _, exists := h.trackers[orderID]; exists {
|
|
return
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
h.trackers[orderID] = cancel
|
|
go fn(ctx)
|
|
}
|
|
|
|
func (h *Hub) StopTracker(orderID uuid.UUID) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if cancel, exists := h.trackers[orderID]; exists {
|
|
cancel()
|
|
delete(h.trackers, orderID)
|
|
}
|
|
}
|