1
0
mirror of https://github.com/nxshock/gron.git synced 2024-11-27 03:41:00 +05:00

Prevent parallel websocket writes

This commit is contained in:
nxshock 2024-04-03 13:47:20 +05:00
parent 252141ae01
commit 0c5aff496b
2 changed files with 33 additions and 5 deletions

View File

@ -16,7 +16,7 @@ import (
) )
type WsConnections struct { type WsConnections struct {
connections map[*websocket.Conn]struct{} connections map[*WsConnection]struct{}
mutex sync.Mutex mutex sync.Mutex
} }
@ -24,19 +24,24 @@ func (wc *WsConnections) Add(c *websocket.Conn) {
wc.mutex.Lock() wc.mutex.Lock()
defer wc.mutex.Unlock() defer wc.mutex.Unlock()
wc.connections[c] = struct{}{} wc.connections[NewWsConnection(c)] = struct{}{}
} }
func (wc *WsConnections) Delete(c *websocket.Conn) { func (wc *WsConnections) Delete(c *websocket.Conn) {
wc.mutex.Lock() wc.mutex.Lock()
defer wc.mutex.Unlock() defer wc.mutex.Unlock()
delete(wc.connections, c) for k := range wc.connections {
if k.w == c {
delete(wc.connections, k)
break
}
}
} }
func (wc *WsConnections) Send(message interface{}) { func (wc *WsConnections) Send(message interface{}) {
for conn := range wc.connections { for conn := range wc.connections {
go func(conn *websocket.Conn) { _ = conn.WriteJSON(message) }(conn) go func(conn *WsConnection) { _ = conn.Send(message) }(conn)
} }
} }
@ -49,7 +54,7 @@ var upgrader = websocket.Upgrader{
} }
var wsConnections = &WsConnections{ var wsConnections = &WsConnections{
connections: make(map[*websocket.Conn]struct{})} connections: make(map[*WsConnection]struct{})}
func httpServer(listenAddress string) { func httpServer(listenAddress string) {
if listenAddress == "none" { if listenAddress == "none" {

23
wconn.go Normal file
View File

@ -0,0 +1,23 @@
package main
import (
"sync"
"github.com/gorilla/websocket"
)
type WsConnection struct {
w *websocket.Conn
mu sync.Mutex
}
func NewWsConnection(w *websocket.Conn) *WsConnection {
return &WsConnection{w: w}
}
func (w *WsConnection) Send(message any) error {
w.mu.Lock()
defer w.mu.Unlock()
return w.w.WriteJSON(message)
}