74 lines
1.5 KiB
Go
74 lines
1.5 KiB
Go
package ws
|
|
|
|
import (
|
|
"log"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type Client struct {
|
|
Conn *websocket.Conn
|
|
RecipientID int64
|
|
}
|
|
|
|
type NotificationHub struct {
|
|
Clients map[*Client]bool
|
|
Broadcast chan interface{}
|
|
Register chan *Client
|
|
Unregister chan *Client
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewNotificationHub() *NotificationHub {
|
|
return &NotificationHub{
|
|
Clients: make(map[*Client]bool),
|
|
Broadcast: make(chan interface{}, 1000),
|
|
Register: make(chan *Client),
|
|
Unregister: make(chan *Client),
|
|
}
|
|
}
|
|
|
|
func (h *NotificationHub) Run() {
|
|
for {
|
|
select {
|
|
case client := <-h.Register:
|
|
h.mu.Lock()
|
|
h.Clients[client] = true
|
|
h.mu.Unlock()
|
|
log.Printf("Client registered: %d", client.RecipientID)
|
|
case client := <-h.Unregister:
|
|
h.mu.Lock()
|
|
if _, ok := h.Clients[client]; ok {
|
|
delete(h.Clients, client)
|
|
client.Conn.Close()
|
|
}
|
|
h.mu.Unlock()
|
|
log.Printf("Client unregistered: %d", client.RecipientID)
|
|
case message := <-h.Broadcast:
|
|
h.mu.Lock()
|
|
for client := range h.Clients {
|
|
if payload, ok := message.(map[string]interface{}); ok {
|
|
if recipient, ok := payload["recipient_id"].(int64); ok && recipient == client.RecipientID {
|
|
err := client.Conn.WriteJSON(payload)
|
|
if err != nil {
|
|
delete(h.Clients, client)
|
|
client.Conn.Close()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
h.mu.Unlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
var Upgrader = websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true
|
|
},
|
|
}
|