nnet/conn/conn_mgr.go

package conn

import (
	"git.noahlan.cn/noahlan/nnet/entity"
	"sync"
)

type Manager struct {
	sync.RWMutex

	// 分组
	groups map[string]*Group
	// 所有 Connection
	conns map[int64]entity.NetworkEntity
}

func NewManager() *Manager {
	return &Manager{
		RWMutex: sync.RWMutex{},
		groups:  make(map[string]*Group),
		conns:   make(map[int64]entity.NetworkEntity),
	}
}

// Store 保存连接，同时加入到指定分组，若给定分组名为空，则不进行分组操作
func (m *Manager) Store(groupName string, s entity.NetworkEntity) error {
	m.Lock()
	m.conns[s.Session().ID()] = s
	m.Unlock()

	group, ok := m.FindGroup(groupName)
	if !ok {
		group = m.NewGroup(groupName)
	}
	return group.Add(s)
}

func (m *Manager) Remove(s entity.NetworkEntity) error {
	m.Lock()
	defer m.Unlock()
	delete(m.conns, s.Session().ID())

	// 从所有group中移除
	for _, group := range m.groups {
		err := group.Leave(s)
		if err != nil {
			return err
		}
	}
	return nil
}

func (m *Manager) RemoveFromGroup(groupName string, s entity.NetworkEntity) error {
	m.Lock()
	delete(m.conns, s.Session().ID())
	m.Unlock()

	group, ok := m.FindGroup(groupName)
	if !ok {
		return nil
	}

	return group.Leave(s)
}

// NewGroup 新增分组，若分组已存在，则返回现有分组
func (m *Manager) 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 *Manager) FindGroup(name string) (*Group, bool) {
	m.RLock()
	defer m.RUnlock()

	g, ok := m.groups[name]
	return g, ok
}

// FindConn 根据连接ID找到连接
func (m *Manager) FindConn(id int64) (entity.NetworkEntity, bool) {
	m.RLock()
	defer m.RUnlock()

	e, ok := m.conns[id]
	return e, ok
}

// FindConnByUID 根据连接绑定的UID找到连接
func (m *Manager) FindConnByUID(uid string) (entity.NetworkEntity, 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 *Manager) PeekConn(fn func(id int64, e entity.NetworkEntity) bool) {
	m.RLock()
	defer m.RUnlock()

	for id, e := range m.conns {
		if fn(id, e) {
			break
		}
	}
}