package realtime import ( "context" "database/sql" "fmt" "log/slog" "net/http" "strings" "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 defaultCloseTimeout = 2 * time.Second defaultSendBuffer = 64 ) // slowCloseCode 是踢掉慢客户端时使用的自定义关闭码。 // 不能用 1008:前端把所有 1008 当「登录失效」并永久停止重连,慢客户端被踢后大屏会冻结。 const slowCloseCode = websocket.StatusCode(4000) // 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 allowedOrigins []string } // 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() } // SetAllowedOrigins 设置 WebSocket 跨站 Origin 白名单(path.Match 模式); // 同源与不带 Origin 的握手始终放行,不受此设置影响。 func (h *Hub) SetAllowedOrigins(origins []string) { h.mu.Lock() h.allowedOrigins = append([]string(nil), origins...) 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 } func (h *Hub) allowedOriginsValue() []string { h.mu.RLock() defer h.mu.RUnlock() return h.allowedOrigins } // 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{ // 同源与不带 Origin 的握手由库放行;跨站需命中 server.allowedOrigins 白名单。 OriginPatterns: h.allowedOriginsValue(), CompressionMode: websocket.CompressionDisabled, }) if err != nil { h.log.Debug("WebSocket 握手失败", "remote", r.RemoteAddr, "origin", r.Header.Get("Origin"), "err", err) return } h.mu.Lock() if h.closed { h.mu.Unlock() _ = conn.Close(websocket.StatusGoingAway, "服务关闭") return } h.connWG.Add(1) h.mu.Unlock() defer h.connWG.Done() // 身份只认会话:/ws 必须被 Authenticate 中间件包住,否则一律按 guest 处理。 role, userID, settingsAdmin := wsIdentity(r.Context()) h.serve(conn, r.RemoteAddr, role, userID, settingsAdmin) }) } // wsIdentity 从请求上下文取会话身份并折算成实时层角色。 // 未登录(guest)只能进大屏房间;recorder/admin 才有管理端房间的订阅权。 // settingsAdmin 单独返回真实的 admin 判定:recorder 能进管理端房间,但读不到系统配置。 func wsIdentity(ctx context.Context) (string, int64, bool) { u := middleware.User(ctx) if u == nil || !u.Role.Allows(auth.RoleRecorder) { return RoleDisplay, 0, false } return RoleAdmin, u.ID, u.Role.Allows(auth.RoleAdmin) } func (h *Hub) serve(conn *websocket.Conn, remoteAddr, sessionRole string, sessionUserID int64, settingsAdmin bool) { conn.SetReadLimit(1 << 20) c := newClient(h, conn, remoteAddr, sessionRole, sessionUserID, settingsAdmin) 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 // 先做关闭握手,再取消读上下文:coder/websocket 在读上下文被取消时会直接 // 关闭底层连接,那样对端只能看到 EOF,收不到策略性关闭码(如 1008)。 // 写循环已因 c.done 退出,读循环会随对端回帧自然结束。 c.mu.Lock() code, reason := c.closeCode, c.closeReason c.mu.Unlock() if code == 0 { code = websocket.StatusNormalClosure } closed := make(chan struct{}) go func() { defer close(closed) _ = conn.Close(code, reason) }() select { case <-closed: case <-time.After(defaultCloseTimeout): // 对端不回应:cancel 会强制关闭底层连接,读循环随之退出。 } cancel() loops.Wait() h.removeClient(c) } 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, p.Screen) 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 } // 快照对所有角色下发,但系统配置只有 admin 能看:非 admin 只保留大屏渲染所需字段。 if !c.settingsAdmin { snap.Settings = snap.Settings.WithoutAdminFields() } 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), nil) } 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) } } // SendCommandToInstances 按实例定向下发 display:command,返回命中的实例数。 // // targets 是要控制的 instanceID(GET /api/display/instances 里的 id);为空时退化成广播, // 公告这类全屏指令靠的就是广播,这条语义不能改。已离线的实例不计入命中, // 调用方据此告诉现场「那块屏没收到」。非法 action 返回错误而不是静默丢弃, // 定向指令需要一个明确的失败信号。 func (h *Hub) SendCommandToInstances(targets []string, cmd DisplayCommand) (int, error) { if !ValidDisplayAction(cmd.Action) { return 0, fmt.Errorf("未知的大屏指令 %q", cmd.Action) } if len(targets) == 0 { h.SendCommand(cmd) return h.onlineDisplayCount(), nil } wanted := make(map[string]bool, len(targets)) for _, id := range targets { if id = strings.TrimSpace(id); id != "" { wanted[id] = true } } h.sendMu.Lock() matched := make(map[string]bool, len(wanted)) for _, c := range h.roomClients(RoomDisplay) { c.mu.Lock() greeted, id := c.greeted, c.instanceID c.mu.Unlock() if !greeted || !wanted[id] { continue } c.sendEvent(EventDisplayCommand, cmd) matched[id] = true } h.sendMu.Unlock() if cmd.Action == CmdGoto { h.setInstancePage(pageFromPayload(cmd.Payload), targets) } return len(matched), nil } // onlineDisplayCount 统计已握手且在线的实例数,用于告知调用方广播命中了几块屏。 func (h *Hub) onlineDisplayCount() int { seen := make(map[string]bool) for _, c := range h.roomClients(RoomDisplay) { c.mu.Lock() greeted, id := c.greeted, c.instanceID c.mu.Unlock() if greeted && id != "" { seen[id] = true } } return len(seen) } // 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(slowCloseCode, "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) } }