80 lines
1.8 KiB
Go
80 lines
1.8 KiB
Go
package ws
|
|
|
|
import (
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type WebSocketPool struct {
|
|
conns map[string][]*websocket.Conn
|
|
sync.RWMutex
|
|
}
|
|
|
|
// 创建新的连接池
|
|
func NewWebSocketPool() *WebSocketPool {
|
|
return &WebSocketPool{
|
|
conns: make(map[string][]*websocket.Conn),
|
|
}
|
|
}
|
|
|
|
// 添加一个 WebSocket 连接到指定 appID 下
|
|
func (p *WebSocketPool) Add(appID string, conn *websocket.Conn) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
p.conns[appID] = append(p.conns[appID], conn)
|
|
}
|
|
|
|
// 从连接池中移除一个连接
|
|
func (p *WebSocketPool) Remove(appID string, conn *websocket.Conn) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
if conns, ok := p.conns[appID]; ok {
|
|
for i, c := range conns {
|
|
if c == conn {
|
|
p.conns[appID] = append(conns[:i], conns[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
if len(p.conns[appID]) == 0 {
|
|
delete(p.conns, appID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 向指定 appID 的所有连接广播消息
|
|
func (p *WebSocketPool) Broadcast(appID string, message string) {
|
|
p.RLock()
|
|
defer p.RUnlock()
|
|
if conns, ok := p.conns[appID]; ok {
|
|
for _, conn := range conns {
|
|
if err := conn.WriteMessage(websocket.TextMessage, []byte(message)); err != nil {
|
|
log.Printf("[WebSocketPool] failed to send to appID=%s: %v", appID, err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 清理掉线或不可用的连接（可定期调用）
|
|
func (p *WebSocketPool) CleanInactive(appID string) {
|
|
p.Lock()
|
|
defer p.Unlock()
|
|
if conns, ok := p.conns[appID]; ok {
|
|
active := make([]*websocket.Conn, 0, len(conns))
|
|
for _, conn := range conns {
|
|
if err := conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(1*time.Second)); err == nil {
|
|
active = append(active, conn)
|
|
} else {
|
|
_ = conn.Close()
|
|
}
|
|
}
|
|
if len(active) > 0 {
|
|
p.conns[appID] = active
|
|
} else {
|
|
delete(p.conns, appID)
|
|
}
|
|
}
|
|
}
|