Codex Release Bot
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
"""Distributed per-server Redis lock for BT Panel SSH bootstrap."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import secrets
|
|
from contextlib import asynccontextmanager
|
|
from typing import AsyncIterator
|
|
|
|
from server.infrastructure.redis.client import get_redis
|
|
|
|
logger = logging.getLogger("nexus.btpanel.bootstrap_lock")
|
|
|
|
_LOCK_KEY_PREFIX = "btpanel:bootstrap:lock"
|
|
_LOCK_TTL_SECONDS = 300
|
|
_LOCK_WAIT_TIMEOUT_SECONDS = 300.0
|
|
_LOCK_POLL_SECONDS = 0.2
|
|
|
|
_RELEASE_SCRIPT = """
|
|
if redis.call("get", KEYS[1]) == ARGV[1] then
|
|
return redis.call("del", KEYS[1])
|
|
end
|
|
return 0
|
|
"""
|
|
|
|
|
|
class BootstrapLockTimeout(RuntimeError):
|
|
"""Raised when another worker is already bootstrapping the same server."""
|
|
|
|
|
|
def _lock_key(server_id: int) -> str:
|
|
return f"{_LOCK_KEY_PREFIX}:{server_id}"
|
|
|
|
|
|
@asynccontextmanager
|
|
async def bootstrap_lock(
|
|
server_id: int,
|
|
*,
|
|
wait_timeout_seconds: float = _LOCK_WAIT_TIMEOUT_SECONDS,
|
|
) -> AsyncIterator[None]:
|
|
redis = get_redis()
|
|
key = _lock_key(server_id)
|
|
token = secrets.token_urlsafe(24)
|
|
loop = asyncio.get_running_loop()
|
|
deadline = loop.time() + max(wait_timeout_seconds, 0.0)
|
|
|
|
while True:
|
|
acquired = bool(await redis.set(key, token, ex=_LOCK_TTL_SECONDS, nx=True))
|
|
if acquired:
|
|
break
|
|
if wait_timeout_seconds <= 0 or loop.time() >= deadline:
|
|
raise BootstrapLockTimeout(f"BT Panel bootstrap is already running for server {server_id}")
|
|
await asyncio.sleep(_LOCK_POLL_SECONDS)
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
try:
|
|
await redis.eval(_RELEASE_SCRIPT, 1, key, token)
|
|
except Exception:
|
|
logger.warning("failed to release bt panel bootstrap lock server=%s", server_id, exc_info=True)
|