95 lines
2.0 KiB
Go
95 lines
2.0 KiB
Go
package ws
|
|
|
|
import (
|
|
"log"
|
|
"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) {
|
|
if message == nil || message.SenderID == 0 {
|
|
log.Println("Attempt to broadcast invalid message")
|
|
return
|
|
}
|
|
|
|
if message.ReceiverID == 0 {
|
|
log.Printf("Message %d has no receiver", message.ID)
|
|
return
|
|
}
|
|
|
|
h.broadcast <- message
|
|
}
|
|
|
|
func (h *Hub) Run() {
|
|
for {
|
|
select {
|
|
case client := <-h.register:
|
|
h.mu.Lock()
|
|
// Закрываем предыдущее соединение если есть
|
|
if existing, ok := h.clients[client.UserID]; ok {
|
|
close(existing.Send)
|
|
}
|
|
h.clients[client.UserID] = client
|
|
h.mu.Unlock()
|
|
|
|
case client := <-h.unregister:
|
|
h.mu.Lock()
|
|
if c, ok := h.clients[client.UserID]; ok && c == client {
|
|
close(c.Send)
|
|
delete(h.clients, client.UserID)
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
case message := <-h.broadcast:
|
|
h.mu.RLock()
|
|
// Отправляем всем клиентам, кто участвует в этом чате
|
|
for _, client := range h.clients {
|
|
if client.UserID == message.SenderID || client.UserID == message.ReceiverID {
|
|
select {
|
|
case client.Send <- message:
|
|
default:
|
|
// Если канал полон, закрываем соединение
|
|
close(client.Send)
|
|
h.mu.RUnlock()
|
|
h.mu.Lock()
|
|
delete(h.clients, client.UserID)
|
|
h.mu.Unlock()
|
|
h.mu.RLock()
|
|
}
|
|
}
|
|
}
|
|
h.mu.RUnlock()
|
|
}
|
|
}
|
|
}
|