318 lines
7.1 KiB
Go
318 lines
7.1 KiB
Go
package websocket
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RoomConfig 房间配置参数
|
|
type RoomConfig struct {
|
|
MaxConnections int // 最大连接数
|
|
MessageBuffer int // 消息缓冲区大小
|
|
PingInterval time.Duration // ping消息发送间隔
|
|
}
|
|
|
|
// DefaultRoomConfig 返回默认的房间配置
|
|
func DefaultRoomConfig() *RoomConfig {
|
|
return &RoomConfig{
|
|
MaxConnections: 100,
|
|
MessageBuffer: 256,
|
|
PingInterval: 30 * time.Second,
|
|
}
|
|
}
|
|
|
|
// BroadcastPayload 广播消息的负载结构
|
|
type BroadcastPayload struct {
|
|
Message *WSMessage // 要广播的消息
|
|
}
|
|
|
|
// DirectPayload 私信消息的负载结构
|
|
type DirectPayload struct {
|
|
To string // 接收者ID
|
|
Message *WSMessage // 要发送的消息
|
|
}
|
|
|
|
// 全局房间管理
|
|
var rooms = make(map[string]*Room)
|
|
var roomsMutex sync.Mutex
|
|
|
|
// DeleteRoom 从全局房间管理器中删除指定房间
|
|
func DeleteRoom(roomID string) {
|
|
roomsMutex.Lock()
|
|
defer roomsMutex.Unlock()
|
|
delete(rooms, roomID)
|
|
log.Printf("房间[%s]已从全局房间管理器删除", roomID)
|
|
}
|
|
|
|
// Room 表示一个聊天房间
|
|
type Room struct {
|
|
ID string // 房间ID
|
|
Connections sync.Map // 房间内的连接
|
|
OnlineUsers sync.Map // 在线用户
|
|
Register chan *Connection // 注册通道
|
|
Unregister chan *Connection // 注销通道
|
|
Broadcast chan *BroadcastPayload // 广播通道
|
|
Direct chan *DirectPayload // 私信通道
|
|
Quit chan struct{} // 退出通道
|
|
config *RoomConfig // 房间配置
|
|
stats *RoomStats // 房间统计信息
|
|
closed bool // 房间是否已关闭
|
|
mu sync.RWMutex // 读写锁
|
|
}
|
|
|
|
// RoomStats 房间统计信息
|
|
type RoomStats struct {
|
|
TotalMessages int64 // 总消息数
|
|
TotalConnections int64 // 总连接数
|
|
StartTime time.Time // 开始时间
|
|
mu sync.RWMutex // 读写锁
|
|
}
|
|
|
|
// NewRoomStats 创建新的房间统计对象
|
|
func NewRoomStats() *RoomStats {
|
|
return &RoomStats{
|
|
StartTime: time.Now(),
|
|
}
|
|
}
|
|
|
|
// IncrementMessages 增加消息计数
|
|
func (rs *RoomStats) IncrementMessages() {
|
|
rs.mu.Lock()
|
|
defer rs.mu.Unlock()
|
|
rs.TotalMessages++
|
|
}
|
|
|
|
// IncrementConnections 增加连接计数
|
|
func (rs *RoomStats) IncrementConnections() {
|
|
rs.mu.Lock()
|
|
defer rs.mu.Unlock()
|
|
rs.TotalConnections++
|
|
}
|
|
|
|
// GetStats 获取统计信息
|
|
func (rs *RoomStats) GetStats() map[string]interface{} {
|
|
rs.mu.RLock()
|
|
defer rs.mu.RUnlock()
|
|
return map[string]interface{}{
|
|
"total_messages": rs.TotalMessages,
|
|
"total_connections": rs.TotalConnections,
|
|
"uptime": time.Since(rs.StartTime).String(),
|
|
}
|
|
}
|
|
|
|
// NewRoom 创建新的房间
|
|
func NewRoom(id string, maxConnections int) *Room {
|
|
config := DefaultRoomConfig()
|
|
if maxConnections > 0 {
|
|
config.MaxConnections = maxConnections
|
|
}
|
|
|
|
r := &Room{
|
|
ID: id,
|
|
Register: make(chan *Connection),
|
|
Unregister: make(chan *Connection),
|
|
Broadcast: make(chan *BroadcastPayload),
|
|
Direct: make(chan *DirectPayload),
|
|
Quit: make(chan struct{}),
|
|
config: config,
|
|
stats: NewRoomStats(),
|
|
}
|
|
go r.Run()
|
|
|
|
go func() {
|
|
<-r.Quit
|
|
log.Printf("房间[%s]销毁成功", r.ID)
|
|
DeleteRoom(r.ID) // 从全局管理器删除
|
|
}()
|
|
|
|
return r
|
|
}
|
|
|
|
// Run 运行房间的主循环
|
|
func (r *Room) Run() {
|
|
ticker := time.NewTicker(r.config.PingInterval)
|
|
defer func() {
|
|
ticker.Stop()
|
|
r.cleanup()
|
|
}()
|
|
|
|
for {
|
|
select {
|
|
case conn := <-r.Register:
|
|
r.mu.RLock()
|
|
if r.closed {
|
|
r.mu.RUnlock()
|
|
return
|
|
}
|
|
r.mu.RUnlock()
|
|
|
|
log.Printf("客户端 %s 注册到房间 %s", conn.UID, r.ID)
|
|
if r.GetConnectionCount() >= r.config.MaxConnections {
|
|
log.Printf("房间 %s: 已达到最大连接数", r.ID)
|
|
conn.Send <- &WSMessage{
|
|
Type: "error",
|
|
Content: "房间已满",
|
|
}
|
|
conn.Conn.Close()
|
|
continue
|
|
}
|
|
r.Connections.Store(conn, true)
|
|
r.OnlineUsers.Store(conn.UID, true)
|
|
r.stats.IncrementConnections()
|
|
r.broadcastRoomInfo()
|
|
|
|
case conn := <-r.Unregister:
|
|
r.mu.RLock()
|
|
if r.closed {
|
|
r.mu.RUnlock()
|
|
return
|
|
}
|
|
r.mu.RUnlock()
|
|
|
|
log.Printf("客户端 %s 从房间 %s 注销", conn.UID, r.ID)
|
|
r.Connections.Delete(conn)
|
|
r.OnlineUsers.Delete(conn.UID)
|
|
r.broadcastRoomInfo()
|
|
|
|
case payload := <-r.Broadcast:
|
|
r.mu.RLock()
|
|
if r.closed {
|
|
r.mu.RUnlock()
|
|
return
|
|
}
|
|
r.mu.RUnlock()
|
|
|
|
log.Printf("在房间 %s 中广播消息: %s", r.ID, payload.Message.Content)
|
|
r.stats.IncrementMessages()
|
|
r.broadcastMessage(payload.Message)
|
|
|
|
case payload := <-r.Direct:
|
|
r.mu.RLock()
|
|
if r.closed {
|
|
r.mu.RUnlock()
|
|
return
|
|
}
|
|
r.mu.RUnlock()
|
|
|
|
log.Printf("发送私信给 %s: %s", payload.To, payload.Message.Content)
|
|
r.stats.IncrementMessages()
|
|
r.sendDirectMessage(payload)
|
|
|
|
case <-ticker.C:
|
|
r.mu.RLock()
|
|
if r.closed {
|
|
r.mu.RUnlock()
|
|
return
|
|
}
|
|
r.mu.RUnlock()
|
|
r.broadcastRoomInfo()
|
|
|
|
case <-r.Quit:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// broadcastMessage 广播消息给房间内所有用户
|
|
func (r *Room) broadcastMessage(msg *WSMessage) {
|
|
r.Connections.Range(func(key, value interface{}) bool {
|
|
conn := key.(*Connection)
|
|
log.Printf("广播消息给客户端 %s", conn.UID)
|
|
select {
|
|
case conn.Send <- msg:
|
|
log.Printf("消息已发送给客户端 %s", conn.UID)
|
|
default:
|
|
log.Printf("发送消息给客户端 %s 失败", conn.UID)
|
|
r.Connections.Delete(conn)
|
|
close(conn.Send)
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
|
|
// sendDirectMessage 发送私信给指定用户
|
|
func (r *Room) sendDirectMessage(payload *DirectPayload) {
|
|
// 获取hub实例
|
|
hub := GetHub()
|
|
if hub == nil {
|
|
log.Printf("无法获取hub实例")
|
|
return
|
|
}
|
|
|
|
// 使用hub发送跨房间私信
|
|
hub.SendDirectMessage(payload.Message.From, payload.To, payload.Message)
|
|
}
|
|
|
|
// broadcastRoomInfo 广播房间信息
|
|
func (r *Room) broadcastRoomInfo() {
|
|
info := map[string]interface{}{
|
|
"type": "room_info",
|
|
"room_id": r.ID,
|
|
"connections": r.GetConnectionCount(),
|
|
"online_users": r.GetOnlineUserCount(),
|
|
"stats": r.stats.GetStats(),
|
|
}
|
|
|
|
infoJSON, _ := json.Marshal(info)
|
|
msg := &WSMessage{
|
|
Type: "room_info",
|
|
Content: string(infoJSON),
|
|
Time: time.Now(),
|
|
}
|
|
r.broadcastMessage(msg)
|
|
}
|
|
|
|
// GetConnectionCount 获取房间内的连接数
|
|
func (r *Room) GetConnectionCount() int {
|
|
count := 0
|
|
r.Connections.Range(func(key, value interface{}) bool {
|
|
count++
|
|
return true
|
|
})
|
|
return count
|
|
}
|
|
|
|
// GetOnlineUserCount 获取房间内的在线用户数
|
|
func (r *Room) GetOnlineUserCount() int {
|
|
count := 0
|
|
r.OnlineUsers.Range(func(key, value interface{}) bool {
|
|
count++
|
|
return true
|
|
})
|
|
return count
|
|
}
|
|
|
|
// cleanup 清理房间资源
|
|
func (r *Room) cleanup() {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
|
|
if r.closed {
|
|
return
|
|
}
|
|
r.closed = true
|
|
|
|
r.Connections.Range(func(key, value interface{}) bool {
|
|
conn := key.(*Connection)
|
|
r.Connections.Delete(conn)
|
|
return true
|
|
})
|
|
|
|
close(r.Register)
|
|
close(r.Unregister)
|
|
close(r.Broadcast)
|
|
close(r.Direct)
|
|
close(r.Quit)
|
|
}
|
|
|
|
// Shutdown 关闭房间
|
|
func (r *Room) Shutdown() {
|
|
r.mu.Lock()
|
|
if !r.closed {
|
|
close(r.Quit)
|
|
}
|
|
r.mu.Unlock()
|
|
}
|