88 lines
1.7 KiB
Go
88 lines
1.7 KiB
Go
package ws
|
||
|
||
import (
|
||
"sync"
|
||
"tailly_back_v2/internal/domain"
|
||
)
|
||
|
||
type Client struct {
|
||
UserID int
|
||
Send chan *domain.Message
|
||
once sync.Once // Для гарантии однократного закрытия
|
||
}
|
||
|
||
func (c *Client) Close() {
|
||
c.once.Do(func() {
|
||
close(c.Send)
|
||
})
|
||
}
|
||
|
||
type ChatHub struct {
|
||
clients map[int]*Client
|
||
register chan *Client
|
||
unregister chan *Client
|
||
broadcast chan *domain.Message
|
||
mutex sync.RWMutex
|
||
}
|
||
|
||
func NewChatHub() *ChatHub {
|
||
return &ChatHub{
|
||
clients: make(map[int]*Client),
|
||
register: make(chan *Client),
|
||
unregister: make(chan *Client),
|
||
broadcast: make(chan *domain.Message),
|
||
}
|
||
}
|
||
|
||
// Register добавляет нового клиента в хаб
|
||
func (h *ChatHub) Register(client *Client) {
|
||
if h == nil || client == nil {
|
||
return
|
||
}
|
||
h.register <- client
|
||
}
|
||
|
||
func (h *ChatHub) Unregister(client *Client) {
|
||
if h == nil || client == nil {
|
||
return
|
||
}
|
||
h.unregister <- client
|
||
}
|
||
|
||
func (h *ChatHub) Broadcast(message *domain.Message) {
|
||
if h == nil || message == nil {
|
||
return
|
||
}
|
||
h.broadcast <- message
|
||
}
|
||
|
||
func (h *ChatHub) Run() {
|
||
for {
|
||
select {
|
||
case client := <-h.unregister:
|
||
h.mutex.Lock()
|
||
if c, ok := h.clients[client.UserID]; ok && c == client {
|
||
client.Close() // Используем безопасное закрытие
|
||
delete(h.clients, client.UserID)
|
||
}
|
||
h.mutex.Unlock()
|
||
|
||
case message := <-h.broadcast:
|
||
h.mutex.RLock()
|
||
if sender, ok := h.clients[message.SenderID]; ok {
|
||
select {
|
||
case sender.Send <- message: // Не блокируется, если канал закрыт
|
||
default:
|
||
}
|
||
}
|
||
if receiver, ok := h.clients[message.ReceiverID]; ok {
|
||
select {
|
||
case receiver.Send <- message:
|
||
default:
|
||
}
|
||
}
|
||
h.mutex.RUnlock()
|
||
}
|
||
}
|
||
}
|