83 lines
1.6 KiB
Go
83 lines
1.6 KiB
Go
// ws/hub.go
|
|
package ws
|
|
|
|
import (
|
|
"sync"
|
|
"tailly_back_v2/internal/domain"
|
|
)
|
|
|
|
type Client struct {
|
|
UserID int
|
|
Send chan *domain.Message
|
|
}
|
|
|
|
type Hub struct {
|
|
clients map[int]*Client
|
|
register chan *Client
|
|
unregister chan *Client
|
|
broadcast chan *domain.Message
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
clients: make(map[int]*Client),
|
|
register: make(chan *Client),
|
|
unregister: make(chan *Client),
|
|
broadcast: make(chan *domain.Message, 100),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) RegisterClient(client *Client) {
|
|
h.register <- client
|
|
}
|
|
|
|
func (h *Hub) UnregisterClient(client *Client) {
|
|
h.unregister <- client
|
|
}
|
|
|
|
func (h *Hub) Broadcast(message *domain.Message) {
|
|
h.broadcast <- message
|
|
}
|
|
|
|
func (h *Hub) Run() {
|
|
for {
|
|
select {
|
|
case client := <-h.register:
|
|
h.mu.Lock()
|
|
h.clients[client.UserID] = client
|
|
h.mu.Unlock()
|
|
|
|
case client := <-h.unregister:
|
|
h.mu.Lock()
|
|
if c, ok := h.clients[client.UserID]; ok {
|
|
close(c.Send)
|
|
delete(h.clients, client.UserID)
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
case message := <-h.broadcast:
|
|
h.mu.RLock()
|
|
// Отправляем сообщение отправителю
|
|
if sender, ok := h.clients[message.SenderID]; ok {
|
|
select {
|
|
case sender.Send <- message:
|
|
default:
|
|
close(sender.Send)
|
|
delete(h.clients, sender.UserID)
|
|
}
|
|
}
|
|
// Отправляем сообщение получателю
|
|
if receiver, ok := h.clients[message.ReceiverID]; ok {
|
|
select {
|
|
case receiver.Send <- message:
|
|
default:
|
|
close(receiver.Send)
|
|
delete(h.clients, receiver.UserID)
|
|
}
|
|
}
|
|
h.mu.RUnlock()
|
|
}
|
|
}
|
|
}
|