fix(collab): 二进制帧的宽松限流档名存实亡

两档共用一个令牌桶:桶按严格档的 20 初始化,之后每条文本帧(含 30 秒
一次的心跳)都会 Math.min(RATE_BURST, ...) 把它压回 20,二进制帧那档
标的 200 突发根本拿不到。实测连打 150 帧在第 101 帧被 1008 踢下线——
正是设计里要避免的「协作编辑时打字把自己踢掉」。

改成两个独立的桶。文本帧仍是 20 / 每秒 2,二进制帧 200 / 每秒 100。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016DHxhKNxXfG89JnVzHbvgj
This commit is contained in:
2026-08-30 08:29:36 -06:00
parent de4e745d90
commit c7132025f7

View File

@@ -55,14 +55,27 @@ export function isAllowedWebSocketOrigin(origin: string | null, url: URL) {
)
}
interface RateBucket {
tokens: number
updatedAt: number
}
export interface SubmissionSocketData {
userId: number
/** 同一个 Bun.serve 只能挂一个 websocket handler用它区分通道 */
kind: "submissions" | "config" | "collab"
/** 握手时那张会话的 token留着定期确认它还没被登出 / 过期,见 sweepSessions */
token: string
/** 令牌桶open 时初始化,见 allowMessage */
rate?: { tokens: number; updatedAt: number }
/** 文本控制帧的令牌桶open 时初始化,见 allowMessage */
rate?: RateBucket
/**
* collab 二进制帧的令牌桶,和 rate 分开。
*
* 共用一个桶的话宽松档名存实亡:每条文本帧(含 30 秒一次的心跳)都会
* `Math.min(RATE_BURST, ...)` 把桶压回 20二进制帧再怎么标 200 突发也拿不到。
* 实测连打 150 帧会在第 101 帧被 1008 踢下线。
*/
binaryRate?: RateBucket
/** 握手时从会话里读,三种 kind 都会填collab 通道用它判断老师身份、拼 room_open 里的姓名 */
username?: string
adminType?: string
@@ -89,21 +102,31 @@ const RATE_REFILL_PER_SECOND = 2
const COLLAB_BINARY_BURST = 200
const COLLAB_BINARY_REFILL_PER_SECOND = 100
function allowMessage(
ws: Bun.ServerWebSocket<SubmissionSocketData>,
burst = RATE_BURST,
refillPerSecond = RATE_REFILL_PER_SECOND,
) {
function consume(bucket: RateBucket, burst: number, refillPerSecond: number) {
const now = Date.now()
const rate = (ws.data.rate ??= { tokens: burst, updatedAt: now })
const refill = ((now - rate.updatedAt) / 1000) * refillPerSecond
rate.tokens = Math.min(burst, rate.tokens + refill)
rate.updatedAt = now
if (rate.tokens < 1) return false
rate.tokens -= 1
const refill = ((now - bucket.updatedAt) / 1000) * refillPerSecond
bucket.tokens = Math.min(burst, bucket.tokens + refill)
bucket.updatedAt = now
if (bucket.tokens < 1) return false
bucket.tokens -= 1
return true
}
/** 文本帧:严格档。会查库,走这一档的都按最坏情况算 */
function allowMessage(ws: Bun.ServerWebSocket<SubmissionSocketData>) {
const bucket = (ws.data.rate ??= { tokens: RATE_BURST, updatedAt: Date.now() })
return consume(bucket, RATE_BURST, RATE_REFILL_PER_SECOND)
}
/** collab 二进制帧:宽松档,独立的桶 —— 见 binaryRate 的注释 */
function allowCollabBinary(ws: Bun.ServerWebSocket<SubmissionSocketData>) {
const bucket = (ws.data.binaryRate ??= {
tokens: COLLAB_BINARY_BURST,
updatedAt: Date.now(),
})
return consume(bucket, COLLAB_BINARY_BURST, COLLAB_BINARY_REFILL_PER_SECOND)
}
/**
* 当前挂着的连接。Bun 不提供遍历连接的接口,要定期巡检就得自己登记。
* open 时加入、close 时移除,见 sweepSessions。
@@ -190,6 +213,7 @@ export function submissionWebSocketHandler(): Bun.WebSocketHandler<SubmissionSoc
liveSockets.add(ws)
ws.data.rate = { tokens: RATE_BURST, updatedAt: Date.now() }
if (ws.data.kind === "collab") {
ws.data.binaryRate = { tokens: COLLAB_BINARY_BURST, updatedAt: Date.now() }
handleCollabOpen(ws)
return
}
@@ -203,7 +227,7 @@ export function submissionWebSocketHandler(): Bun.WebSocketHandler<SubmissionSoc
message(ws, message) {
if (ws.data.kind === "collab") {
if (typeof message !== "string") {
if (!allowMessage(ws, COLLAB_BINARY_BURST, COLLAB_BINARY_REFILL_PER_SECOND)) {
if (!allowCollabBinary(ws)) {
ws.close(1008, "Too many messages")
return
}