|
|
package conn
|
|
|
|
|
|
import (
|
|
|
"sync"
|
|
|
)
|
|
|
|
|
|
type ConnManager struct {
|
|
|
sync.RWMutex
|
|
|
|
|
|
// 分组
|
|
|
groups map[string]*Group
|
|
|
// 所有 Connection
|
|
|
conns map[int64]*Connection
|
|
|
}
|
|
|
|
|
|
func NewConnManager() *ConnManager {
|
|
|
return &ConnManager{
|
|
|
RWMutex: sync.RWMutex{},
|
|
|
groups: make(map[string]*Group),
|
|
|
conns: make(map[int64]*Connection),
|
|
|
}
|
|
|
}
|
|
|
|
|
|
// Store 保存连接，同时加入到指定分组，若给定分组名为空，则不进行分组操作
|
|
|
func (m *ConnManager) Store(groupName string, c *Connection) error {
|
|
|
m.Lock()
|
|
|
m.conns[c.Session().ID()] = c
|
|
|
m.Unlock()
|
|
|
|
|
|
group, ok := m.FindGroup(groupName)
|
|
|
if !ok {
|
|
|
group = m.NewGroup(groupName)
|
|
|
}
|
|
|
return group.Add(c)
|
|
|
}
|
|
|
|
|
|
func (m *ConnManager) Remove(c *Connection) error {
|
|
|
m.Lock()
|
|
|
defer m.Unlock()
|
|
|
delete(m.conns, c.Session().ID())
|
|
|
|
|
|
// 从所有group中移除
|
|
|
for _, group := range m.groups {
|
|
|
err := group.Leave(c)
|
|
|
if err != nil {
|
|
|
return err
|
|
|
}
|
|
|
}
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
func (m *ConnManager) RemoveFromGroup(groupName string, c *Connection) error {
|
|
|
m.Lock()
|
|
|
delete(m.conns, c.Session().ID())
|
|
|
m.Unlock()
|
|
|
|
|
|
group, ok := m.FindGroup(groupName)
|
|
|
if !ok {
|
|
|
return nil
|
|
|
}
|
|
|
|
|
|
return group.Leave(c)
|
|
|
}
|
|
|
|
|
|
// NewGroup 新增分组，若分组已存在，则返回现有分组
|
|
|
func (m *ConnManager) NewGroup(name string) *Group {
|
|
|
m.Lock()
|
|
|
defer m.Unlock()
|
|
|
|
|
|
group, ok := m.groups[name]
|
|
|
if ok {
|
|
|
return group
|
|
|
}
|
|
|
|
|
|
group = NewGroup(name)
|
|
|
m.groups[name] = group
|
|
|
|
|
|
return group
|
|
|
}
|
|
|
|
|
|
// FindGroup 查找分组
|
|
|
func (m *ConnManager) FindGroup(name string) (*Group, bool) {
|
|
|
m.RLock()
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
g, ok := m.groups[name]
|
|
|
return g, ok
|
|
|
}
|
|
|
|
|
|
// FindConn 根据连接ID找到连接
|
|
|
func (m *ConnManager) FindConn(id int64) (*Connection, bool) {
|
|
|
m.RLock()
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
e, ok := m.conns[id]
|
|
|
return e, ok
|
|
|
}
|
|
|
|
|
|
// FindConnByUID 根据连接绑定的UID找到连接
|
|
|
func (m *ConnManager) FindConnByUID(uid string) (*Connection, bool) {
|
|
|
m.RLock()
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
for _, e := range m.conns {
|
|
|
if e.Session().UID() == uid {
|
|
|
return e, true
|
|
|
}
|
|
|
}
|
|
|
return nil, false
|
|
|
}
|
|
|
|
|
|
// PeekConn 循环所有连接
|
|
|
// fn 返回true跳过循环，反之一直循环
|
|
|
func (m *ConnManager) PeekConn(fn func(id int64, c *Connection) bool) {
|
|
|
m.RLock()
|
|
|
defer m.RUnlock()
|
|
|
|
|
|
for id, c := range m.conns {
|
|
|
if fn(id, c) {
|
|
|
break
|
|
|
}
|
|
|
}
|
|
|
}
|