68 lines
1.2 KiB
Go
68 lines
1.2 KiB
Go
package ws
|
|
|
|
import (
|
|
"sync"
|
|
)
|
|
|
|
var wsMgr = newManager()
|
|
|
|
type wsManager struct {
|
|
wsConnAll sync.Map // cmap.ConcurrentMap[uint32, *wsConnect]
|
|
}
|
|
|
|
func newManager() *wsManager {
|
|
return &wsManager{
|
|
//wsConnAll: cmap.NewWithCustomShardingFunction[uint32, *wsConnect](func(key uint32) uint32 {
|
|
// return key
|
|
//}),
|
|
}
|
|
}
|
|
|
|
func (m *wsManager) Add(conn *wsConnect) {
|
|
m.wsConnAll.Store(conn.id, conn)
|
|
}
|
|
|
|
func (m *wsManager) SetUserId(connId uint32, userId int64) {
|
|
userMgr.Add(connId, userId)
|
|
}
|
|
|
|
func (m *wsManager) Remove(conn *wsConnect) {
|
|
if conn.UserId() > 0 {
|
|
userMgr.Remove(conn.UserId())
|
|
}
|
|
m.wsConnAll.Delete(conn.id)
|
|
}
|
|
|
|
func (m *wsManager) Get(connId uint32) (*wsConnect, bool) {
|
|
v, ok := m.wsConnAll.Load(connId)
|
|
if ok {
|
|
conn, ok := v.(*wsConnect)
|
|
return conn, ok
|
|
}
|
|
return nil, false
|
|
}
|
|
|
|
func (m *wsManager) FindByUserId(userId int64) (*wsConnect, bool) {
|
|
connId := userMgr.GetConnId(userId)
|
|
return m.Get(connId)
|
|
}
|
|
|
|
func (m *wsManager) Rang(cb func(conn IConn) bool) {
|
|
userMgr.Rang(func(_, v any) bool {
|
|
connId := v.(uint32)
|
|
if conn, ok := m.Get(connId); ok {
|
|
return cb(conn)
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
|
|
func (m *wsManager) Count() int {
|
|
count := 0
|
|
m.wsConnAll.Range(func(k, v interface{}) bool {
|
|
count++
|
|
return true
|
|
})
|
|
return count
|
|
}
|