service/library/melody/hub.go

package melody

import (
	"service/library/logger"
	"sync"
)

type hub struct {
	sidSessionMap  map[string]*Session
	sessionMap     map[*Session]bool
	midSessionsMap map[int64]map[*Session]struct{} // mid: map[session]{}
	didSessionMap  map[string]*Session             // did: session
	broadcast      chan *CMsg
	register       chan *Session
	unregister     chan *Session
	exit           chan *CMsg
	open           bool
	rwmutex        *sync.RWMutex
}

func newHub() *hub {
	return &hub{
		sidSessionMap:  make(map[string]*Session),
		sessionMap:     make(map[*Session]bool),
		midSessionsMap: make(map[int64]map[*Session]struct{}, 1024*20),
		didSessionMap:  make(map[string]*Session),
		broadcast:      make(chan *CMsg),
		register:       make(chan *Session),
		unregister:     make(chan *Session),
		exit:           make(chan *CMsg),
		open:           true,
		rwmutex:        &sync.RWMutex{},
	}
}

func (h *hub) run() {
loop:
	for {
		select {
		case s := <-h.register:
			h.rwmutex.Lock()
			h.sessionMap[s] = true
			h.didSessionMap[s.Did] = s
			if h.midSessionsMap[s.Mid] == nil {
				h.midSessionsMap[s.Mid] = map[*Session]struct{}{s: {}}
			} else {
				h.midSessionsMap[s.Mid][s] = struct{}{}
			}
			h.rwmutex.Unlock()

			for mid, sMap := range h.midSessionsMap {
				logger.Info("reg, mid: %v, len(s): %v", mid, len(sMap))
			}
		case s := <-h.unregister:
			if _, ok := h.sessionMap[s]; ok {
				h.rwmutex.Lock()
				delete(h.sessionMap, s)
				delete(h.didSessionMap, s.Did)
				if len(h.midSessionsMap[s.Mid]) > 0 {
					delete(h.midSessionsMap[s.Mid], s)
					if len(h.midSessionsMap[s.Mid]) <= 0 {
						delete(h.midSessionsMap, s.Mid)
					}
				}
				h.rwmutex.Unlock()

				for mid, sMap := range h.midSessionsMap {
					logger.Info("unreg, mid: %v, len(s): %v", mid, len(sMap))
				}
			}
		case m := <-h.broadcast:
			h.rwmutex.RLock()
			for s := range h.sessionMap {
				if m.filter != nil {
					if m.filter(s) {
						s.writeMessage(m)
					}
				} else {
					s.writeMessage(m)
				}
			}
			h.rwmutex.RUnlock()
		case m := <-h.exit:
			h.rwmutex.Lock()
			for s := range h.sessionMap {
				s.writeMessage(m)
				delete(h.sessionMap, s)
				s.Close()
			}
			h.open = false
			h.rwmutex.Unlock()
			break loop
		}
	}
}

func (h *hub) closed() bool {
	h.rwmutex.RLock()
	defer h.rwmutex.RUnlock()
	return !h.open
}

func (h *hub) len() int {
	h.rwmutex.RLock()
	defer h.rwmutex.RUnlock()

	return len(h.sessionMap)
}

func (h *hub) all() []*Session {
	h.rwmutex.RLock()
	defer h.rwmutex.RUnlock()

	s := make([]*Session, 0, len(h.sessionMap))
	for k := range h.sessionMap {
		s = append(s, k)
	}
	return s
}

func (h *hub) updateSession(s *Session) {
	if len(s.Sid) > 0 {
		h.rwmutex.Lock()
		h.sidSessionMap[s.Sid] = s
		h.rwmutex.Unlock()
	}
	for i, v := range h.sidSessionMap {
		logger.Info("%d, %v", i, v.Sid)
	}
}