47 lines
1.0 KiB
Go
47 lines
1.0 KiB
Go
package models
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type WebSocketManager struct {
|
|
connections map[*websocket.Conn]string // 连接指针:任务ID
|
|
mutex sync.RWMutex
|
|
}
|
|
|
|
// 创建一个新的连接池
|
|
func NewWebSocketManager() *WebSocketManager {
|
|
return &WebSocketManager{
|
|
connections: make(map[*websocket.Conn]string),
|
|
mutex: sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
// 向连接池加入一个新连接
|
|
func (mgr *WebSocketManager) AddConnection(conn *websocket.Conn, task string) {
|
|
mgr.mutex.Lock()
|
|
defer mgr.mutex.Unlock()
|
|
mgr.connections[conn] = task
|
|
}
|
|
|
|
// 从连接池中移除一个连接
|
|
func (mgr *WebSocketManager) RemoveConnection(conn *websocket.Conn) {
|
|
mgr.mutex.Lock()
|
|
defer mgr.mutex.Unlock()
|
|
delete(mgr.connections, conn)
|
|
}
|
|
|
|
// 任务状态变化时, 向监听此任务的所有连接发送消息
|
|
func (mgr *WebSocketManager) NotifyTaskChange(task string, data interface{}) {
|
|
mgr.mutex.Lock()
|
|
defer mgr.mutex.Unlock()
|
|
|
|
for conn, value := range mgr.connections {
|
|
if value == task {
|
|
conn.WriteJSON(data)
|
|
}
|
|
}
|
|
}
|