小鱼开发
64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
"""
|
|
全局并发槽位管理器
|
|
==================
|
|
|
|
基于 Redis SET + Lua 脚本实现严格原子的槽位申请与释放。
|
|
"""
|
|
|
|
import logging
|
|
|
|
from redis.asyncio import Redis
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Lua 脚本：原子执行 SADD -> SCARD -> 条件 SREM
|
|
_ACQUIRE_LUA = """
|
|
local key = KEYS[1]
|
|
local slot_id = ARGV[1]
|
|
local max_slots = tonumber(ARGV[2])
|
|
redis.call('sadd', key, slot_id)
|
|
local count = redis.call('scard', key)
|
|
if count > max_slots then
|
|
redis.call('srem', key, slot_id)
|
|
return 0
|
|
end
|
|
redis.call('expire', key, 1800)
|
|
return 1
|
|
"""
|
|
|
|
|
|
class SlotManager:
|
|
"""全局并发槽位管理器"""
|
|
|
|
def __init__(self, redis: Redis):
|
|
self.redis = redis
|
|
|
|
async def acquire(self, slot_key: str, slot_id: str, max_slots: int) -> bool:
|
|
"""申请一个槽位。返回 True 表示成功，False 表示槽位已满。"""
|
|
try:
|
|
result = await self.redis.eval(_ACQUIRE_LUA, 1, slot_key, slot_id, str(max_slots))
|
|
acquired = result == 1
|
|
if acquired:
|
|
logger.debug(f"Slot acquired: {slot_key}/{slot_id} (max={max_slots})")
|
|
else:
|
|
logger.debug(f"Slot full: {slot_key}/{slot_id} (max={max_slots})")
|
|
return acquired
|
|
except Exception as e:
|
|
logger.warning(f"Slot acquire error: {slot_key}/{slot_id}: {e}")
|
|
return False
|
|
|
|
async def release(self, slot_key: str, slot_id: str) -> None:
|
|
"""释放一个槽位。"""
|
|
try:
|
|
await self.redis.srem(slot_key, slot_id)
|
|
logger.debug(f"Slot released: {slot_key}/{slot_id}")
|
|
except Exception as e:
|
|
logger.warning(f"Slot release error: {slot_key}/{slot_id}: {e}")
|
|
|
|
async def count(self, slot_key: str) -> int:
|
|
"""获取当前已占用的槽位数量。"""
|
|
try:
|
|
return await self.redis.scard(slot_key)
|
|
except Exception:
|
|
return 0
|