ForcePilot/backend/package/yuxi/channels/adapters/wechat/polling_lease.py

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, 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,
        }