99 lines
2.9 KiB
Python
99 lines
2.9 KiB
Python
|
"""Channel worker entrypoint with embedded health endpoint."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import os
|
||
|
|
import signal
|
||
|
|
import sys
|
||
|
|
|
||
|
|
if sys.platform == "win32":
|
||
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
|
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
||
|
|
|
||
|
|
from aiohttp import web
|
||
|
|
|
||
|
|
from yuxi.channel.infrastructure.agent.agent_adapter import AgentAdapter
|
||
|
|
from yuxi.channel.startup import init_channel, shutdown_channel
|
||
|
|
from yuxi.storage.postgres.manager import pg_manager
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
_HEALTH_PORT = int(os.getenv("CHANNEL_HEALTH_PORT", "5051"))
|
||
|
|
|
||
|
|
|
||
|
|
async def _health_handler(request: web.Request) -> web.Response:
|
||
|
|
checks: dict[str, str] = {"status": "ok"}
|
||
|
|
container = request.app["container"]
|
||
|
|
|
||
|
|
try:
|
||
|
|
redis = container.cache_port
|
||
|
|
await redis.ping()
|
||
|
|
checks["redis"] = "ok"
|
||
|
|
except Exception as e:
|
||
|
|
checks["redis"] = f"error: {e}"
|
||
|
|
checks["status"] = "degraded"
|
||
|
|
|
||
|
|
worker_pool = getattr(container, "worker_pool", None)
|
||
|
|
if worker_pool and getattr(worker_pool, "is_running", False):
|
||
|
|
checks["worker_pool"] = "ok"
|
||
|
|
else:
|
||
|
|
checks["worker_pool"] = "not_running"
|
||
|
|
checks["status"] = "degraded"
|
||
|
|
|
||
|
|
status_code = 200 if checks["status"] == "ok" else 503
|
||
|
|
return web.json_response(checks, status=status_code)
|
||
|
|
|
||
|
|
|
||
|
|
async def _start_health_server(container: object) -> web.AppRunner:
|
||
|
|
app = web.Application()
|
||
|
|
app["container"] = container
|
||
|
|
app.router.add_get("/healthz", _health_handler)
|
||
|
|
|
||
|
|
runner = web.AppRunner(app)
|
||
|
|
await runner.setup()
|
||
|
|
site = web.TCPSite(runner, "0.0.0.0", _HEALTH_PORT)
|
||
|
|
await site.start()
|
||
|
|
logger.info("health endpoint started on port %d", _HEALTH_PORT)
|
||
|
|
return runner
|
||
|
|
|
||
|
|
|
||
|
|
async def main():
|
||
|
|
redis_url = os.getenv("REDIS_URL", "redis://redis:6379/0")
|
||
|
|
config_path = os.getenv("CHANNEL_CONFIG_PATH", "channel_config.yaml")
|
||
|
|
|
||
|
|
logger.info("starting channel worker, redis=%s", redis_url)
|
||
|
|
|
||
|
|
container = await init_channel(
|
||
|
|
redis_url=redis_url,
|
||
|
|
agent_port=AgentAdapter(pg_manager.get_async_session_context),
|
||
|
|
config_yaml_path=config_path,
|
||
|
|
)
|
||
|
|
|
||
|
|
health_runner = await _start_health_server(container)
|
||
|
|
|
||
|
|
stop_event = asyncio.Event()
|
||
|
|
|
||
|
|
def _signal_handler():
|
||
|
|
stop_event.set()
|
||
|
|
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
if sys.platform == "win32":
|
||
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
||
|
|
signal.signal(sig, lambda *_: _signal_handler())
|
||
|
|
else:
|
||
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
||
|
|
loop.add_signal_handler(sig, _signal_handler)
|
||
|
|
|
||
|
|
logger.info("channel worker running, waiting for stop signal")
|
||
|
|
await stop_event.wait()
|
||
|
|
|
||
|
|
await health_runner.cleanup()
|
||
|
|
logger.info("health endpoint stopped")
|
||
|
|
|
||
|
|
await shutdown_channel(container)
|
||
|
|
logger.info("channel worker stopped")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(main())
|