93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import hashlib
|
||
|
|
import os
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
|
||
|
|
class PollingLease:
|
||
|
|
def __init__(self, config: dict[str, Any] | None = None):
|
||
|
|
cfg = config or {}
|
||
|
|
self._token_hash = hashlib.sha256(cfg.get("bot_token", "").encode()).hexdigest()[:8]
|
||
|
|
self._instance_id = os.urandom(6).hex()
|
||
|
|
self._lease_ttl = cfg.get("polling_lease_ttl", 30)
|
||
|
|
self._acquired = False
|
||
|
|
self._renew_task: asyncio.Task | None = None
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_acquired(self) -> bool:
|
||
|
|
return self._acquired
|
||
|
|
|
||
|
|
async def acquire(self) -> bool:
|
||
|
|
try:
|
||
|
|
from yuxi.storage.redis import get_redis
|
||
|
|
|
||
|
|
redis = get_redis()
|
||
|
|
except ImportError:
|
||
|
|
logger.error(
|
||
|
|
"[Telegram] Redis not available - lease disabled. "
|
||
|
|
"Multiple instances may poll simultaneously, causing duplicate message processing."
|
||
|
|
)
|
||
|
|
self._acquired = True
|
||
|
|
return True
|
||
|
|
|
||
|
|
lease_key = f"telegram:polling_lease:{self._token_hash}"
|
||
|
|
|
||
|
|
try:
|
||
|
|
acquired = await redis.set(lease_key, self._instance_id, nx=True, ex=self._lease_ttl)
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(
|
||
|
|
f"[Telegram] Redis lease check failed: {e}. "
|
||
|
|
"Proceeding without lease - multiple instances may poll simultaneously."
|
||
|
|
)
|
||
|
|
self._acquired = True
|
||
|
|
return True
|
||
|
|
|
||
|
|
if acquired:
|
||
|
|
self._acquired = True
|
||
|
|
self._renew_task = asyncio.create_task(self._renew_loop(lease_key))
|
||
|
|
logger.info(f"[Telegram] Polling lease acquired (instance={self._instance_id}, ttl={self._lease_ttl}s)")
|
||
|
|
return True
|
||
|
|
|
||
|
|
logger.debug("[Telegram] Polling lease held by another instance")
|
||
|
|
return False
|
||
|
|
|
||
|
|
async def release(self) -> None:
|
||
|
|
self._acquired = False
|
||
|
|
if self._renew_task and not self._renew_task.done():
|
||
|
|
self._renew_task.cancel()
|
||
|
|
try:
|
||
|
|
await self._renew_task
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
pass
|
||
|
|
self._renew_task = None
|
||
|
|
|
||
|
|
try:
|
||
|
|
from yuxi.storage.redis import get_redis
|
||
|
|
|
||
|
|
redis = get_redis()
|
||
|
|
lease_key = f"telegram:polling_lease:{self._token_hash}"
|
||
|
|
await redis.delete(lease_key)
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
async def _renew_loop(self, key: str) -> None:
|
||
|
|
try:
|
||
|
|
from yuxi.storage.redis import get_redis
|
||
|
|
|
||
|
|
redis = get_redis()
|
||
|
|
except ImportError:
|
||
|
|
return
|
||
|
|
|
||
|
|
while self._acquired:
|
||
|
|
await asyncio.sleep(self._lease_ttl * 0.5)
|
||
|
|
try:
|
||
|
|
await redis.expire(key, self._lease_ttl)
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
break
|
||
|
|
except Exception:
|
||
|
|
pass
|