116 lines
2.3 KiB
Go
116 lines
2.3 KiB
Go
package websocket
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
type BroadcastPayload struct {
|
|
Message *WSMessage
|
|
}
|
|
|
|
type DirectPayload struct {
|
|
To string
|
|
Message *WSMessage
|
|
}
|
|
|
|
var rooms = make(map[string]*Room)
|
|
var roomsMutex sync.Mutex
|
|
|
|
func DeleteRoom(roomID string) {
|
|
roomsMutex.Lock()
|
|
defer roomsMutex.Unlock()
|
|
delete(rooms, roomID)
|
|
log.Printf("房间[%s]已从全局房间管理器删除", roomID)
|
|
}
|
|
|
|
type Room struct {
|
|
ID string
|
|
Connections map[*Connection]bool
|
|
OnlineUsers map[string]bool
|
|
Register chan *Connection
|
|
Unregister chan *Connection
|
|
Broadcast chan *BroadcastPayload
|
|
Direct chan *DirectPayload
|
|
Quit chan struct{}
|
|
}
|
|
|
|
func NewRoom(id string) *Room {
|
|
r := &Room{
|
|
ID: id,
|
|
Connections: make(map[*Connection]bool),
|
|
OnlineUsers: make(map[string]bool),
|
|
Register: make(chan *Connection),
|
|
Unregister: make(chan *Connection),
|
|
Broadcast: make(chan *BroadcastPayload),
|
|
Direct: make(chan *DirectPayload),
|
|
Quit: make(chan struct{}),
|
|
}
|
|
go r.Run()
|
|
|
|
go func() {
|
|
<-r.Quit
|
|
log.Printf("房间[%s]销毁成功", r.ID)
|
|
DeleteRoom(r.ID) // 从全局管理器删除
|
|
}()
|
|
|
|
return r
|
|
}
|
|
|
|
func (r *Room) Run() {
|
|
for {
|
|
select {
|
|
case conn := <-r.Register:
|
|
r.Connections[conn] = true
|
|
r.broadcastRoomInfo()
|
|
case conn := <-r.Unregister:
|
|
if _, ok := r.Connections[conn]; ok {
|
|
delete(r.Connections, conn)
|
|
close(conn.Send)
|
|
r.broadcastRoomInfo()
|
|
if len(r.Connections) == 0 {
|
|
log.Printf("房间[%s]没人了，准备销毁", r.ID)
|
|
close(r.Quit)
|
|
return
|
|
}
|
|
}
|
|
case payload := <-r.Broadcast:
|
|
for conn := range r.Connections {
|
|
select {
|
|
case conn.Send <- payload.Message:
|
|
default:
|
|
close(conn.Send)
|
|
delete(r.Connections, conn)
|
|
}
|
|
}
|
|
case payload := <-r.Direct:
|
|
for conn := range r.Connections {
|
|
if conn.UID == payload.To {
|
|
select {
|
|
case conn.Send <- payload.Message:
|
|
default:
|
|
close(conn.Send)
|
|
delete(r.Connections, conn)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (r *Room) broadcastRoomInfo() {
|
|
infoMsg := &WSMessage{
|
|
Type: "room_info",
|
|
Content: fmt.Sprintf("当前房间人数: %d", len(r.Connections)),
|
|
}
|
|
for conn := range r.Connections {
|
|
select {
|
|
case conn.Send <- infoMsg:
|
|
default:
|
|
close(conn.Send)
|
|
delete(r.Connections, conn)
|
|
}
|
|
}
|
|
}
|