nnet/connection/group.go

package connection

import (
	"errors"
	"git.noahlan.cn/noahlan/ntool/nlog"
	"sync"
	"sync/atomic"
)

const groupKey = "NNET_GROUP#"
const DefaultGroupName = "DEFAULT_GROUP"

const (
	groupStatusWorking = 0
	groupStatusClosed  = 1
)

var (
	ErrCloseClosedGroup        = errors.New("close closed group")
	ErrClosedGroup             = errors.New("group closed")
	DeleteDefaultGroupNotAllow = errors.New("delete default group not allow")
)

type Group struct {
	mu sync.RWMutex

	status int32  // group current status
	name   string // group name
	conns  map[int64]*Connection
}

func NewGroup(name string) *Group {
	return &Group{
		mu:     sync.RWMutex{},
		status: groupStatusWorking,
		name:   name,
		conns:  make(map[int64]*Connection),
	}
}

// Member returns connection by specified uid
func (g *Group) Member(uid string) (*Connection, bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	for _, e := range g.conns {
		if e.Session().UID() == uid {
			return e, true
		}
	}

	return nil, false
}

// MemberBySID returns specified sId's connection
func (g *Group) MemberBySID(id int64) (*Connection, bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

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

func (g *Group) Members() []*Connection {
	var resp []*Connection
	g.PeekMembers(func(_ int64, c *Connection) bool {
		resp = append(resp, c)
		return false
	})
	return resp
}

// PeekMembers returns all members in current group
// fn 返回true跳过循环，反之一直循环
func (g *Group) PeekMembers(fn func(sId int64, c *Connection) bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()

	for sId, c := range g.conns {
		if fn(sId, c) {
			break
		}
	}
}

// Contains check whether a UID is contained in current group or not
func (g *Group) Contains(uid string) bool {
	_, ok := g.Member(uid)
	return ok
}

// Add session to group
func (g *Group) Add(c *Connection) error {
	if g.isClosed() {
		return ErrClosedGroup
	}

	g.mu.Lock()
	defer g.mu.Unlock()

	sess := c.Session()
	id := sess.ID()

	// group attribute
	if sess.Exists(groupKey) {
		groups, ok := sess.Attribute(groupKey).([]string)
		if !ok {
			groups = make([]string, 0)
			sess.SetAttribute(groupKey, groups)
		}
		contains := false
		for _, group := range groups {
			if group == g.name {
				contains = true
				break
			}
		}
		if !contains {
			groups = append(groups, g.name)
			sess.SetAttribute(groupKey, groups)
		}
	} else {
		sess.SetAttribute(groupKey, []string{g.name})
	}

	if _, ok := g.conns[id]; !ok {
		g.conns[id] = c
	}

	nlog.Debugf("Add connection to group %s, ID=%d, UID=%s", g.name, sess.ID(), sess.UID())
	return nil
}

// Leave remove specified UID related session from group
func (g *Group) Leave(c *Connection) error {
	if g.isClosed() {
		return ErrClosedGroup
	}
	if c == nil {
		return nil
	}
	sess := c.Session()
	nlog.Debugf("Remove connection from group %s, UID=%s", g.name, sess.UID())

	g.mu.Lock()
	defer g.mu.Unlock()

	if sess.Exists(groupKey) {
		groups, ok := sess.Attribute(groupKey).([]string)
		if !ok {
			groups = make([]string, 0)
			sess.SetAttribute(groupKey, groups)
		}
		groups = g.removeGroupAttr(groups)

		if len(groups) == 0 {
			sess.RemoveAttribute(groupKey)
		} else {
			sess.SetAttribute(groupKey, groups)
		}
	}

	delete(g.conns, sess.ID())
	return nil
}

func (g *Group) LeaveByUID(uid string) error {
	if g.isClosed() {
		return ErrClosedGroup
	}
	member, _ := g.Member(uid)
	return g.Leave(member)
}

// LeaveAll clear all sessions in the group
func (g *Group) LeaveAll() error {
	if g.isClosed() {
		return ErrClosedGroup
	}

	g.mu.Lock()
	defer g.mu.Unlock()

	for _, e := range g.conns {
		sess := e.Session()

		groups, ok := sess.Attribute(groupKey).([]string)
		if !ok {
			groups = make([]string, 0)
			sess.SetAttribute(groupKey, groups)
		}
		groups = g.removeGroupAttr(groups)

		if len(groups) == 0 {
			sess.RemoveAttribute(groupKey)
		} else {
			sess.SetAttribute(groupKey, groups)
		}
	}
	g.conns = make(map[int64]*Connection)
	return nil
}

// 使用移位法移除group中与name匹配的元素
func (g *Group) removeGroupAttr(groups []string) []string {
	j := 0
	for _, v := range groups {
		if v != g.name {
			groups[j] = v
			j++
		}
	}
	return groups[:j]
}

// Count get current member amount in the group
func (g *Group) Count() int {
	g.mu.RLock()
	defer g.mu.RUnlock()

	return len(g.conns)
}

func (g *Group) isClosed() bool {
	if atomic.LoadInt32(&g.status) == groupStatusClosed {
		return true
	}
	return false
}

// Close destroy group, which will release all resource in the group
func (g *Group) Close() error {
	if g.isClosed() {
		return ErrCloseClosedGroup
	}
	if g.name == DefaultGroupName {
		// 默认分组不允许删除
		return DeleteDefaultGroupNotAllow
	}

	_ = g.LeaveAll()

	atomic.StoreInt32(&g.status, groupStatusClosed)
	return nil
}