feat(watcher): watched bots report capture uptime like the userbot
This commit is contained in:
1 parent
f62c345e2b
commit
7f7fc1296e
2 files changed
+21
-3
No files matched your search
@@ -113,6 +113,12 @@ class BotRegistry:
|
||||
def running(self) -> list[int]:
|
||||
return list(self._running)
|
||||
|
||||
@property
|
||||
def connected_accounts(self) -> list[int]:
|
||||
return [
|
||||
bot.account_id for bot in self._running.values() if bot.client.is_connected
|
||||
]
|
||||
|
||||
async def sync(self) -> None:
|
||||
async with self._lock:
|
||||
await self._adopt_dropped_in()
|
||||
|
||||
@@ -6,6 +6,7 @@ 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
|
||||
@@ -16,6 +17,7 @@ from watcher.registry import BotRegistry
|
||||
setup_logging()
|
||||
|
||||
RETRY_INTERVAL = 300
|
||||
UPTIME_INTERVAL = 60
|
||||
|
||||
|
||||
async def _retry_loop(registry: BotRegistry) -> None:
|
||||
@@ -27,6 +29,14 @@ async def _retry_loop(registry: BotRegistry) -> None:
|
||||
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())
|
||||
@@ -58,6 +68,7 @@ async def runner() -> None:
|
||||
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:
|
||||
@@ -66,9 +77,10 @@ async def runner() -> None:
|
||||
logger.info("[green]Bot watcher running.[/]")
|
||||
await stop.wait()
|
||||
finally:
|
||||
retry.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await retry
|
||||
for task in (retry, beat):
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
await listener.stop()
|
||||
await registry.close()
|
||||
await container.close()
|
||||
|
||||
Reference in new issue
Block a user