133 lines
2.8 KiB
Go
133 lines
2.8 KiB
Go
package ws
|
|
|
|
import (
|
|
"encoding/json"
|
|
"sync"
|
|
)
|
|
|
|
// OnlineUser 在线成员信息
|
|
type OnlineUser struct {
|
|
UserID int64 `json:"user_id"`
|
|
Nickname string `json:"nickname"`
|
|
Avatar string `json:"avatar"`
|
|
}
|
|
|
|
// Room 按 family_id 分组的聊天室
|
|
type Room struct {
|
|
clients map[int64]*Client // user_id → Client
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// Hub 管理所有 Room
|
|
type Hub struct {
|
|
rooms map[int64]*Room // family_id → Room
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
var defaultHub = &Hub{
|
|
rooms: make(map[int64]*Room),
|
|
}
|
|
|
|
func GetHub() *Hub {
|
|
return defaultHub
|
|
}
|
|
|
|
// getOrCreateRoom 获取或创建 Room
|
|
func (h *Hub) getOrCreateRoom(familyID int64) *Room {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
room, ok := h.rooms[familyID]
|
|
if !ok {
|
|
room = &Room{
|
|
clients: make(map[int64]*Client),
|
|
}
|
|
h.rooms[familyID] = room
|
|
}
|
|
return room
|
|
}
|
|
|
|
// Join 用户加入房间，广播 online 事件
|
|
func (h *Hub) Join(familyID int64, client *Client) {
|
|
room := h.getOrCreateRoom(familyID)
|
|
|
|
room.mu.Lock()
|
|
room.clients[client.userID] = client
|
|
onlineUsers := h.collectOnlineUsers(room)
|
|
room.mu.Unlock()
|
|
|
|
h.broadcastOnlineUsers(familyID, onlineUsers)
|
|
}
|
|
|
|
// Leave 用户离开房间，广播 offline 事件
|
|
func (h *Hub) Leave(familyID int64, userID int64) {
|
|
room := h.getOrCreateRoom(familyID)
|
|
|
|
room.mu.Lock()
|
|
delete(room.clients, userID)
|
|
onlineUsers := h.collectOnlineUsers(room)
|
|
room.mu.Unlock()
|
|
|
|
// 广播 offline 给房间内其他人
|
|
msg, _ := json.Marshal(map[string]any{
|
|
"type": "offline",
|
|
"user_id": userID,
|
|
})
|
|
h.broadcastRaw(familyID, msg)
|
|
|
|
// 广播更新后的在线列表
|
|
h.broadcastOnlineUsers(familyID, onlineUsers)
|
|
}
|
|
|
|
// Broadcast 向房间内所有连接广播消息
|
|
func (h *Hub) Broadcast(familyID int64, msg any) {
|
|
data, err := json.Marshal(msg)
|
|
if err != nil {
|
|
return
|
|
}
|
|
h.broadcastRaw(familyID, data)
|
|
}
|
|
|
|
// broadcastRaw 向房间内所有连接广播原始字节
|
|
func (h *Hub) broadcastRaw(familyID int64, data []byte) {
|
|
h.mu.RLock()
|
|
room, ok := h.rooms[familyID]
|
|
h.mu.RUnlock()
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
room.mu.RLock()
|
|
defer room.mu.RUnlock()
|
|
|
|
for _, client := range room.clients {
|
|
select {
|
|
case client.send <- data:
|
|
default:
|
|
// 发送通道满则跳过
|
|
}
|
|
}
|
|
}
|
|
|
|
// collectOnlineUsers 收集在线用户列表（调用方需持有 room.mu 锁）
|
|
func (h *Hub) collectOnlineUsers(room *Room) []OnlineUser {
|
|
users := make([]OnlineUser, 0, len(room.clients))
|
|
for _, client := range room.clients {
|
|
users = append(users, OnlineUser{
|
|
UserID: client.userID,
|
|
Nickname: client.nickname,
|
|
Avatar: client.avatar,
|
|
})
|
|
}
|
|
return users
|
|
}
|
|
|
|
// broadcastOnlineUsers 广播在线用户列表
|
|
func (h *Hub) broadcastOnlineUsers(familyID int64, users []OnlineUser) {
|
|
msg, _ := json.Marshal(map[string]any{
|
|
"type": "online",
|
|
"users": users,
|
|
})
|
|
h.broadcastRaw(familyID, msg)
|
|
}
|