/** * 令牌桶数学(纯函数、零 IO、单一事实来源)。 * * 内存 store 与未来的 redis store 都复用这里的公式,保证两种后端行为一致。 * 算法变更只动本文件,两边自动同步。时间一律以毫秒整数(`nowMs`)注入, * 不在本文件调用 `Date.now()`,保证测试可精确推进时间、结果确定。 * * 模型:桶容量 `capacity`(突发上限),稳态速率 `refillPerSec`(每秒补充令牌)。 * 每次操作先按 `(nowMs - lastRefillMs)` 经过的时间补充令牌(不超过 capacity), * 再尝试扣减 `cost`。补满靠突发吃掉容量,之后退回 refill 速率——这就是「平滑限速 + * 允许小幅突发」。 * * @typedef {Object} BucketState * @property {number} tokens 当前令牌余额(可为小数,补充按速率连续累积) * @property {number} lastRefillMs 上次结算时间戳(毫秒) */ /** * 桶的初始状态:满桶。首次使用即拥有完整突发额度。 * @param {{ capacity: number }} rule * @param {number} nowMs * @returns {BucketState} */ export function initialState(rule, nowMs) { return { tokens: rule.capacity, lastRefillMs: nowMs } } /** * 按经过时间把状态补充到「截至 nowMs」,不扣令牌、不越过 capacity。 * `refillPerSec <= 0` 视为不补充(固定桶,配额耗尽不恢复)。 * `nowMs` 早于 `lastRefillMs`(时钟回拨)时不倒扣,按 0 经过处理。 * @param {BucketState} state * @param {{ capacity: number, refillPerSec: number }} rule * @param {number} nowMs * @returns {BucketState} */ export function refill(state, rule, nowMs) { const elapsedMs = Math.max(0, nowMs - state.lastRefillMs) const added = (elapsedMs / 1000) * rule.refillPerSec const tokens = Math.min(rule.capacity, state.tokens + added) return { tokens, lastRefillMs: nowMs } } /** * 还差 `deficit` 个令牌时,按 refill 速率需要等待的毫秒数(向上取整)。 * `refillPerSec <= 0`(不补充)时永远等不到,返回 `Number.POSITIVE_INFINITY`。 * @param {number} deficit 缺口(正数) * @param {{ refillPerSec: number }} rule * @returns {number} */ export function waitMsFor(deficit, rule) { if (deficit <= 0) return 0 if (rule.refillPerSec <= 0) return Number.POSITIVE_INFINITY return Math.ceil((deficit / rule.refillPerSec) * 1000) } /** * 在「截至 nowMs」的状态上尝试扣 `cost` 个令牌。 * 余额足够则扣减并放行;不足则不动余额、拒绝,并算出补够缺口所需等待。 * 返回新状态与判定,store 负责把 `next` 写回。 * @param {BucketState} state 任意时间点的状态(内部会先 refill 到 nowMs) * @param {number} cost * @param {{ capacity: number, refillPerSec: number }} rule * @param {number} nowMs * @returns {{ next: BucketState, allowed: boolean, remaining: number, retryAfterMs: number }} */ export function take(state, cost, rule, nowMs) { const filled = refill(state, rule, nowMs) if (filled.tokens >= cost) { const next = { tokens: filled.tokens - cost, lastRefillMs: nowMs } return { next, allowed: true, remaining: Math.floor(next.tokens), retryAfterMs: 0, } } const deficit = cost - filled.tokens return { next: filled, allowed: false, remaining: Math.floor(filled.tokens), retryAfterMs: waitMsFor(deficit, rule), } } /** * 只读判定:把状态补充到 nowMs,看 `cost` 是否能放行,但不写回、不扣减。 * 余额与等待时间按「截至 nowMs」计算。 * @param {BucketState} state * @param {number} cost * @param {{ capacity: number, refillPerSec: number }} rule * @param {number} nowMs * @returns {{ allowed: boolean, remaining: number, retryAfterMs: number }} */ export function peek(state, cost, rule, nowMs) { const filled = refill(state, rule, nowMs) const allowed = filled.tokens >= cost const deficit = allowed ? 0 : cost - filled.tokens return { allowed, remaining: Math.floor(filled.tokens), retryAfterMs: waitMsFor(deficit, rule), } }