snake-python/server/Server.py

from quart_common.web.logger import build_logger, await_log
from quart_common.web.env import env_bool, env_int

from snakes import SnakeBuilder

from server.database import (
  GameplayDatabase,
  GameplayBackendBuilder,
  StorageLoader,
)
from server.metrics import (
  MetricsStoreBuilder,
  MetricsCollector,
)

import asyncio, signal, logging, os, re, time
from quart import Quart

from server.blueprints import (
  create_battlesnake_blueprint,
  create_metrics_blueprint,
  create_dashboard_blueprint,
)
from server.services import (
  DashboardEventsService,
  DashboardWebSocketHub,
  GameRuntimeService,
  GameplayTrackingService,
  DashboardQueryService,
)

class Server:
  def __init__(self, data_path:str, snake_type:str, storage_type:str, debug:bool=False, check_tls_security:bool=False, metrics_backend:str='memory', metrics_redis_url:str='redis://localhost:6379/0', metrics_ttl_sec:int|None=None, gameplay_db_enabled:bool=True, gameplay_db_backend:str='sqlite', gameplay_db_path:str|None=None, gameplay_db_busy_timeout_ms:int=5000, gameplay_db_pg_dsn:str|None=None):
    self.debug = debug
    self.data_path = data_path

    self.snake_type = snake_type
    self.storage_type = storage_type

    self.config_file = os.path.join(data_path, 'data', 'snake-config.json')
    self.check_tls_security = check_tls_security

    self.store_game_state = False
    metrics_backend_normalized = (metrics_backend or 'memory').strip().lower()
    self.metrics_backend_normalized = metrics_backend_normalized
    self.metrics_redis_url = metrics_redis_url
    self.stale_game_timeout_sec = self._get_stale_game_timeout_sec()

    self.game_runtime = GameRuntimeService(
      snake_type=self.snake_type,
      stale_game_timeout_sec=self.stale_game_timeout_sec,
    )
    self.dashboard_ws_hub = DashboardWebSocketHub()

    self.metrics_collector = MetricsCollector(
      metrics_manager=MetricsStoreBuilder.build(
        backend=metrics_backend_normalized,
        redis_url=metrics_redis_url,
        ttl_seconds=metrics_ttl_sec,
        key_prefix=os.environ.get('METRICS_REDIS_KEY_PREFIX', 'snake:metrics:worker'),
      ),
      metrics_backend=metrics_backend_normalized,
      stale_game_timeout_sec=self.stale_game_timeout_sec,
      game_last_seen_unix=self.game_runtime.game_last_seen_unix,
      game_move_counts=self.game_runtime.game_move_counts,
    )

    self.game_runtime.attach_metrics_collector(self.metrics_collector)
    self._startup_worker_metrics_cleared = False

    self.logger = build_logger('Battlesnake', debug_env_var='DEBUG_SERVER')
    self.snake_version = self._get_snake_version()

    self.gameplay_database = None
    if gameplay_db_enabled:
      db_path = gameplay_db_path or os.path.join(data_path, 'data', 'database', 'gameplay.sqlite3')
      self.gameplay_database = GameplayDatabase(
        backend=GameplayBackendBuilder.build(
          backend=gameplay_db_backend,
          db_path=db_path,
          busy_timeout_ms=gameplay_db_busy_timeout_ms,
          pg_dsn=gameplay_db_pg_dsn,
        )
      )

    self.gameplay_tracking = GameplayTrackingService(
      gameplay_database=self.gameplay_database,
      logger=self.logger,
    )
    self.dashboard_query = DashboardQueryService(
      gameplay_database=self.gameplay_database,
      ws_hub=self.dashboard_ws_hub,
      logger=self.logger,
      dashboard_running_game_stale_sec=600,
    )
    self.dashboard_events_service = DashboardEventsService(
      enabled=(self.metrics_backend_normalized == 'redis' and env_bool('DASHBOARD_EVENTS_ENABLED', True)),
      redis_url=self.metrics_redis_url,
      channel= os.getenv('DASHBOARD_EVENTS_CHANNEL', 'snake:dashboard:events'),
      event_origin=f'worker-{os.getpid()}-{int(time.time() * 1000)}',
      shutdown_event=self.dashboard_ws_hub.shutdown_event,
      on_notice=self._on_dashboard_games_update_notice,
      logger=self.logger,
    )
    self.dashboard_query.set_publish_notice(self.dashboard_events_service.publish_notice)

    self.app = Quart('Battlesnake', template_folder=os.path.join(data_path, 'templates', 'side'), static_folder=os.path.join(data_path, 'templates', 'files'))

    self.app.register_blueprint(create_battlesnake_blueprint(self))
    self.app.register_blueprint(create_metrics_blueprint(self))
    self.app.register_blueprint(create_dashboard_blueprint(self))

    @self.app.after_request
    async def identify_server(response):
      response.headers.set('server', 'battlesnake/gitea/snake-python')
      return response

    @self.app.before_serving
    async def clear_startup_worker_metrics_once():
      if self._startup_worker_metrics_cleared:
        return
      self._startup_worker_metrics_cleared = True

      if env_bool('METRICS_CLEAR_WORKERS_ON_STARTUP', True):
        should_clear = await self.metrics_collector.should_clear_worker_metrics_on_startup(env_int('METRICS_STARTUP_CLEANUP_LOCK_TTL_SEC', 300))
        if should_clear:
          await self.metrics_collector.clear_worker_metrics()
      await self.dashboard_events_service.start_listener()

      if self.gameplay_database is not None:
        await self.gameplay_database.initialize()

    @self.app.after_serving
    async def shutdown_state_storage():
      await self.dashboard_events_service.stop_listener()
      await self.metrics_collector.close()
      if self.gameplay_database is not None:
        await self.gameplay_database.close()

  async def run(self, host:str='0.0.0.0', port:int=8000, debug:bool=False):
    logging.getLogger('werkzeug').setLevel(logging.ERROR)
    loop = asyncio.get_running_loop()

    installed_signal_handlers:list[signal.Signals] = []
    shutdown_event = asyncio.Event()

    def on_shutdown_signal() -> None:
      self.dashboard_ws_hub.request_shutdown()
      shutdown_event.set()

    async def shutdown_trigger() -> None:
      await shutdown_event.wait()

    for shutdown_signal in (signal.SIGINT, signal.SIGTERM):
      try:
        loop.add_signal_handler(shutdown_signal, on_shutdown_signal)
        installed_signal_handlers.append(shutdown_signal)
      except (NotImplementedError, RuntimeError):
        continue

    await await_log(self.logger.info(f'Running Battlesnake at http://{host}:{port} with the {' '.join(re.findall('[A-Z][^A-Z]*', self.snake_type))}'))
    try:
      await self.app.run_task(host=host, port=port, debug=debug, shutdown_trigger=shutdown_trigger)
    finally:
      self.dashboard_ws_hub.request_shutdown()
      for shutdown_signal in installed_signal_handlers:
        try:
          loop.remove_signal_handler(shutdown_signal)
        except Exception:
          continue

  def _get_snake_version(self) -> str:
    configured_version = SnakeBuilder.get_version(self.snake_type)
    if configured_version is None:
      return str(SnakeBuilder.get_version('TemplateSnake'))
    return str(configured_version)

  def _get_stale_game_timeout_sec(self) -> int:
    return max(30, env_int('SNAKE_STUCK_GAME_TIMEOUT_SEC', 180))

  def enable_store_game_state(self):
    self.store_game_state = True

  def _cleanup_database(self):
    storage = StorageLoader.build(self.storage_type)
    return storage.cleanup()

  async def _on_dashboard_games_update_notice(self, trigger:str) -> None:
    await self.dashboard_query.on_dashboard_games_update_notice(trigger)