106 lines
2.1 KiB
Go
106 lines
2.1 KiB
Go
|
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
|
||
|
}
|