package service import ( "log/slog" "sync" ) type WSMessage struct { Type string `json:"type"` Payload any `json:"payload,omitempty"` } type WSClient struct { id string send chan WSMessage done chan struct{} once sync.Once } func (client *WSClient) ID() string { if client == nil { return "" } return client.id } func (client *WSClient) Messages() <-chan WSMessage { if client == nil { return nil } return client.send } func (client *WSClient) Done() <-chan struct{} { if client == nil { return nil } return client.done } func (client *WSClient) Send(message WSMessage) bool { if client == nil { return false } select { case <-client.done: return false case client.send <- message: return true default: return false } } func (client *WSClient) Close() { if client == nil { return } client.once.Do(func() { close(client.done) }) } type WSHub struct { name string mu sync.RWMutex clients map[string]*WSClient } func NewWSHub(name string) *WSHub { return &WSHub{ name: name, clients: make(map[string]*WSClient), } } func (h *WSHub) Register(id string) *WSClient { client := &WSClient{ id: id, send: make(chan WSMessage, 16), done: make(chan struct{}), } h.mu.Lock() if existing := h.clients[id]; existing != nil { slog.Debug("ws replacing existing connection", "hub", h.name, "id", id) existing.Close() } h.clients[id] = client count := len(h.clients) h.mu.Unlock() slog.Debug("ws connection registered", "hub", h.name, "id", id, "client_count", count) return client } func (h *WSHub) Unregister(client *WSClient) { if client == nil { return } h.mu.Lock() if current := h.clients[client.id]; current == client { delete(h.clients, client.id) } count := len(h.clients) h.mu.Unlock() client.Close() slog.Debug("ws connection unregistered", "hub", h.name, "id", client.id, "client_count", count) } func (h *WSHub) Disconnect(id string) { h.mu.Lock() client := h.clients[id] if client != nil { delete(h.clients, id) } count := len(h.clients) h.mu.Unlock() if client != nil { client.Close() slog.Debug("ws connection forcefully disconnected", "hub", h.name, "id", id, "client_count", count) } } func (h *WSHub) IsConnected(id string) bool { h.mu.RLock() client := h.clients[id] h.mu.RUnlock() if client == nil { return false } select { case <-client.done: return false default: return true } } func (h *WSHub) SendMessage(id string, message WSMessage) bool { h.mu.RLock() client := h.clients[id] h.mu.RUnlock() if client == nil { return false } ok := client.Send(message) if !ok { slog.Debug("ws send queued message failed", "hub", h.name, "id", id, "type", message.Type) } return ok } type WSBroadcastResult struct { ClientCount int `json:"client_count"` SuccessCount int `json:"success_count"` FailedIDs []string `json:"failed_ids"` } func (h *WSHub) Broadcast(message WSMessage) WSBroadcastResult { h.mu.RLock() clients := make([]*WSClient, 0, len(h.clients)) for _, client := range h.clients { clients = append(clients, client) } h.mu.RUnlock() var result WSBroadcastResult result.ClientCount = len(clients) for _, client := range clients { if client.Send(message) { result.SuccessCount++ continue } result.FailedIDs = append(result.FailedIDs, client.ID()) } return result }