SFM-loan-Rest-Api/websocket_queue.go

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()
}