131 lines
2.8 KiB
Go
131 lines
2.8 KiB
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)
|
|
}
|
|
}
|