package main import ( "github.com/brianvoe/gofakeit/v6" "github.com/gorilla/websocket" log "github.com/sirupsen/logrus" "sync" ) type wsQ struct { wsMutex sync.Mutex Clients map[*websocket.Conn]string ClientsBySocketId map[string]*websocket.Conn ClientToSessionId map[*websocket.Conn]string } var wsq = wsQ{ Clients: make(map[*websocket.Conn]string), ClientsBySocketId: make(map[string]*websocket.Conn), ClientToSessionId: make(map[*websocket.Conn]string), } func (m *wsQ) addConnection(c *websocket.Conn) (socketId string) { m.wsMutex.Lock() socketId = gofakeit.UUID() m.Clients[c] = socketId m.ClientsBySocketId[socketId] = c m.wsMutex.Unlock() m.debugNumOfClients() m.informClient(c, socketId) return } func (m *wsQ) informClient(conn *websocket.Conn, socketId string) { msg := make(map[string]string) msg["T"] = "assign-socketId" msg["socketId"] = socketId _ = wsSendByConn(conn, msg) } func (m *wsQ) debugNumOfClients() { m.wsMutex.Lock() numOfClients := len(m.Clients) numOfClientsById := len(m.ClientsBySocketId) m.wsMutex.Unlock() log.Info("number of concurrent websocket is ", numOfClients, numOfClientsById) } func (m *wsQ) MapSessionToConnection(sid string, c *websocket.Conn) { m.wsMutex.Lock() m.ClientToSessionId[c] = sid m.wsMutex.Unlock() } func (m *wsQ) RemoveSession(sessionId string, c *websocket.Conn) { m.wsMutex.Lock() for conn, v := range m.ClientToSessionId { if v == sessionId { delete(m.ClientToSessionId, conn) // remove connection to session map if socketId, exist := m.Clients[conn]; exist { delete(m.Clients, conn) delete(m.ClientsBySocketId, socketId) } } } m.wsMutex.Unlock() } func (m *wsQ) has(conn *websocket.Conn) (yes bool) { m.wsMutex.Lock() _, ok := m.Clients[conn] m.wsMutex.Unlock() return ok } func (m *wsQ) getConnBySessionId(sid string) (conn *websocket.Conn) { conn = nil m.wsMutex.Lock() for k, v := range m.ClientsBySocketId { if k == sid { conn = v break } } m.wsMutex.Unlock() return } func (m *wsQ) del(conn *websocket.Conn) { m.wsMutex.Lock() if sid, exist := m.Clients[conn]; exist { delete(m.Clients, conn) delete(m.ClientsBySocketId, sid) } m.wsMutex.Unlock() m.debugNumOfClients() } func (m *wsQ) delBySessionId(sid string) { m.wsMutex.Lock() if conn, exist := m.ClientsBySocketId[sid]; exist { delete(m.ClientsBySocketId, sid) delete(m.Clients, conn) } m.wsMutex.Unlock() m.debugNumOfClients() } func (m *wsQ) wsDoBroadcast(msg string, except string) { toBeRemoved := make([]*websocket.Conn, 0, 20) m.wsMutex.Lock() for k, _ := range m.Clients { if m.ClientsBySocketId[except] == k { continue } err := k.WriteMessage(websocket.TextMessage, []byte(msg)) if err != nil { toBeRemoved = append(toBeRemoved, k) } } // remove all those have errors for communication for _, v := range toBeRemoved { delete(m.Clients, v) } m.wsMutex.Unlock() }