145 lines
3.2 KiB
Go
145 lines
3.2 KiB
Go
package websocket
|
|
|
|
import (
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
// Hub 管理所有活动的房间
|
|
type Hub struct {
|
|
Rooms map[string]*Room // 房间映射表
|
|
Lock sync.RWMutex // 读写锁
|
|
MaxRooms int // 最大房间数
|
|
MaxRoomSize int // 每个房间的最大容量
|
|
}
|
|
|
|
// HubConfig Hub的配置参数
|
|
type HubConfig struct {
|
|
MaxRooms int // 最大房间数
|
|
MaxRoomSize int // 每个房间的最大容量
|
|
}
|
|
|
|
// DefaultHubConfig 返回默认的Hub配置
|
|
func DefaultHubConfig() *HubConfig {
|
|
return &HubConfig{
|
|
MaxRooms: 1000, // 默认最多1000个房间
|
|
MaxRoomSize: 100, // 每个房间默认最多100人
|
|
}
|
|
}
|
|
|
|
// NewHub 创建新的Hub实例
|
|
func NewHub(config *HubConfig) *Hub {
|
|
if config == nil {
|
|
config = DefaultHubConfig()
|
|
}
|
|
return &Hub{
|
|
Rooms: make(map[string]*Room),
|
|
MaxRooms: config.MaxRooms,
|
|
MaxRoomSize: config.MaxRoomSize,
|
|
}
|
|
}
|
|
|
|
// GetOrCreateRoom GetRoom 获取指定名称的房间，如果不存在则创建
|
|
func (h *Hub) GetOrCreateRoom(name string) *Room {
|
|
h.Lock.Lock()
|
|
defer h.Lock.Unlock()
|
|
|
|
room, ok := h.Rooms[name]
|
|
if !ok {
|
|
if len(h.Rooms) >= h.MaxRooms {
|
|
log.Printf("已达到最大房间数限制 (%d)", h.MaxRooms)
|
|
return nil
|
|
}
|
|
room = NewRoom(name, h.MaxRoomSize)
|
|
h.Rooms[name] = room
|
|
log.Printf("创建新房间: %s", name)
|
|
}
|
|
return room
|
|
}
|
|
|
|
// DeleteRoom 从Hub中删除指定房间
|
|
func (h *Hub) DeleteRoom(name string) {
|
|
h.Lock.Lock()
|
|
defer h.Lock.Unlock()
|
|
|
|
if room, exists := h.Rooms[name]; exists {
|
|
room.Shutdown()
|
|
delete(h.Rooms, name)
|
|
log.Printf("删除房间: %s", name)
|
|
}
|
|
}
|
|
|
|
// GetRoomCount 获取当前房间总数
|
|
func (h *Hub) GetRoomCount() int {
|
|
h.Lock.RLock()
|
|
defer h.Lock.RUnlock()
|
|
return len(h.Rooms)
|
|
}
|
|
|
|
// GetRoomInfo 获取所有房间的信息
|
|
func (h *Hub) GetRoomInfo() map[string]int {
|
|
h.Lock.RLock()
|
|
defer h.Lock.RUnlock()
|
|
|
|
info := make(map[string]int)
|
|
for name, room := range h.Rooms {
|
|
info[name] = room.GetConnectionCount()
|
|
}
|
|
return info
|
|
}
|
|
|
|
// Shutdown 优雅地关闭所有房间
|
|
func (h *Hub) Shutdown() {
|
|
h.Lock.Lock()
|
|
defer h.Lock.Unlock()
|
|
|
|
for name, room := range h.Rooms {
|
|
room.Shutdown()
|
|
delete(h.Rooms, name)
|
|
log.Printf("关闭房间: %s", name)
|
|
}
|
|
}
|
|
|
|
// SendDirectMessage 发送跨房间私信
|
|
func (h *Hub) SendDirectMessage(fromUID string, toUID string, message *WSMessage) {
|
|
h.Lock.RLock()
|
|
defer h.Lock.RUnlock()
|
|
|
|
// 确保消息的From字段被正确设置
|
|
message.From = fromUID
|
|
message.To = toUID
|
|
|
|
// 遍历所有房间查找接收者
|
|
for _, room := range h.Rooms {
|
|
room.Connections.Range(func(key, value interface{}) bool {
|
|
conn := key.(*Connection)
|
|
if conn.UID == toUID {
|
|
log.Printf("在房间 %s 中找到接收者 %s", room.ID, toUID)
|
|
select {
|
|
case conn.Send <- message:
|
|
log.Printf("私信已发送给客户端 %s", toUID)
|
|
default:
|
|
log.Printf("发送私信给客户端 %s 失败", toUID)
|
|
}
|
|
return false // 找到接收者后停止遍历
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
}
|
|
|
|
var (
|
|
// defaultHub 是默认的hub实例
|
|
defaultHub *Hub
|
|
)
|
|
|
|
// GetHub 返回默认的hub实例
|
|
func GetHub() *Hub {
|
|
return defaultHub
|
|
}
|
|
|
|
// SetHub 设置默认的hub实例
|
|
func SetHub(hub *Hub) {
|
|
defaultHub = hub
|
|
}
|