service/library/melody/hub.go

package melody

import (
	"fmt"
	"sync"
)

type hub struct {
	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{
		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()
		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()
			}
		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) getSessionBySid(sid string) *Session {
	return h.sidSessionMap[sid]
}

func (h *hub) getSessionByMid(mid int64) *Session {
	return h.midSessionsMap[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")
}

func (h *hub) midSessions(mid int64) ([]*Session, error) {
	sids := h.midSidsMap[mid]
	if len(sids) > 0 {
		sessions := make([]*Session, 0)
		for sid := range sids {
			sessions = append(sessions, h.sidSessionMap[sid])
		}
		return sessions, nil
	}
	return nil, fmt.Errorf("get mid sessionMap fail")
}