169 lines
3.9 KiB
Go
169 lines
3.9 KiB
Go
package realtime
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"github.com/coder/websocket"
|
||
)
|
||
|
||
// client 是一条 WebSocket 连接。读写各一个 goroutine，
|
||
// 写通过带缓冲的 send 通道串行化，通道满即视为慢客户端。
|
||
type client struct {
|
||
hub *Hub
|
||
conn *websocket.Conn
|
||
remoteAddr string
|
||
|
||
send chan []byte
|
||
done chan struct{}
|
||
|
||
closeOnce sync.Once
|
||
closeCode websocket.StatusCode
|
||
closeReason string
|
||
|
||
mu sync.Mutex
|
||
greeted bool
|
||
// role/userID 是 hello 之后实际生效的身份；sessionRole/sessionUserID 是握手时从会话取到的上限，
|
||
// 客户端自报的 role/userId 不能越过它。
|
||
role string
|
||
userID int64
|
||
sessionRole string
|
||
sessionUserID int64
|
||
instanceID string
|
||
|
||
ackVersion atomic.Int64 // 客户端已确认应用到的版本
|
||
lastSent atomic.Int64 // 服务端已发出的最大版本
|
||
slow atomic.Bool
|
||
}
|
||
|
||
func newClient(h *Hub, conn *websocket.Conn, remoteAddr, sessionRole string, sessionUserID int64) *client {
|
||
return &client{
|
||
hub: h,
|
||
conn: conn,
|
||
remoteAddr: remoteAddr,
|
||
send: make(chan []byte, defaultSendBuffer),
|
||
done: make(chan struct{}),
|
||
sessionRole: sessionRole,
|
||
sessionUserID: sessionUserID,
|
||
}
|
||
}
|
||
|
||
// stop 只负责让两条 goroutine 退出；真正的关闭动作由 serve 统一执行，
|
||
// 避免在持有广播锁时做可能阻塞的关闭握手。
|
||
func (c *client) stop(code websocket.StatusCode, reason string) {
|
||
c.closeOnce.Do(func() {
|
||
c.mu.Lock()
|
||
c.closeCode = code
|
||
c.closeReason = reason
|
||
c.mu.Unlock()
|
||
close(c.done)
|
||
})
|
||
}
|
||
|
||
func (c *client) sendEvent(event string, data any) bool {
|
||
msg, err := encode(event, data)
|
||
if err != nil {
|
||
c.hub.log.Error("编码实时消息失败", "event", event, "err", err)
|
||
return false
|
||
}
|
||
return c.sendRaw(msg)
|
||
}
|
||
|
||
func (c *client) sendRaw(msg []byte) bool {
|
||
select {
|
||
case <-c.done:
|
||
return false
|
||
default:
|
||
}
|
||
select {
|
||
case c.send <- msg:
|
||
return true
|
||
default:
|
||
c.hub.markSlow(c)
|
||
return false
|
||
}
|
||
}
|
||
|
||
// readLoop 解析 C→S 消息；首条消息受 helloTimeout 限制，防止空连接占用资源。
|
||
func (c *client) readLoop(ctx context.Context) {
|
||
defer c.stop(websocket.StatusNormalClosure, "")
|
||
|
||
for {
|
||
c.mu.Lock()
|
||
greeted := c.greeted
|
||
c.mu.Unlock()
|
||
|
||
readCtx := ctx
|
||
var cancel context.CancelFunc
|
||
if !greeted {
|
||
readCtx, cancel = context.WithTimeout(ctx, c.hub.helloTimeout())
|
||
}
|
||
_, data, err := c.conn.Read(readCtx)
|
||
if cancel != nil {
|
||
cancel()
|
||
}
|
||
if err != nil {
|
||
return
|
||
}
|
||
c.handleMessage(ctx, data)
|
||
}
|
||
}
|
||
|
||
func (c *client) handleMessage(ctx context.Context, data []byte) {
|
||
var env envelope
|
||
if err := json.Unmarshal(data, &env); err != nil {
|
||
c.hub.log.Debug("忽略无法解析的实时消息", "remote", c.remoteAddr, "err", err)
|
||
return
|
||
}
|
||
|
||
switch env.Event {
|
||
case EventHello:
|
||
var p helloPayload
|
||
if err := json.Unmarshal(env.Data, &p); err != nil {
|
||
c.hub.log.Debug("hello 载荷不合法", "remote", c.remoteAddr, "err", err)
|
||
return
|
||
}
|
||
c.hub.handleHello(ctx, c, p)
|
||
|
||
case EventAck:
|
||
var p ackPayload
|
||
if err := json.Unmarshal(env.Data, &p); err != nil {
|
||
c.hub.log.Debug("ack 载荷不合法", "remote", c.remoteAddr, "err", err)
|
||
return
|
||
}
|
||
if p.Version > c.ackVersion.Load() {
|
||
c.ackVersion.Store(p.Version)
|
||
}
|
||
c.hub.touchInstance(c, time.Now())
|
||
|
||
case EventResync:
|
||
c.hub.sendInitial(ctx, c, nil)
|
||
|
||
default:
|
||
c.hub.log.Debug("忽略未知实时事件", "event", env.Event, "remote", c.remoteAddr)
|
||
}
|
||
}
|
||
|
||
// writeLoop 是唯一的写者，保证同一连接上的帧顺序。
|
||
func (c *client) writeLoop(ctx context.Context) {
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
case <-c.done:
|
||
return
|
||
case msg := <-c.send:
|
||
writeCtx, cancel := context.WithTimeout(ctx, c.hub.writeTimeout())
|
||
err := c.conn.Write(writeCtx, websocket.MessageText, msg)
|
||
cancel()
|
||
if err != nil {
|
||
c.stop(websocket.StatusInternalError, "写入失败")
|
||
return
|
||
}
|
||
}
|
||
}
|
||
}
|