95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
import asyncio
|
|
import contextlib
|
|
import signal
|
|
|
|
import asyncpg
|
|
import uvloop
|
|
|
|
from dependencies.container import container
|
|
from userbot.modules.uptime import repository as uptime
|
|
from utils.listen import Listener
|
|
from utils.logging import logger, setup_logging
|
|
from utils.read.bots import BOTS_CHANGED_CHANNEL
|
|
from utils.read.watches import WATCHES_CHANGED_CHANNEL
|
|
from utils.storage import ContentAddressedStorage
|
|
from watcher.registry import BotRegistry
|
|
|
|
setup_logging()
|
|
|
|
RETRY_INTERVAL = 300
|
|
UPTIME_INTERVAL = 60
|
|
|
|
|
|
async def _retry_loop(registry: BotRegistry) -> None:
|
|
while True:
|
|
await asyncio.sleep(RETRY_INTERVAL)
|
|
try:
|
|
await registry.sync()
|
|
except Exception:
|
|
logger.exception("[red]Bot sync failed[/]")
|
|
|
|
|
|
async def _keep_uptime(registry: BotRegistry, pool: asyncpg.Pool) -> None:
|
|
while True:
|
|
for account_id in registry.connected_accounts:
|
|
with contextlib.suppress(Exception):
|
|
await uptime.heartbeat(pool, account_id)
|
|
await asyncio.sleep(UPTIME_INTERVAL)
|
|
|
|
|
|
def _change_listener(registry: BotRegistry, tasks: set[asyncio.Task]) -> Listener:
|
|
def on_change(_payload: str) -> None:
|
|
task = asyncio.create_task(registry.sync())
|
|
tasks.add(task)
|
|
task.add_done_callback(tasks.discard)
|
|
|
|
def on_watch(_payload: str) -> None:
|
|
task = asyncio.create_task(registry.refresh_watches())
|
|
tasks.add(task)
|
|
task.add_done_callback(tasks.discard)
|
|
|
|
def on_reconnect() -> None:
|
|
on_change("")
|
|
on_watch("")
|
|
|
|
return Listener(
|
|
"watcher",
|
|
{BOTS_CHANGED_CHANNEL: on_change, WATCHES_CHANGED_CHANNEL: on_watch},
|
|
on_reconnect,
|
|
)
|
|
|
|
|
|
async def runner() -> None:
|
|
pool = await container.get(asyncpg.Pool)
|
|
storage = await container.get(ContentAddressedStorage)
|
|
registry = BotRegistry(pool, storage)
|
|
tasks: set[asyncio.Task] = set()
|
|
listener = _change_listener(registry, tasks)
|
|
stop = asyncio.Event()
|
|
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
|
|
retry = asyncio.create_task(_retry_loop(registry))
|
|
beat = asyncio.create_task(_keep_uptime(registry, pool))
|
|
try:
|
|
await registry.sync()
|
|
if not registry.running:
|
|
logger.warning("[yellow]No bots yet. Add a bot token in the web UI.[/]")
|
|
await listener.start()
|
|
logger.info("[green]Bot watcher running.[/]")
|
|
await stop.wait()
|
|
finally:
|
|
for task in (retry, beat):
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
await listener.stop()
|
|
await registry.close()
|
|
await container.close()
|
|
|
|
|
|
def main() -> None:
|
|
uvloop.install()
|
|
logger.info("Starting bot watcher...")
|
|
with contextlib.suppress(KeyboardInterrupt):
|
|
asyncio.run(runner())
|
|
logger.info("[red]Bot watcher stopped.[/]")
|