from __future__ import annotations import asyncio import os import time from typing import Any import redis.asyncio as aioredis from yuxi.utils.logging_config import logger class PollingLease: _LOCK_KEY_PREFIX = "wechat:polling_lease:" def __init__( self, channel_id: str = "", lease_ttl: float = 15.0, renew_interval: float = 0.5, redis_url: str | None = None, ): self._channel_id = channel_id or os.environ.get("CHANNEL_INSTANCE_ID", "wechat-default") self._lease_ttl = lease_ttl self._renew_interval = renew_interval self._acquired_at: float = 0.0 self._last_renew: float = 0.0 self._active = False self._renew_task: asyncio.Task | None = None self._lock_key = f"{self._LOCK_KEY_PREFIX}{self._channel_id}" self._lock_value = f"{os.getpid()}:{time.monotonic()}" self._redis: aioredis.Redis | None = None self._redis_url = redis_url or os.environ.get("REDIS_URL", "redis://redis:6379/0") @property def is_active(self) -> bool: if not self._active: return False if self._lease_ttl > 0 and time.monotonic() - self._last_renew > self._lease_ttl * 1.5: self._active = False logger.warning(f"[PollingLease/{self._channel_id}] Lease expired (stale)") return False return self._active @property def lease_id(self) -> str: return self._channel_id async def _ensure_redis(self) -> aioredis.Redis: if self._redis is None: self._redis = aioredis.from_url(self._redis_url, decode_responses=False) return self._redis async def try_acquire(self, force: bool = False) -> bool: if self._active and not force: return True try: redis_client = await self._ensure_redis() acquired = await redis_client.set( self._lock_key, self._lock_value, nx=True, px=int(self._lease_ttl * 1000), ) if not acquired: logger.warning( f"[PollingLease/{self._channel_id}] Lease held by another instance" ) return False except Exception as e: logger.warning( f"[PollingLease/{self._channel_id}] Redis unavailable, " f"falling back to local mode: {e}" ) self._acquired_at = time.monotonic() self._last_renew = self._acquired_at self._active = True logger.info( f"[PollingLease/{self._channel_id}] Lease acquired " f"(ttl={self._lease_ttl}s, renew_interval={self._renew_interval}s)" ) return True async def start_renew(self) -> None: if self._renew_task and not self._renew_task.done(): return self._renew_task = asyncio.create_task(self._renew_loop()) async def stop_renew(self) -> None: 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 async def _renew_loop(self) -> None: while self._active: try: await asyncio.sleep(self._renew_interval) if self._redis is not None: try: await self._redis.pexpire( self._lock_key, int(self._lease_ttl * 1000), ) except Exception as e: logger.error( f"[PollingLease/{self._channel_id}] Renew failed: {e}" ) self._last_renew = time.monotonic() except asyncio.CancelledError: break async def release(self) -> None: await self.stop_renew() self._active = False if self._redis is not None: try: script = """ if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("DEL", KEYS[1]) else return 0 end """ await self._redis.eval(script, 1, self._lock_key, self._lock_value) except Exception as e: logger.error(f"[PollingLease/{self._channel_id}] Release failed: {e}") logger.info(f"[PollingLease/{self._channel_id}] Lease released") async def close(self) -> None: if self._redis is not None: await self._redis.aclose() self._redis = None def get_snapshot(self) -> dict[str, Any]: return { "channel_id": self._channel_id, "active": self._active, "acquired_at": self._acquired_at, "last_renew": self._last_renew, "lease_ttl": self._lease_ttl, }