@@ -0,0 +1,256 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type Event struct {
|
||||
Type string `json:"type"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
Data any `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
Conn *websocket.Conn
|
||||
TaskID string
|
||||
ClientID string
|
||||
All bool
|
||||
send chan []byte
|
||||
closed chan struct{}
|
||||
onWrite func(int)
|
||||
onDrop func()
|
||||
closeOnce sync.Once
|
||||
}
|
||||
|
||||
func NewClient(conn *websocket.Conn, taskID, clientID string, all bool) *Client {
|
||||
return &Client{Conn: conn, TaskID: taskID, ClientID: clientID, All: all, send: make(chan []byte, 32), closed: make(chan struct{})}
|
||||
}
|
||||
func (c *Client) start() { go c.writeLoop() }
|
||||
func (c *Client) writeLoop() {
|
||||
for {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return
|
||||
case b := <-c.send:
|
||||
_ = c.Conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
||||
if err := c.Conn.WriteMessage(websocket.TextMessage, b); err != nil {
|
||||
c.Close()
|
||||
return
|
||||
}
|
||||
if c.onWrite != nil {
|
||||
c.onWrite(len(b))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
func (c *Client) Enqueue(v any) bool {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return c.EnqueueBytes(b)
|
||||
}
|
||||
func (c *Client) EnqueueBytes(b []byte) bool {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return false
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case c.send <- b:
|
||||
return true
|
||||
default:
|
||||
if c.onDrop != nil {
|
||||
c.onDrop()
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
func (c *Client) Close() { c.closeOnce.Do(func() { close(c.closed); _ = c.Conn.Close() }) }
|
||||
|
||||
// Ping sends a real WebSocket control ping. Browsers answer this automatically
|
||||
// with a pong, which refreshes the server-side read deadline. Do not replace
|
||||
// this with an application JSON message: that caused every otherwise-idle
|
||||
// browser connection to be closed after the 90s read deadline.
|
||||
func (c *Client) Ping() error {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return context.Canceled
|
||||
default:
|
||||
}
|
||||
if err := c.Conn.WriteControl(websocket.PingMessage, []byte("nh"), time.Now().Add(5*time.Second)); err != nil {
|
||||
c.Close()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type pointUpdate struct {
|
||||
ClientID string
|
||||
Data any
|
||||
}
|
||||
type rateBucket struct {
|
||||
unix int64
|
||||
frames, bytes, drops uint64
|
||||
}
|
||||
|
||||
type Hub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[*Client]struct{}
|
||||
pendingMu sync.Mutex
|
||||
pending map[string]map[string]any
|
||||
frames atomic.Uint64
|
||||
bytes atomic.Uint64
|
||||
drops atomic.Uint64
|
||||
rateMu sync.Mutex
|
||||
rate [60]rateBucket
|
||||
}
|
||||
|
||||
func New() *Hub { return &Hub{clients: map[*Client]struct{}{}, pending: map[string]map[string]any{}} }
|
||||
func (h *Hub) Add(c *Client) {
|
||||
c.onWrite = h.recordWrite
|
||||
c.onDrop = h.recordDrop
|
||||
// Same-session websocket reconnects are intentionally allowed by runtime.
|
||||
// Keep only the newest user socket for a ClientID so a half-dead TCP
|
||||
// connection cannot continue receiving frames beside its replacement.
|
||||
var stale []*Client
|
||||
h.mu.Lock()
|
||||
if !c.All && c.ClientID != "" {
|
||||
for old := range h.clients {
|
||||
if !old.All && old.ClientID == c.ClientID {
|
||||
delete(h.clients, old)
|
||||
stale = append(stale, old)
|
||||
}
|
||||
}
|
||||
}
|
||||
h.clients[c] = struct{}{}
|
||||
h.mu.Unlock()
|
||||
for _, old := range stale {
|
||||
old.Close()
|
||||
}
|
||||
c.start()
|
||||
}
|
||||
func (h *Hub) Remove(c *Client) { h.mu.Lock(); delete(h.clients, c); h.mu.Unlock(); c.Close() }
|
||||
func (h *Hub) Connected() int { h.mu.RLock(); defer h.mu.RUnlock(); return len(h.clients) }
|
||||
|
||||
func (h *Hub) Publish(ctx context.Context, e Event) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
b, err := json.Marshal(e)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
h.broadcastBytes(e.TaskID, b)
|
||||
return nil
|
||||
}
|
||||
|
||||
// PublishPoint coalesces repeated improvements for the same client. Every 250ms
|
||||
// all changed points for a task are emitted as one frame per websocket instead
|
||||
// of one synchronous write per guess per client.
|
||||
func (h *Hub) PublishPoint(taskID, clientID string, data any) {
|
||||
h.pendingMu.Lock()
|
||||
m := h.pending[taskID]
|
||||
if m == nil {
|
||||
m = map[string]any{}
|
||||
h.pending[taskID] = m
|
||||
}
|
||||
m[clientID] = data
|
||||
h.pendingMu.Unlock()
|
||||
}
|
||||
func (h *Hub) flushPoints() {
|
||||
h.pendingMu.Lock()
|
||||
pending := h.pending
|
||||
h.pending = map[string]map[string]any{}
|
||||
h.pendingMu.Unlock()
|
||||
for taskID, m := range pending {
|
||||
arr := make([]any, 0, len(m))
|
||||
for _, v := range m {
|
||||
arr = append(arr, v)
|
||||
}
|
||||
b, err := json.Marshal(Event{Type: "points", TaskID: taskID, Data: arr})
|
||||
if err == nil {
|
||||
h.broadcastBytes(taskID, b)
|
||||
}
|
||||
}
|
||||
}
|
||||
func (h *Hub) broadcastBytes(taskID string, b []byte) {
|
||||
h.mu.RLock()
|
||||
clients := make([]*Client, 0, len(h.clients))
|
||||
for c := range h.clients {
|
||||
if c.All || taskID == "" || c.TaskID == taskID {
|
||||
clients = append(clients, c)
|
||||
}
|
||||
}
|
||||
h.mu.RUnlock()
|
||||
for _, c := range clients {
|
||||
c.EnqueueBytes(b)
|
||||
}
|
||||
}
|
||||
func (h *Hub) recordWrite(n int) {
|
||||
h.frames.Add(1)
|
||||
h.bytes.Add(uint64(n))
|
||||
now := time.Now().Unix()
|
||||
h.rateMu.Lock()
|
||||
i := now % 60
|
||||
if h.rate[i].unix != now {
|
||||
h.rate[i] = rateBucket{unix: now}
|
||||
}
|
||||
h.rate[i].frames++
|
||||
h.rate[i].bytes += uint64(n)
|
||||
h.rateMu.Unlock()
|
||||
}
|
||||
func (h *Hub) recordDrop() {
|
||||
h.drops.Add(1)
|
||||
now := time.Now().Unix()
|
||||
h.rateMu.Lock()
|
||||
i := now % 60
|
||||
if h.rate[i].unix != now {
|
||||
h.rate[i] = rateBucket{unix: now}
|
||||
}
|
||||
h.rate[i].drops++
|
||||
h.rateMu.Unlock()
|
||||
}
|
||||
|
||||
type Metrics struct {
|
||||
Connected int `json:"connected"`
|
||||
FramesPerSec float64 `json:"frames_per_sec"`
|
||||
BytesPerSec float64 `json:"bytes_per_sec"`
|
||||
DroppedPerSec float64 `json:"dropped_per_sec"`
|
||||
FramesTotal uint64 `json:"frames_total"`
|
||||
BytesTotal uint64 `json:"bytes_total"`
|
||||
DroppedTotal uint64 `json:"dropped_total"`
|
||||
}
|
||||
|
||||
func (h *Hub) Metrics() Metrics {
|
||||
now := time.Now().Unix()
|
||||
var f, b, d uint64
|
||||
h.rateMu.Lock()
|
||||
for _, x := range h.rate {
|
||||
if x.unix > now-5 {
|
||||
f += x.frames
|
||||
b += x.bytes
|
||||
d += x.drops
|
||||
}
|
||||
}
|
||||
h.rateMu.Unlock()
|
||||
return Metrics{Connected: h.Connected(), FramesPerSec: float64(f) / 5, BytesPerSec: float64(b) / 5, DroppedPerSec: float64(d) / 5, FramesTotal: h.frames.Load(), BytesTotal: h.bytes.Load(), DroppedTotal: h.drops.Load()}
|
||||
}
|
||||
func (h *Hub) Run(ctx context.Context) {
|
||||
t := time.NewTicker(250 * time.Millisecond)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
h.flushPoints()
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user