宜博 / vokoworld

/**
 * 令牌桶数学（纯函数、零 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),
  }
}