578 lines
14 KiB
Go
578 lines
14 KiB
Go
package realtime
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"log/slog"
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/coder/websocket"
|
||
|
||
"git.yeij.top/AskaEth/Scoreboard/server/internal/auth"
|
||
"git.yeij.top/AskaEth/Scoreboard/server/internal/middleware"
|
||
)
|
||
|
||
// 默认参数；现场可通过 Set* 调整。
|
||
const (
|
||
defaultHeartbeatInterval = 10 * time.Second
|
||
defaultMaxCatchup = 500
|
||
defaultSlowLag = 500
|
||
defaultHelloTimeout = 15 * time.Second
|
||
defaultWriteTimeout = 10 * time.Second
|
||
defaultSendBuffer = 64
|
||
)
|
||
|
||
// Hub 是实时同步中心：维护房间、连接、大屏实例注册表与版本号。
|
||
//
|
||
// 版本号等于 change_log 的最大 id；写入方提交事务后调用 Broadcast 推送变更，
|
||
// 断线客户端用 hello.lastVersion 请求增量补齐。
|
||
type Hub struct {
|
||
db *sql.DB
|
||
log *slog.Logger
|
||
|
||
mu sync.RWMutex
|
||
rooms map[string]map[*client]struct{}
|
||
clients map[*client]struct{}
|
||
instances map[string]*instanceState
|
||
version int64
|
||
closed bool
|
||
|
||
// sendMu 串行化「版本推进 + 下发」，保证同一客户端收到的 change 与版本严格递增，
|
||
// 也避免初始快照与并发广播交错导致漏变更。
|
||
sendMu sync.Mutex
|
||
|
||
ctx context.Context
|
||
cancel context.CancelFunc
|
||
hbWG sync.WaitGroup // 心跳循环
|
||
connWG sync.WaitGroup // 连接处理
|
||
|
||
heartbeatInterval time.Duration
|
||
heartbeatReset chan struct{}
|
||
maxCatchup int
|
||
slowLag int64
|
||
helloTO time.Duration
|
||
writeTO time.Duration
|
||
}
|
||
|
||
// NewHub 创建实时中心并启动心跳循环；db 是赛事库（change_log 所在库）。
|
||
// 调用方在服务关闭时应调用 Hub.Close。
|
||
func NewHub(db *sql.DB, log *slog.Logger) *Hub {
|
||
if log == nil {
|
||
log = slog.Default()
|
||
}
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
h := &Hub{
|
||
db: db,
|
||
log: log,
|
||
rooms: make(map[string]map[*client]struct{}),
|
||
clients: make(map[*client]struct{}),
|
||
instances: make(map[string]*instanceState),
|
||
ctx: ctx,
|
||
cancel: cancel,
|
||
heartbeatInterval: defaultHeartbeatInterval,
|
||
heartbeatReset: make(chan struct{}, 1),
|
||
maxCatchup: defaultMaxCatchup,
|
||
slowLag: defaultSlowLag,
|
||
helloTO: defaultHelloTimeout,
|
||
writeTO: defaultWriteTimeout,
|
||
}
|
||
if db != nil {
|
||
if v, err := h.currentVersion(context.Background()); err == nil {
|
||
h.version = v
|
||
} else {
|
||
log.Warn("读取初始版本号失败", "err", err)
|
||
}
|
||
}
|
||
|
||
h.hbWG.Add(1)
|
||
go h.heartbeatLoop()
|
||
return h
|
||
}
|
||
|
||
// Close 停止心跳、断开所有连接并等待处理中的连接退出。
|
||
func (h *Hub) Close() error {
|
||
h.mu.Lock()
|
||
if h.closed {
|
||
h.mu.Unlock()
|
||
return nil
|
||
}
|
||
h.closed = true
|
||
h.mu.Unlock()
|
||
|
||
h.cancel()
|
||
for _, c := range h.roomClients("") {
|
||
c.stop(websocket.StatusGoingAway, "服务关闭")
|
||
}
|
||
h.connWG.Wait()
|
||
h.hbWG.Wait()
|
||
return nil
|
||
}
|
||
|
||
// Version 返回内存中记录的当前版本号。
|
||
func (h *Hub) Version() int64 {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return h.version
|
||
}
|
||
|
||
func (h *Hub) setVersion(v int64) {
|
||
if v <= 0 {
|
||
return
|
||
}
|
||
h.mu.Lock()
|
||
if v > h.version {
|
||
h.version = v
|
||
}
|
||
h.mu.Unlock()
|
||
}
|
||
|
||
// SetHeartbeatInterval 调整心跳间隔并立即生效；用于现场调优与测试。
|
||
func (h *Hub) SetHeartbeatInterval(d time.Duration) {
|
||
if d <= 0 {
|
||
return
|
||
}
|
||
h.mu.Lock()
|
||
h.heartbeatInterval = d
|
||
h.mu.Unlock()
|
||
|
||
select {
|
||
case h.heartbeatReset <- struct{}{}:
|
||
default:
|
||
}
|
||
}
|
||
|
||
// SetMaxCatchup 设置增量补发的条数上限；超过就回落全量快照。
|
||
func (h *Hub) SetMaxCatchup(n int) {
|
||
if n <= 0 {
|
||
return
|
||
}
|
||
h.mu.Lock()
|
||
h.maxCatchup = n
|
||
h.mu.Unlock()
|
||
}
|
||
|
||
// SetSlowLag 设置「慢客户端」的版本落后阈值；<=0 表示不判定。
|
||
func (h *Hub) SetSlowLag(n int64) {
|
||
h.mu.Lock()
|
||
h.slowLag = n
|
||
h.mu.Unlock()
|
||
}
|
||
|
||
func (h *Hub) heartbeatIntervalValue() time.Duration {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return h.heartbeatInterval
|
||
}
|
||
|
||
func (h *Hub) maxCatchupValue() int {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return h.maxCatchup
|
||
}
|
||
|
||
func (h *Hub) slowLagValue() int64 {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return h.slowLag
|
||
}
|
||
|
||
func (h *Hub) helloTimeout() time.Duration {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return h.helloTO
|
||
}
|
||
|
||
func (h *Hub) writeTimeout() time.Duration {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return h.writeTO
|
||
}
|
||
|
||
// Handler 返回 WebSocket 处理器，由调用方挂到 /ws 上。
|
||
func (h *Hub) Handler() http.Handler {
|
||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
|
||
// 大屏常以 IP/别名访问且与管理端同端口，Origin 校验交给反向代理与鉴权层。
|
||
InsecureSkipVerify: true,
|
||
CompressionMode: websocket.CompressionDisabled,
|
||
})
|
||
if err != nil {
|
||
h.log.Debug("WebSocket 握手失败", "remote", r.RemoteAddr, "err", err)
|
||
return
|
||
}
|
||
|
||
h.mu.Lock()
|
||
if h.closed {
|
||
h.mu.Unlock()
|
||
_ = conn.CloseNow()
|
||
return
|
||
}
|
||
h.connWG.Add(1)
|
||
h.mu.Unlock()
|
||
defer h.connWG.Done()
|
||
|
||
// 身份只认会话：/ws 必须被 Authenticate 中间件包住，否则一律按 guest 处理。
|
||
role, userID := wsIdentity(r.Context())
|
||
h.serve(conn, r.RemoteAddr, role, userID)
|
||
})
|
||
}
|
||
|
||
// wsIdentity 从请求上下文取会话身份并折算成实时层角色。
|
||
// 未登录（guest）只能进大屏房间；recorder/admin 才有管理端房间的订阅权。
|
||
func wsIdentity(ctx context.Context) (string, int64) {
|
||
u := middleware.User(ctx)
|
||
if u == nil || !u.Role.Allows(auth.RoleRecorder) {
|
||
return RoleDisplay, 0
|
||
}
|
||
return RoleAdmin, u.ID
|
||
}
|
||
|
||
func (h *Hub) serve(conn *websocket.Conn, remoteAddr, sessionRole string, sessionUserID int64) {
|
||
conn.SetReadLimit(1 << 20)
|
||
|
||
c := newClient(h, conn, remoteAddr, sessionRole, sessionUserID)
|
||
h.addClient(c)
|
||
|
||
ctx, cancel := context.WithCancel(h.ctx)
|
||
defer cancel()
|
||
|
||
var loops sync.WaitGroup
|
||
loops.Add(2)
|
||
go func() { defer loops.Done(); c.readLoop(ctx) }()
|
||
go func() { defer loops.Done(); c.writeLoop(ctx) }()
|
||
|
||
<-c.done
|
||
cancel()
|
||
loops.Wait()
|
||
h.removeClient(c)
|
||
// 客户端会自动重连并走 catchup，这里直接断开，避免关闭握手阻塞。
|
||
_ = conn.CloseNow()
|
||
}
|
||
|
||
func (h *Hub) addClient(c *client) {
|
||
h.mu.Lock()
|
||
h.clients[c] = struct{}{}
|
||
h.mu.Unlock()
|
||
}
|
||
|
||
func (h *Hub) removeClient(c *client) {
|
||
h.mu.Lock()
|
||
delete(h.clients, c)
|
||
for room, set := range h.rooms {
|
||
delete(set, c)
|
||
if len(set) == 0 {
|
||
delete(h.rooms, room)
|
||
}
|
||
}
|
||
h.mu.Unlock()
|
||
|
||
h.releaseDisplay(c)
|
||
}
|
||
|
||
// join 把连接加入房间；房间用内存 map 维护，不引入实时框架。
|
||
func (h *Hub) join(c *client, rooms ...string) {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
for _, room := range rooms {
|
||
if room == "" {
|
||
continue
|
||
}
|
||
set := h.rooms[room]
|
||
if set == nil {
|
||
set = make(map[*client]struct{})
|
||
h.rooms[room] = set
|
||
}
|
||
set[c] = struct{}{}
|
||
}
|
||
}
|
||
|
||
// roomClients 返回房间成员的快照；room 为空表示全部连接。
|
||
func (h *Hub) roomClients(room string) []*client {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
|
||
if room == "" {
|
||
out := make([]*client, 0, len(h.clients))
|
||
for c := range h.clients {
|
||
out = append(out, c)
|
||
}
|
||
return out
|
||
}
|
||
set := h.rooms[room]
|
||
out := make([]*client, 0, len(set))
|
||
for c := range set {
|
||
out = append(out, c)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// handleHello 只信任会话身份：客户端自报的 role/userId 会被忽略，
|
||
// 管理端房间必须由登录会话授予，未登录连接请求管理端直接断开。
|
||
func (h *Hub) handleHello(ctx context.Context, c *client, p helloPayload) {
|
||
c.mu.Lock()
|
||
sessionRole, sessionUserID := c.sessionRole, c.sessionUserID
|
||
c.mu.Unlock()
|
||
|
||
if p.Role != "" && p.Role != RoleDisplay && p.Role != RoleAdmin {
|
||
h.log.Debug("hello 角色非法", "role", p.Role, "remote", c.remoteAddr)
|
||
c.stop(websocket.StatusPolicyViolation, "role 非法")
|
||
return
|
||
}
|
||
if p.Role == RoleAdmin && sessionRole != RoleAdmin {
|
||
h.log.Warn("未登录连接请求管理端房间，已断开", "remote", c.remoteAddr)
|
||
c.stop(websocket.StatusPolicyViolation, "无权订阅管理端房间")
|
||
return
|
||
}
|
||
|
||
// 不写 role 时按会话上限走；显式 role=display 允许管理端用户降级成大屏。
|
||
role := sessionRole
|
||
if p.Role == RoleDisplay {
|
||
role = RoleDisplay
|
||
}
|
||
|
||
c.mu.Lock()
|
||
c.role = role
|
||
c.userID = 0
|
||
if role == RoleAdmin {
|
||
c.userID = sessionUserID
|
||
}
|
||
c.mu.Unlock()
|
||
|
||
if role == RoleAdmin {
|
||
rooms := []string{RoomAdmin}
|
||
if c.userID > 0 {
|
||
rooms = append(rooms, AdminRoom(c.userID))
|
||
}
|
||
h.join(c, rooms...)
|
||
} else {
|
||
id := p.InstanceID
|
||
if id == "" {
|
||
id = p.DisplayKey
|
||
}
|
||
h.registerDisplay(c, id, p.Page)
|
||
h.join(c, RoomDisplay)
|
||
}
|
||
|
||
h.sendInitial(ctx, c, p.LastVersion)
|
||
}
|
||
|
||
// sendInitial 决定首包是 snapshot 还是 catchup：
|
||
// 没有 lastVersion、版本超前、或增量条数超阈值，一律回落全量快照。
|
||
func (h *Hub) sendInitial(ctx context.Context, c *client, lastVersion *int64) {
|
||
h.sendMu.Lock()
|
||
defer h.sendMu.Unlock()
|
||
|
||
c.mu.Lock()
|
||
c.greeted = true
|
||
c.mu.Unlock()
|
||
|
||
cur, err := h.currentVersion(ctx)
|
||
if err != nil {
|
||
h.log.Error("读取版本号失败，无法下发首包", "err", err)
|
||
c.stop(websocket.StatusInternalError, "版本号不可用")
|
||
return
|
||
}
|
||
h.setVersion(cur)
|
||
|
||
if lastVersion == nil || *lastVersion <= 0 || *lastVersion > cur {
|
||
h.sendSnapshotLocked(ctx, c)
|
||
return
|
||
}
|
||
|
||
n, err := h.countChangesSince(ctx, *lastVersion)
|
||
if err != nil {
|
||
h.log.Warn("统计增量失败，回落快照", "err", err)
|
||
h.sendSnapshotLocked(ctx, c)
|
||
return
|
||
}
|
||
if n > h.maxCatchupValue() {
|
||
h.sendSnapshotLocked(ctx, c)
|
||
return
|
||
}
|
||
|
||
changes, err := h.changesSince(ctx, *lastVersion, 0)
|
||
if err != nil {
|
||
h.log.Warn("读取增量失败，回落快照", "err", err)
|
||
h.sendSnapshotLocked(ctx, c)
|
||
return
|
||
}
|
||
if !c.sendEvent(EventCatchup, catchupPayload{FromVersion: *lastVersion, Version: cur, Changes: changes}) {
|
||
return
|
||
}
|
||
c.lastSent.Store(cur)
|
||
c.ackVersion.Store(cur)
|
||
}
|
||
|
||
// sendSnapshotLocked 必须在持有 sendMu 时调用：快照读取与「已发送版本」的推进
|
||
// 必须原子完成，否则并发的 change 会插到快照前面。
|
||
func (h *Hub) sendSnapshotLocked(ctx context.Context, c *client) {
|
||
snap, err := h.Snapshot(ctx)
|
||
if err != nil {
|
||
h.log.Error("构建快照失败", "err", err)
|
||
c.stop(websocket.StatusInternalError, "快照不可用")
|
||
return
|
||
}
|
||
if !c.sendEvent(EventSnapshot, snap) {
|
||
return
|
||
}
|
||
c.lastSent.Store(snap.Version)
|
||
c.ackVersion.Store(snap.Version)
|
||
h.setVersion(snap.Version)
|
||
}
|
||
|
||
// Broadcast 把变更推给所有已握手的连接（大屏与管理端）。
|
||
func (h *Hub) Broadcast(changes []Change) { h.fanout("", changes) }
|
||
|
||
// BroadcastTo 把变更推给指定房间（如 RoomAdmin、AdminRoom(uid)）。
|
||
func (h *Hub) BroadcastTo(room string, changes []Change) { h.fanout(room, changes) }
|
||
|
||
// fanout 在 sendMu 下按版本去重后下发：已经发过的版本不再重复投递。
|
||
func (h *Hub) fanout(room string, changes []Change) {
|
||
if len(changes) == 0 {
|
||
return
|
||
}
|
||
|
||
h.sendMu.Lock()
|
||
defer h.sendMu.Unlock()
|
||
|
||
version := int64(0)
|
||
for _, ch := range changes {
|
||
if ch.Version > version {
|
||
version = ch.Version
|
||
}
|
||
}
|
||
if version == 0 {
|
||
version = h.currentVersionCached(h.ctx)
|
||
}
|
||
h.setVersion(version)
|
||
|
||
for _, c := range h.roomClients(room) {
|
||
c.mu.Lock()
|
||
greeted := c.greeted
|
||
c.mu.Unlock()
|
||
if !greeted {
|
||
continue
|
||
}
|
||
|
||
last := c.lastSent.Load()
|
||
filtered := make([]Change, 0, len(changes))
|
||
for _, ch := range changes {
|
||
if ch.Version > 0 && ch.Version <= last {
|
||
continue
|
||
}
|
||
filtered = append(filtered, ch)
|
||
}
|
||
if len(filtered) == 0 {
|
||
continue
|
||
}
|
||
if !c.sendEvent(EventChange, changePayload{Version: version, Changes: filtered}) {
|
||
continue
|
||
}
|
||
c.lastSent.Store(version)
|
||
h.checkSlow(c, version)
|
||
}
|
||
}
|
||
|
||
// SendCommand 向所有大屏下发 display:command。
|
||
func (h *Hub) SendCommand(cmd DisplayCommand) { h.SendCommandTo(RoomDisplay, cmd) }
|
||
|
||
// SendCommandTo 向指定房间下发 display:command；非法 action 直接丢弃并记录。
|
||
func (h *Hub) SendCommandTo(room string, cmd DisplayCommand) {
|
||
if !ValidDisplayAction(cmd.Action) {
|
||
h.log.Warn("忽略非法的大屏指令", "action", cmd.Action, "room", room)
|
||
return
|
||
}
|
||
if cmd.Action == CmdGoto && (room == "" || room == RoomDisplay) {
|
||
h.setInstancePage(pageFromPayload(cmd.Payload))
|
||
}
|
||
|
||
h.sendMu.Lock()
|
||
defer h.sendMu.Unlock()
|
||
for _, c := range h.roomClients(room) {
|
||
c.mu.Lock()
|
||
greeted := c.greeted
|
||
c.mu.Unlock()
|
||
if !greeted {
|
||
continue
|
||
}
|
||
c.sendEvent(EventDisplayCommand, cmd)
|
||
}
|
||
}
|
||
|
||
// checkSlow 用 ack 版本差识别慢客户端；超阈值直接断开，由客户端重连后补齐。
|
||
func (h *Hub) checkSlow(c *client, version int64) {
|
||
lag := h.slowLagValue()
|
||
if lag <= 0 {
|
||
return
|
||
}
|
||
if version-c.ackVersion.Load() <= lag {
|
||
return
|
||
}
|
||
h.markSlow(c)
|
||
}
|
||
|
||
func (h *Hub) markSlow(c *client) {
|
||
if c.slow.Swap(true) {
|
||
return
|
||
}
|
||
c.mu.Lock()
|
||
id := c.instanceID
|
||
c.mu.Unlock()
|
||
|
||
h.mu.Lock()
|
||
if st := h.instances[id]; st != nil {
|
||
st.slow = true
|
||
}
|
||
h.mu.Unlock()
|
||
|
||
h.log.Warn("慢客户端：ack 长时间落后，断开后由重连补齐",
|
||
"instance", id, "ackVersion", c.ackVersion.Load(), "remote", c.remoteAddr)
|
||
c.stop(websocket.StatusPolicyViolation, "ack 落后")
|
||
}
|
||
|
||
// heartbeatLoop 每 10s 下发一次心跳（version + serverTime）。
|
||
func (h *Hub) heartbeatLoop() {
|
||
defer h.hbWG.Done()
|
||
|
||
timer := time.NewTimer(h.heartbeatIntervalValue())
|
||
defer timer.Stop()
|
||
|
||
for {
|
||
select {
|
||
case <-h.ctx.Done():
|
||
return
|
||
case <-timer.C:
|
||
h.beat()
|
||
timer.Reset(h.heartbeatIntervalValue())
|
||
case <-h.heartbeatReset:
|
||
if !timer.Stop() {
|
||
select {
|
||
case <-timer.C:
|
||
default:
|
||
}
|
||
}
|
||
timer.Reset(h.heartbeatIntervalValue())
|
||
}
|
||
}
|
||
}
|
||
|
||
func (h *Hub) beat() {
|
||
version := h.currentVersionCached(h.ctx)
|
||
now := time.Now()
|
||
data := heartbeatPayload{Version: version, ServerTime: nowISO(now)}
|
||
|
||
for _, c := range h.roomClients("") {
|
||
c.mu.Lock()
|
||
greeted := c.greeted
|
||
c.mu.Unlock()
|
||
if !greeted {
|
||
continue
|
||
}
|
||
c.sendEvent(EventHeartbeat, data)
|
||
h.touchInstance(c, now)
|
||
h.checkSlow(c, version)
|
||
}
|
||
}
|