service/library/melody/hub.go

package melody

import (
	"fmt"
	"sync"
)

type hub struct {
	sessions     map[*Session]bool
	sess         map[string]*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 *CMsg
	register     chan *Session
	unregister   chan *Session
	exit         chan *CMsg
	open         bool
	rwmutex      *sync.RWMutex
}

func newHub() *hub {
	return &hub{
		sessions:     make(map[*Session]bool),
		sess:         make(map[string]*Session),
		midsessions:  make(map[int64]*Session),
		didsessions:  make(map[string]*Session),
		roomsessions: make(map[int64]map[int64]bool),
		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.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
}

func (h *hub) getSessionBySid(sid string) *Session {
	return h.sess[sid]
}

func (h *hub) getSessionByMid(mid int64) *Session {
	return h.midsessions[mid]
}

func (h *hub) mustGetSession(sid string, mid int64) (*Session, error) {
	sessSid := h.getSessionBySid(sid)
	if sessSid != nil {
		return sessSid, nil
	}
	sessMid := h.getSessionByMid(mid)
	if sessMid != nil {
		return sessMid, nil
	}
	return nil, fmt.Errorf("get session fail")
}