浅念
58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
"""WS 每 IP 并发连接上限(spec §6.1 P2)——进程内 (租户, IP) 计数闸。
|
|
|
|
动机: 平台化后单 IP 可对所有租户域名各开一圈 WS 耗尽单进程资源;
|
|
REST 消息接口已有的 rate_limit.py 管不到长连接握手。key 是
|
|
(tenant_slug, ip) 二元组——各租户名额独立,不共享预算。
|
|
|
|
约束与取舍(与 rate_limit.py 同一族):
|
|
- 单进程内存态, ``--workers 1`` 约束下准确; 多 worker 各自一份计数,
|
|
实际上限 = limit * workers。
|
|
- 纯同步 dict 操作、方法内无 await——单事件循环下 check-then-set
|
|
原子,无需锁;不得在多线程里使用。
|
|
- 计数归零即删 key, 长期运行不随历史 (租户, IP) 对积累内存;
|
|
release 未持有的 key 是安全 no-op(不出负数)。
|
|
- limit 可在 try_acquire 时按租户配置覆盖: 模块单例服务全部租户,而
|
|
ws_conn_limit 是 engine 组的租户级配置——按次传入让各租户差异实时
|
|
生效(含 registry_reload 后), 又不必维护"每租户一个 limiter 实例"
|
|
的第二层簿记(实例缓存还有配置改了不生效的滞留问题)。
|
|
"""
|
|
from typing import Optional
|
|
|
|
DEFAULT_WS_CONN_LIMIT = 10
|
|
|
|
|
|
class ConnectionLimiter:
|
|
"""每 (tenant, ip) 并发计数器: 满则 try_acquire 拒绝, release 归零清键。"""
|
|
|
|
def __init__(self, limit: int = DEFAULT_WS_CONN_LIMIT) -> None:
|
|
self._limit = limit
|
|
self._counts: dict[tuple[str, str], int] = {}
|
|
|
|
def try_acquire(self, tenant: str, ip: str, *,
|
|
limit: Optional[int] = None) -> bool:
|
|
"""未满则记账并返回 True; 已满返回 False(不记账)。"""
|
|
cap = self._limit if limit is None else limit
|
|
key = (tenant, ip)
|
|
n = self._counts.get(key, 0)
|
|
if n >= cap:
|
|
return False
|
|
self._counts[key] = n + 1
|
|
return True
|
|
|
|
def release(self, tenant: str, ip: str) -> None:
|
|
"""归还一个名额; 计数归零删 key; 未持有的 key 是安全 no-op。"""
|
|
key = (tenant, ip)
|
|
n = self._counts.get(key, 0)
|
|
if n <= 1:
|
|
self._counts.pop(key, None)
|
|
else:
|
|
self._counts[key] = n - 1
|
|
|
|
|
|
# app 级单例: chat.py 的 WS 路由共用这一份进程内计数(--workers 1 约束)
|
|
limiter = ConnectionLimiter()
|
|
|
|
|
|
def _reset_for_test() -> None:
|
|
limiter._counts.clear()
|