fix: Lua 原子闸门堵死并发 TOCTOU 击穿

根因:_check_request_block 的"读 usage + 比较 + 放行"是纯读检查,
与结算 incrby 之间无原子性。并发请求在同一 usage 快照上各自判定
"未超限"集体放行,叠加消耗击穿限额(08-18 13:14:41-50,5 个并发
请求从 1.49M 冲到 2.11M,超支 11.2 万 token)。

修复:Lua 原子闸门把"读 usage + 读 inflight + 判 limit + INCRBY
预扣"合并为单个 Redis 原子操作,并发请求串行看到彼此的在途预留:
  - 放行:INCRBY inflight reserve,返回 granted
  - 拦截:不预扣,返回 blocked

配对释放:pre-hook 放行时把 reserve 写入 data["metadata"]
(PREBLOCK_RESERVE_META_KEY),litellm 透传到 log_success_event
的 kwargs["litellm_params"]["metadata"],结算时 _release_inflight
配对 DECRBY,usage 只计实际消耗。

兜底:inflight key 设 PREBLOCK_INFLIGHT_TTL(1h),失败/丢失的
预扣由 TTL 自动归零,不永久占额度。

验证:打桩 redis/litellm 跑 5 场景,含复现 08-18 击穿的并发测试
(5 并发 reserve,旧代码全放行击穿,新代码仅放行 2 个且
usage+inflight < 2M)。ALL PASS ✅

volc_limit_sync.py

@@ -109,6 +109,35 @@ PREBLOCK_RESERVE_TOKENS = int(os.getenv("PREBLOCK_RESERVE_TOKENS", "50000"))
 # max(声明值, 本下限) 兜住思考开销；2M 限额下 4096 偏差仅 0.2%。
 THINKING_RESERVE_FLOOR = int(os.getenv("THINKING_RESERVE_FLOOR", "4096"))
 
+# 在途预留（inflight）TTL：pre-hook 放行时预扣的 reserve 记入独立
+# daily_usage_inflight key，正常请求在结算回调中配对释放；若请求
+# 失败/进程崩溃/回调丢失，预留量由本 TTL 自动归零，防止额度被永久
+# 卡死。1 小时远大于正常结算延迟（秒级），兜底释放足够及时。
+PREBLOCK_INFLIGHT_TTL = int(os.getenv("PREBLOCK_INFLIGHT_TTL", "3600"))
+
+# pre-hook 写入请求 metadata 的预留量键名，结算回调据此配对释放 inflight。
+# 以 "_" 前缀避免与用户/代理自定义 metadata 冲突。
+PREBLOCK_RESERVE_META_KEY = "_volc_preblock_reserve"
+
+# Lua 原子闸门：把"读 usage + 读 inflight + 判 limit + 预扣"合并为
+# 一个 Redis 原子操作，堵住并发请求在同一 usage 快照上集体撞线的
+# TOCTOU 窗口（见 docs/volc-quota-notes.md 并发击穿分析）。
+# KEYS[1]=usage_key, KEYS[2]=inflight_key
+# ARGV[1]=effective_limit, ARGV[2]=reserve, ARGV[3]=inflight TTL
+# 返回: 1=放行(已预扣), 0=拦截(未预扣)
+GATE_SCRIPT = """
+local usage = tonumber(redis.call('GET', KEYS[1]) or '0')
+local inflight = tonumber(redis.call('GET', KEYS[2]) or '0')
+local limit = tonumber(ARGV[1])
+local reserve = tonumber(ARGV[2])
+if usage + inflight + reserve >= limit then
+    return 0
+end
+redis.call('INCRBY', KEYS[2], reserve)
+redis.call('EXPIRE', KEYS[2], tonumber(ARGV[3]))
+return 1
+"""
+
 # stable_limit 缺失时（首次同步前/新增模型/TTL 过期/Redis 数据丢失）的保守兜底限额。
 # 免费池常态 ~2M（首开赠送/昨日采多少返多少），若 fallback 裸奔到采集上限 2M，
 # 余额不足或已耗尽的模型会被烧穿扣费。900k 对 0 余额模型仍是 9 倍优惠，
@@ -375,6 +404,16 @@ def get_usage_key(clean_model, api_base=None):
     return f"daily_usage:{date_str}:{clean_model}{suffix}"
 
 
+def get_inflight_key(clean_model, api_base=None):
+    """在途预留 key：与 usage key 同日期轮换，独立账户，结算配对释放。
+
+    pre-hook 放行时 INCRBY reserve（Lua 原子闸门内完成），请求成功结算时
+    按该请求的 reserve DECRBY 释放，usage 只计实际消耗；失败/丢失的预留
+    由 PREBLOCK_INFLIGHT_TTL 自动归零。
+    """
+    return get_usage_key(clean_model, api_base) + ":inflight"
+
+
 def get_static_limit_for_hash(clean_model, current_hash=None):
     """Check static configuration in HARD_DAILY_LIMITS using an optional endpoint hash."""
     if current_hash:
@@ -862,13 +901,60 @@ def startup_audit():
         log(f"❌ Startup audit failed: {e}", level="ERROR")
 
 
+def _atomic_reserve(clean_model, api_base, effective_limit, reserve):
+    """Lua 原子闸门：usage + inflight + reserve >= limit 则拦截，否则预扣并放行。
+
+    返回 (granted, usage_total)。granted=False 时 usage_total = usage + inflight
+    （供日志展示）；granted=True 时预扣已入 inflight key。
+    Redis 不可用/脚本失败时抛异常由调用方 fail-closed 处理。
+    """
+    usage_key = get_usage_key(clean_model, api_base)
+    inflight_key = get_inflight_key(clean_model, api_base)
+    granted = redis_client.eval(
+        GATE_SCRIPT,
+        2,
+        usage_key,
+        inflight_key,
+        int(float(effective_limit)),
+        int(reserve),
+        int(PREBLOCK_INFLIGHT_TTL),
+    )
+    usage = float(redis_client.get(usage_key) or 0)
+    inflight = float(redis_client.get(inflight_key) or 0)
+    return bool(granted), usage + inflight
+
+
+def _release_inflight(clean_model, api_base, reserve):
+    """结算配对释放：请求完成后把 pre-hook 预扣的 reserve 从 inflight 账户扣除。
+
+    reserve<=0 或 inflight 已过期（TTL 兜底释放）时无操作。失败/丢失的预留
+    由 TTL 自动归零，这里只做正常结算路径的精确回收。
+    """
+    if not redis_client or not reserve:
+        return
+    try:
+        inflight_key = get_inflight_key(clean_model, api_base)
+        cur = float(redis_client.get(inflight_key) or 0)
+        if cur <= 0:
+            return
+        # DECRBY 到负数说明 inflight 已部分过期，直接 DEL 归零即可
+        new_val = redis_client.decrby(inflight_key, int(reserve))
+        if int(new_val) <= 0:
+            redis_client.delete(inflight_key)
+    except Exception as e:
+        log(f"❌ _release_inflight failed: {e}", level="ERROR")
+
+
 def _check_request_block(model_name, reserve=0.0):
     """
     请求前闸门：对外模型名 → 判定是否超限。
     只拦截 MONITORED_MODEL_PAIRS 内的火山监控条目；newapi 等非监控渠道天然放行。
-    reserve: 本次请求的预估消耗（max_tokens 或常量兜底），usage + reserve >= limit 即拦截，
+    reserve: 本次请求的预估消耗（max_tokens 或常量兜底）。
-    避免并发请求在闸门检查后集体撞线（限额是低频快照，用量是实时累计）。
+    判定用 Lua 原子闸门：usage + inflight(在途预留) + reserve >= limit 即拦截，
-    返回 (clean_model, usage, limit) 表示超限，None 表示放行。
+    inflight 账户在 Redis 内原子预扣，并发请求串行看到彼此的在途量，堵死
+    "同一 usage 快照集体撞线"的 TOCTOU 击穿。放行时 reserve 已预扣入 inflight，
+    由结算回调配对释放（见 _release_inflight）。
+    返回 (clean_model, usage, limit) 表示超限，None 表示放行（已预扣）。
     fail-closed：Redis 不可用 / 限额获取失败 / 读取异常时，对监控条目返回
     (clean_model, -1, -1) 标记拒绝（宁可错杀，不可漏放）；非监控条目不受影响。
     手动软踢：Redis key manual_evict:{config model 字段全名} 存在时返回
@@ -975,16 +1061,24 @@ def _check_request_block(model_name, reserve=0.0):
                 return (clean, -1, -1)
 
             usage_key = get_usage_key(clean, api_base)
-            val = redis_client.get(usage_key)
-            usage = float(val or 0)
 
             # 已被软踢的条目恒 429(无论 reserve)，避免踢后灰色区间(如 1.95M~2M)
             # 小请求被放行 → 路由原名不存在 → 400 路由错误
             if entry_name == f"{model_name}-limit-reached" and "_daily_limit_backup" in entry:
+                usage = float(redis_client.get(usage_key) or 0)
                 return (clean, usage, effective_limit)
 
-            if usage + float(reserve) >= float(effective_limit):
+            # Lua 原子闸门：usage + inflight + reserve >= limit 即拒绝（不预扣），
-                return (clean, usage, effective_limit)
+            # 否则把 reserve 预扣进 inflight 账户并放行。判定与预扣在 Redis 内
+            # 原子完成，并发请求串行看到彼此的在途预留，堵死"同一 usage 快照上
+            # 集体撞线"的 TOCTOU 击穿（08-18 13:14:41-50 根因，见诊断）。
+            granted, _inflight_total = _atomic_reserve(
+                clean, api_base, effective_limit, reserve
+            )
+            if not granted:
+                return (clean, _inflight_total, effective_limit)
+            # 放行：预扣已入 inflight，结算时配对释放
+            return None
 
         return None
     except Exception as e:
@@ -1081,7 +1175,19 @@ class DailyLimitHandler(CustomLogger):
 
         blocked = _check_request_block(data["model"], reserve=reserve)
         if not blocked:
-            return None
+            # 放行：Lua 原子闸门已把 reserve 预扣进 inflight 账户。
+            # 把本次预留量记入请求 metadata，随 litellm 流转到 log_success_event
+            # 的 kwargs["litellm_params"]["metadata"]，结算时配对释放（否则
+            # inflight 只会增不减，额度被已放行请求永久占死）。
+            try:
+                meta = data.get("metadata") or {}
+                if not isinstance(meta, dict):
+                    meta = {}
+                meta[PREBLOCK_RESERVE_META_KEY] = int(reserve)
+                data["metadata"] = meta
+            except Exception:
+                pass  # 写 metadata 失败仅影响结算释放，TTL 兜底仍会回收
+            return data
 
         clean, usage, limit = blocked
         if usage < 0 or limit < 0:
@@ -1157,6 +1263,20 @@ class DailyLimitHandler(CustomLogger):
 
             if clean_model and total_tokens > 0:
                 key = get_usage_key(clean_model, api_base)
+
+                # 结算配对释放：读回 pre-hook 预扣的 reserve（经 metadata 透传），
+                # 从 inflight 账户扣除。实际消耗 total_tokens 照常记入 usage；
+                # 预扣量 = 预估值，与实耗差额天然由 usage 侧的"只计实际"修正，
+                # inflight 只代表"已放行未结算"的在途量，结算即清零。
+                released = 0
+                try:
+                    meta = litellm_params.get("metadata") or {}
+                    if isinstance(meta, dict):
+                        released = int(meta.get(PREBLOCK_RESERVE_META_KEY, 0) or 0)
+                except (TypeError, ValueError):
+                    released = 0
+                _release_inflight(clean_model, api_base, released)
+
                 current_total = redis_client.incrby(key, int(total_tokens))
                 redis_client.expire(key, 172800)