triangular_arbitrage_bot/fh_ob/__main__.py

87 lines
2.6 KiB
Python

import asyncio
import signal
from pathlib import Path
import structlog
import uvicorn
from common.config import Settings
from common.log import configure_logging
from fh_ob.book_store import BookStore
from fh_ob.rest_server import create_app
from fh_ob.socket_server import SocketServer
from fh_ob.ws_client import KuCoinWSClient
async def main() -> None:
config_path = Path("config.yaml")
settings = await Settings.from_yaml(config_path) if config_path.exists() else Settings()
configure_logging(settings.fh_ob.log_level, settings.fh_ob.log_file)
log = structlog.get_logger().bind(component="fh_ob")
log.info("fh_ob_starting", symbols=settings.fh_ob.symbols)
book_store = BookStore()
socket_server = SocketServer(settings.fh_ob.socket_path)
async def on_book_update(book):
try:
await socket_server.broadcast(book)
except Exception:
pass
ws_client = KuCoinWSClient(
settings=settings.fh_ob,
book_store=book_store,
on_book_update=on_book_update,
)
rest_app = create_app(
book_store,
get_socket_clients=socket_server.client_count,
get_subscribed_count=ws_client.subscribed_count,
is_connected=ws_client.is_connected,
add_symbol=ws_client.add_symbol,
remove_symbol=ws_client.remove_symbol,
get_symbols=ws_client.get_symbols,
get_reconnect_stats=ws_client.reconnect_stats,
)
rest_config = uvicorn.Config(
rest_app,
host=settings.fh_ob.rest_host,
port=settings.fh_ob.rest_port,
log_level="warning",
)
rest_server = uvicorn.Server(rest_config)
async def shutdown(sig: signal.Signals) -> None:
log.info("shutdown_signal_received", signal=sig.name)
await ws_client.stop()
await socket_server.stop()
rest_config.should_exit = True
loop = asyncio.get_running_loop()
for sig in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(sig, lambda s=sig: asyncio.create_task(shutdown(s)))
ws_task = asyncio.create_task(ws_client.start())
socket_task = asyncio.create_task(socket_server.start())
rest_task = asyncio.create_task(rest_server.serve())
log.info(
"fh_ob_ready",
rest_endpoint=f"http://{settings.fh_ob.rest_host}:{settings.fh_ob.rest_port}",
socket_path=str(settings.fh_ob.socket_path),
)
try:
await asyncio.gather(ws_task, socket_task, rest_task)
except asyncio.CancelledError:
log.info("fh_ob_cancelled")
except Exception as e:
log.error("fh_ob_error", error=str(e))
raise
if __name__ == "__main__":
asyncio.run(main())