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 asyncio
import contextlib import contextlib
import signal
from collections.abc import Callable, Coroutine from collections.abc import Callable, Coroutine
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
@@ -26,6 +27,8 @@ from utils.storage import ContentAddressedStorage
setup_logging() setup_logging()
SESSION_FLUSH_INTERVAL = 60
@dataclass @dataclass
class RunningAccount: class RunningAccount:
@@ -149,6 +152,14 @@ class AccountRegistry:
logger.info(f"[yellow]Account logged out:[/] {path.stem}") 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( async def _listen_changes(
registry: AccountRegistry, tasks: set[asyncio.Task] registry: AccountRegistry, tasks: set[asyncio.Task]
) -> asyncpg.Connection: ) -> asyncpg.Connection:
@@ -193,14 +204,18 @@ async def runner() -> None:
registry = AccountRegistry(pool, storage) registry = AccountRegistry(pool, storage)
tasks: set[asyncio.Task] = set() tasks: set[asyncio.Task] = set()
listen_conn: asyncpg.Connection | None = None 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: try:
await registry.sync() await registry.sync()
if not registry.clients: if not registry.clients:
logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]") logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
listen_conn = await _listen_changes(registry, tasks) listen_conn = await _listen_changes(registry, tasks)
logger.info("[green]Userbot running.[/]") logger.info("[green]Userbot running.[/]")
await asyncio.Event().wait() await stop.wait()
finally: finally:
await _cancel(flush)
if listen_conn is not None: if listen_conn is not None:
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
await listen_conn.close() await listen_conn.close()