SFM-loan-Rest-Api/websocket_queue.go

package main

import (
	"github.com/gorilla/websocket"
	"sync"
)

type wsQ struct {
	wsMutex          sync.Mutex
	Clients          map[*websocket.Conn]struct{}
	ClientsBySession map[string]*websocket.Conn
}

var wsq = wsQ{
	Clients:          make(map[*websocket.Conn]struct{}),
	ClientsBySession: make(map[string]*websocket.Conn),
}

func (m *wsQ) addConnection(c *websocket.Conn) {
	m.wsMutex.Lock()
	m.Clients[c] = struct{}{}
	m.wsMutex.Unlock()
}

func (m *wsQ) MapMidToConnection(mid string, c *websocket.Conn) {
	m.wsMutex.Lock()
	m.ClientsBySession[mid] = c
	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.ClientsBySession {
		if k == sid {
			conn = v
			break
		}
	}
	m.wsMutex.Unlock()
	return
}

func (m *wsQ) del(conn *websocket.Conn) {
	m.wsMutex.Lock()
	delete(m.Clients, conn)
	for k, v := range m.ClientsBySession {
		if v == conn {
			delete(m.ClientsBySession, k)
			break
		}
	}
	m.wsMutex.Unlock()
}

func (m *wsQ) delBySessionId(sid string) {
	m.wsMutex.Lock()
	for k, v := range m.ClientsBySession {
		if k == sid {
			delete(m.ClientsBySession, k)
			break
		}
		delete(m.Clients, v)
	}
	m.wsMutex.Unlock()
}

func (m *wsQ) wsDoBroadcast(msg string, except string) {
	toBeRemoved := make([]*websocket.Conn, 0, 20)
	m.wsMutex.Lock()
	for k, _ := range m.Clients {
		if m.ClientsBySession[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()
}