fix(userbot): stop gracefully on sigterm and flush session peers periodically

This commit is contained in:
hh
2026-08-30 01:56:10 +02:00
parent 8f7d47476c
commit cf277c56cf
+16 -1
View File
@@ -1,5 +1,6 @@
import asyncio
import contextlib
import signal
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from pathlib import Path
@@ -26,6 +27,8 @@ from utils.storage import ContentAddressedStorage
setup_logging()
SESSION_FLUSH_INTERVAL = 60
@dataclass
class RunningAccount:
@@ -149,6 +152,14 @@ class AccountRegistry:
logger.info(f"[yellow]Account logged out:[/] {path.stem}")
async def _flush_sessions(registry: AccountRegistry) -> None:
while True:
await asyncio.sleep(SESSION_FLUSH_INTERVAL)
for client in registry.clients:
with contextlib.suppress(Exception):
await client.storage.save()
async def _listen_changes(
registry: AccountRegistry, tasks: set[asyncio.Task]
) -> asyncpg.Connection:
@@ -193,14 +204,18 @@ async def runner() -> None:
registry = AccountRegistry(pool, storage)
tasks: set[asyncio.Task] = set()
listen_conn: asyncpg.Connection | None = None
stop = asyncio.Event()
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
flush = asyncio.create_task(_flush_sessions(registry))
try:
await registry.sync()
if not registry.clients:
logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
listen_conn = await _listen_changes(registry, tasks)
logger.info("[green]Userbot running.[/]")
await asyncio.Event().wait()
await stop.wait()
finally:
await _cancel(flush)
if listen_conn is not None:
with contextlib.suppress(Exception):
await listen_conn.close()