fix(userbot): stop gracefully on sigterm and flush session peers periodically
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user