62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
package models
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type WebSocketManager struct {
|
|
connections map[string]*websocket.Conn
|
|
listeners map[string]map[chan struct{}]struct{}
|
|
mutex sync.RWMutex
|
|
}
|
|
|
|
func NewWebSocketManager() *WebSocketManager {
|
|
return &WebSocketManager{
|
|
connections: make(map[string]*websocket.Conn),
|
|
mutex: sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
func (mgr *WebSocketManager) AddConnection(conn *websocket.Conn) string {
|
|
mgr.mutex.Lock()
|
|
defer mgr.mutex.Unlock()
|
|
|
|
id := uuid.New().String() // 为每个连接生成一个唯一的 ID
|
|
mgr.connections[id] = conn
|
|
|
|
return id
|
|
}
|
|
|
|
func (mgr *WebSocketManager) RemoveConnection(id string) {
|
|
mgr.mutex.Lock()
|
|
defer mgr.mutex.Unlock()
|
|
delete(mgr.connections, id)
|
|
}
|
|
|
|
func (mgr *WebSocketManager) ListenForChanges(target string, callback func()) {
|
|
notifications := make(chan struct{})
|
|
mgr.mutex.Lock()
|
|
defer mgr.mutex.Unlock()
|
|
|
|
if _, ok := mgr.listeners[target]; !ok {
|
|
mgr.listeners[target] = make(map[chan struct{}]struct{})
|
|
}
|
|
mgr.listeners[target][notifications] = struct{}{}
|
|
|
|
go func() {
|
|
for {
|
|
callback()
|
|
for listener := range mgr.listeners[target] {
|
|
select {
|
|
case listener <- struct{}{}:
|
|
default:
|
|
delete(mgr.listeners[target], listener)
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|