nnet/conn/manager.go

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
		}
	}
}