72 lines
1.2 KiB
Go
Executable File
72 lines
1.2 KiB
Go
Executable File
package main
|
|
|
|
import (
|
|
"io"
|
|
"io/ioutil"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type Conn struct {
|
|
ws websocket.Conn
|
|
lock sync.Mutex
|
|
}
|
|
|
|
type Pool struct {
|
|
lock *sync.RWMutex
|
|
conns *sync.Map //map[string]*Conn
|
|
}
|
|
|
|
func NewPool() *Pool {
|
|
return &Pool{
|
|
conns: &sync.Map{},
|
|
lock: &sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
func (p *Pool) Push(id string, conn *websocket.Conn) {
|
|
p.lock.Lock()
|
|
p.conns.Store(id, &Conn{ws: *conn})
|
|
p.lock.Unlock()
|
|
p.Broadcast(websocket.TextMessage, strings.NewReader(`{"joined":"`+id+`", "uuid":"`+id+`"}`))
|
|
}
|
|
|
|
func (p *Pool) Broadcast(mt int, r io.Reader) error {
|
|
p.lock.RLock()
|
|
defer p.lock.RUnlock()
|
|
// io.MultiWriter exists but I like this
|
|
b, err := ioutil.ReadAll(r)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
n := 1000000
|
|
cnt := 0
|
|
p.conns.Range(func(k, v interface{}) bool {
|
|
k = k.(string)
|
|
conn := v.(*Conn)
|
|
lock := &conn.lock
|
|
lock.Lock()
|
|
defer lock.Unlock()
|
|
ws := &conn.ws
|
|
cnt += 1
|
|
w, err := ws.NextWriter(mt)
|
|
if err != nil {
|
|
p.conns.Delete(k)
|
|
return true
|
|
}
|
|
defer w.Close()
|
|
m, err := w.Write(b)
|
|
if err != nil {
|
|
p.conns.Delete(k)
|
|
return true
|
|
}
|
|
if m < n {
|
|
n = m
|
|
}
|
|
return true
|
|
})
|
|
return nil
|
|
}
|