nnet/internal/hub.go

package internal

import (
	"math/rand"
	"one.com/nnet"
	"sync"
	"time"
)

type Hub struct {
	sync.Mutex
	conf   *nnet.HubConfig
	cbSes  nnet.ISessionCallback
	prot   nnet.IProtocol
	chQuit chan struct{}
	wg     *sync.WaitGroup

	sess map[uint64]nnet.ISession
}

func newHub(cf *nnet.HubConfig, cb nnet.ISessionCallback, p nnet.IProtocol) *Hub {
	return &Hub{
		conf:   cf,
		cbSes:  cb,
		prot:   p,
		chQuit: make(chan struct{}),
		wg:     &sync.WaitGroup{},
		sess:   make(map[uint64]nnet.ISession),
	}
}

func (self *Hub) Wg() *sync.WaitGroup {
	return self.wg
}

func (self *Hub) ChQuit() <-chan struct{} {
	return self.chQuit
}

func (self *Hub) Conf() *nnet.HubConfig {
	return self.conf
}

func (self *Hub) Callback() nnet.ISessionCallback {
	return self.cbSes
}

func (self *Hub) Protocol() nnet.IProtocol {
	return self.prot
}

func (self *Hub) PutSession(id uint64, ses nnet.ISession) error {
	self.Lock()
	//@Notice: 顶替
	self.sess[id] = ses
	self.Unlock()
	return nil
}

func (self *Hub) DelSession(id uint64) error {
	self.Lock()
	defer self.Unlock()

	if _, ok := self.sess[id]; !ok {
		return nnet.ErrNotExists
	}
	delete(self.sess, id)
	return nil
}

func (self *Hub) PeekSession(id uint64) (nnet.ISession, error) {
	self.Lock()
	defer self.Unlock()

	s, ok := self.sess[id]
	if !ok {
		return nil, nnet.ErrNotExists
	}

	delete(self.sess, id)

	return s, nil
}

func (self *Hub) GetSession(id uint64) (nnet.ISession, error) {
	self.Lock()
	defer self.Unlock()

	s, ok := self.sess[id]
	if !ok {
		return nil, nnet.ErrNotExists
	}
	return s, nil
}

func (self *Hub) GetAllSessions() map[uint64]nnet.ISession {
	self.Lock()
	defer self.Unlock()

	return self.sess
}

var (
	s = rand.NewSource(time.Now().UnixNano())
	r = rand.New(s)
)

func Intn(a int) int {
	return r.Intn(a)
}

func (self *Hub) RandSession() (nnet.ISession, error) {
	self.Lock()
	defer self.Unlock()

	sz := len(self.sess)

	sel := Intn(sz)
	counter := 0
	for _, ses := range self.sess {
		if counter == sel {
			return ses, nil
		}
		counter += 1
	}
	return nil, nnet.ErrNotExists
}

func (self *Hub) GetSessionNum() int {
	self.Lock()
	defer self.Unlock()
	return len(self.sess)
}