Scoreboard/server/internal/realtime/hub.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
	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.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)
	}
}