package melody import ( "sync" ) type hub struct { sessions map[*Session]bool sess map[int64]*Session // sid: session midsessions map[int64]*Session // mid: session didsessions map[string]*Session // did: session roomsessions map[int64]map[int64]bool // roomid: {sid: true} broadcast chan *envelope register chan *Session unregister chan *Session exit chan *envelope open bool rwmutex *sync.RWMutex } func newHub() *hub { return &hub{ sessions: make(map[*Session]bool), sess: make(map[int64]*Session), midsessions: make(map[int64]*Session), didsessions: make(map[string]*Session), roomsessions: make(map[int64]map[int64]bool), broadcast: make(chan *envelope), register: make(chan *Session), unregister: make(chan *Session), exit: make(chan *envelope), open: true, rwmutex: &sync.RWMutex{}, } } func (h *hub) run() { loop: for { select { case s := <-h.register: h.rwmutex.Lock() h.sessions[s] = true h.sess[s.Sid] = s h.midsessions[s.Mid] = s h.didsessions[s.Did] = s h.rwmutex.Unlock() case s := <-h.unregister: if _, ok := h.sessions[s]; ok { h.rwmutex.Lock() delete(h.sessions, s) delete(h.sess, s.Sid) delete(h.midsessions, s.Mid) delete(h.didsessions, s.Did) h.rwmutex.Unlock() } case m := <-h.broadcast: h.rwmutex.RLock() for s := range h.sessions { 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.sessions { s.writeMessage(m) delete(h.sessions, 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.sessions) } func (h *hub) all() []*Session { h.rwmutex.RLock() defer h.rwmutex.RUnlock() s := make([]*Session, 0, len(h.sessions)) for k := range h.sessions { s = append(s, k) } return s }