trex-core/trex_core/runtime/daemon.py

"""Slot-core daemon entrypoint (REQ-005, REQ-012/013/014, BAT-007).

Sync + threaded runtime. HTTP `/v1` API has been removed (REQ-025); all
strategy-facing communication runs over RabbitMQ (REQ-021/022/023).

Threads (started in `run()`):
  - intent RPC consumer    — `trex.intent.req`,  prefetch=1, FIFO, dedup via intent_log
  - snapshot RPC consumer  — `trex.snapshot.req` (also refreshes the watchdog)
  - control RPC consumer    — `trex.control.req` (pause/resume/list + kill_switch/resume_kill_switch, REQ-032/043)
  - execution loop         — broker dispatch + poll + slot transitions + event emit
  - ledger writer          — ledger outbox drain (REQ-019)
  - snapshot scheduler      — daily portfolio-hub snapshot refresh (REQ-030)
  - preallocation scheduler — daily slot preallocation in trading TZ (REQ-005, dcy-014)
  - safety-sell worker      — forced_close / stop-loss·take-profit / market-close (REQ-037/038/039)

Each consumer/publisher uses its OWN RabbitMQTransport instance because
pika BlockingConnection is not thread-safe across channels.
Console script: `trex-core-daemon`.
"""
from __future__ import annotations

import logging
import signal
import threading
from datetime import time as time_cls
from zoneinfo import ZoneInfo

from trex_core.adapters.broker.paper import PaperBroker
from trex_core.adapters.ledger import InMemoryLedger, PortfolioHubLedger
from trex_core.adapters.messaging import (
    InMemoryTransport,
    RabbitMQTransport,
    topology,
)
from trex_core.adapters.persistence.migrate import apply_migrations
from trex_core.adapters.persistence.postgresql import (
    PostgreSQLIntentStore,
    PostgreSQLLedgerOutbox,
    PostgreSQLSlotRepository,
)
from trex_core.adapters.portfolio import (
    PortfolioHubReadAdapter,
    PostgreSQLPauseStore,
    StubPortfolioRead,
)
from trex_core.adapters.portfolio.killswitch_store import PostgreSQLKillSwitchStore
from trex_core.application.broker_port import BrokerPort
from trex_core.application.execution import ExecutionService
from trex_core.application.intents import IntentService
from trex_core.application.ledger_outbox import LedgerWriter
from trex_core.application.ledger_port import LedgerPort
from trex_core.application.messaging import MessagingTransport
from trex_core.application.killswitch import KillSwitchService
from trex_core.application.messaging_handlers import (
    make_control_rpc_handler,
    make_intent_rpc_handler,
    make_snapshot_rpc_handler,
    slot_transition_event_payload,
)
from trex_core.application.portfolio import (
    PortfolioHubReadPort,
    SnapshotStore,
)
from trex_core.application.safety_sell import SafetySellService
from trex_core.application.preallocation_scheduler import DailyPreallocationScheduler
from trex_core.application.snapshot_scheduler import SnapshotScheduler
from trex_core.application.watchdog import WatchdogStore
from trex_core.domain.slot import SlotService
from trex_core.runtime.config import DaemonConfig

_log = logging.getLogger(__name__)


def _build_broker(config: DaemonConfig) -> BrokerPort:
    if config.broker == "kis":
        from trex_core.adapters.broker.kis import KISBroker

        k = config.kis
        token_cache = None
        if k.redis_url:
            from trex_core.adapters.broker.token_cache import RedisTokenCache

            token_cache = RedisTokenCache(k.redis_url)
        return KISBroker(
            k.account_ref, app_key=k.app_key, app_secret=k.app_secret,
            base_url=k.base_url, timeout=k.timeout, allow_live=k.allow_live,
            token_cache=token_cache,
        )
    return PaperBroker(
        fill_after_polls=config.paper.fill_after_polls,
        qty_per_poll=config.paper.qty_per_poll,
        reject_symbols=set(config.paper.reject_symbols),
    )


def _build_ledger(config: DaemonConfig) -> LedgerPort:
    """portfolio-hub when configured; otherwise in-memory stub (REQ-017)."""
    pf = config.portfolio_hub
    if pf is None:
        return InMemoryLedger()
    return PortfolioHubLedger(pf.base_url, token=pf.token or None, timeout=pf.timeout)


def _build_transport(config: DaemonConfig) -> MessagingTransport:
    """RabbitMQ when configured; otherwise in-memory stub (REQ-022)."""
    mq = config.rabbitmq
    if mq is None:
        # Stub: lets the daemon boot in environments without a broker (CI,
        # local dev). REQ-024 degrade policy still applies — call_rpc raises.
        return InMemoryTransport()
    return RabbitMQTransport(mq.url, timeout=mq.timeout)


def _build_portfolio_read(config: DaemonConfig) -> PortfolioHubReadPort:
    """portfolio-hub HTTP read adapter when configured; otherwise stub (REQ-027)."""
    pf = config.portfolio_hub
    if pf is None:
        return StubPortfolioRead()
    return PortfolioHubReadAdapter(
        pf.base_url, token=pf.token or None, timeout=pf.timeout,
    )


def _parse_kst_time(spec: str) -> time_cls:
    parts = spec.split(":")
    h = int(parts[0])
    m = int(parts[1]) if len(parts) > 1 else 0
    return time_cls(hour=h, minute=m)


class Daemon:
    """Thread-orchestrating daemon (no HTTP)."""

    def __init__(self, config: DaemonConfig) -> None:
        self.config = config
        self.stop = threading.Event()
        self._threads: list[threading.Thread] = []
        self._transports: list[MessagingTransport] = []

        # Persistence + services
        self.slot_repo = PostgreSQLSlotRepository(config.dsn)
        self.intent_store = PostgreSQLIntentStore(config.dsn)
        self.slot_service = SlotService(self.slot_repo)
        self.intent_service = IntentService(
            self.slot_service, self.slot_repo, self.intent_store,
            resolve_symbol=self._resolve_symbol,
        )
        self.broker = _build_broker(config)
        self.outbox = PostgreSQLLedgerOutbox(config.dsn)
        # BAT-008: portfolio-hub read + snapshot store + durable pause + scheduler.
        # (Created before the ledger writer so option-D write-back can target it.)
        self.snapshot_store = SnapshotStore()
        self.ledger_writer = LedgerWriter(
            self.outbox, _build_ledger(config),
            snapshot_store=self.snapshot_store,  # dcy-018 / REQ-030: option D write-back
            retry_base_sec=config.ledger_retry_base_sec,
            retry_cap_sec=config.ledger_retry_cap_sec,
            max_attempts=config.ledger_max_attempts,
        )
        # BAT-010/011: per-slot watchdog (snapshot RPC writer, safety-sell reader).
        self.watchdog = WatchdogStore()
        self.pause_store = PostgreSQLPauseStore(config.dsn)
        # BAT-012: day-scoped kill-switch exclusion store + service.
        self.killswitch_store = PostgreSQLKillSwitchStore(config.dsn)
        self.kill_switch = KillSwitchService(
            slot_repo=self.slot_repo,
            slot_service=self.slot_service,
            killswitch_store=self.killswitch_store,
        )
        self.portfolio_read = _build_portfolio_read(config)
        self.snapshot_scheduler = SnapshotScheduler(
            read_port=self.portfolio_read,
            store=self.snapshot_store,
            broker=self.broker,
            schedule_kst=_parse_kst_time(config.snapshot_schedule.kst_time),
        )
        # dcy-014 / REQ-005: daily slot preallocation in the configured
        # trading timezone (default KST), so trading_date matches the local
        # calendar and the next day's slots appear without a restart.
        self.preallocation_scheduler = DailyPreallocationScheduler(
            slot_service=self.slot_service,
            total_slots=config.total_slots,
            tz=ZoneInfo(config.trading_timezone),
            schedule_time=_parse_kst_time(config.daily_preallocation_time),
        )

        # Event publisher transport — one dedicated instance for the
        # execution-loop thread (pika is not thread-safe).
        self.event_publisher: MessagingTransport = _build_transport(config)
        self._transports.append(self.event_publisher)

        self.execution = ExecutionService(
            self.slot_service, self.slot_repo, self.broker,
            ledger_outbox=self.outbox,
            event_emitter=self._make_event_emitter(),
            stale_pending_threshold_sec=config.pending_stale_threshold_sec,
        )

        # BAT-011: safety-sell trigger service (forced_close / SL-TP / market-close).
        self.safety_sell = SafetySellService(
            slot_repo=self.slot_repo,
            slot_service=self.slot_service,
            broker=self.broker,
            watchdog=self.watchdog,
            config=config.safety_sell,
            instance_key=config.instance_key,
        )

    # ------------------------------------------------------------ buy symbol resolve
    def _resolve_symbol(self, portfolio_id: str):
        """dcy-011 / REQ-007: resolve symbol from the universe snapshot for a
        portfolio_id at buy admission (None → unknown_portfolio reject).

        SnapshotStore.get returns a SnapshotAsset wrapping the PortfolioItem;
        the symbol lives on `.item.symbol`."""
        asset = self.snapshot_store.get(portfolio_id)
        return asset.item.symbol if asset is not None else None

    # ------------------------------------------------------------ event hook
    def _make_event_emitter(self):
        def emit(event_type: str, slot) -> None:
            payload = slot_transition_event_payload(slot, event_type=event_type)
            rk = topology.event_routing_key(
                event_type, slot.strategy_id or "_", slot.symbol or "_",
            )
            try:
                self.event_publisher.publish_event(rk, payload)
            except Exception:  # noqa: BLE001 - publish failure must not roll back
                _log.warning("event publish failed (rk=%s)", rk, exc_info=True)
        return emit

    # ------------------------------------------------------------ workers
    def _execution_loop(self) -> None:
        while not self.stop.wait(self.config.execution_interval_sec):
            try:
                self.execution.run_once()
            except Exception:  # noqa: BLE001 - loop must survive transient errors
                continue

    def _ledger_loop(self) -> None:
        while not self.stop.wait(self.config.ledger_interval_sec):
            try:
                self.ledger_writer.run_once()
            except Exception:  # noqa: BLE001
                continue

    def _intent_rpc_loop(self) -> None:
        transport = _build_transport(self.config)
        self._transports.append(transport)
        try:
            transport.declare_topology(self.config.instance_key)
        except Exception:  # noqa: BLE001
            pass
        handler = make_intent_rpc_handler(
            self.intent_service,
            slot_service=self.slot_service,
            watchdog_store=self.watchdog,
        )
        transport.consume_rpc_forever(
            topology.intent_queue(self.config.instance_key), handler, self.stop)

    def _snapshot_rpc_loop(self) -> None:
        transport = _build_transport(self.config)
        self._transports.append(transport)
        try:
            transport.declare_topology(self.config.instance_key)
        except Exception:  # noqa: BLE001
            pass
        handler = make_snapshot_rpc_handler(
            self.slot_repo, self.slot_service,
            snapshot_store=self.snapshot_store,
            pause_store=self.pause_store,
            default_total_slots=self.config.total_slots,
            watchdog_store=self.watchdog,
            watchdog_timeout_sec=self.config.safety_sell.watchdog_timeout_sec,
            killswitch_store=self.killswitch_store,
            ledger_outbox=self.outbox,  # dcy-018 / REQ-029: estimated-cash overlay
        )
        transport.consume_rpc_forever(
            topology.snapshot_queue(self.config.instance_key), handler, self.stop)

    def _control_rpc_loop(self) -> None:
        # REQ-043: operator control channel (pause/resume/list + kill_switch/resume_kill_switch).
        transport = _build_transport(self.config)
        self._transports.append(transport)
        try:
            transport.declare_topology(self.config.instance_key)
        except Exception:  # noqa: BLE001
            pass
        handler = make_control_rpc_handler(
            self.pause_store, kill_switch_service=self.kill_switch,
        )
        transport.consume_rpc_forever(
            topology.control_queue(self.config.instance_key), handler, self.stop)

    def _safety_sell_loop(self) -> None:
        # REQ-037/038/039: forced_close / stop-loss·take-profit / market-close.
        while not self.stop.wait(self.safety_sell.base_tick_sec):
            try:
                self.safety_sell.run_once()
            except Exception:  # noqa: BLE001 - loop must survive transient errors
                continue

    def _snapshot_scheduler_loop(self) -> None:
        try:
            self.snapshot_scheduler.run_forever(
                self.stop, tick_sec=self.config.snapshot_schedule.tick_sec,
            )
        except Exception:  # noqa: BLE001
            _log.exception("snapshot scheduler loop crashed")

    def _preallocation_loop(self) -> None:
        try:
            self.preallocation_scheduler.run_forever(self.stop, tick_sec=60.0)
        except Exception:  # noqa: BLE001
            _log.exception("preallocation scheduler loop crashed")

    # ------------------------------------------------------------ lifecycle
    def setup(self) -> None:
        apply_migrations(self.config.dsn)
        # dcy-014 / REQ-005: preallocate against the configured trading timezone
        # (default KST), not the system UTC date. Boot-time run_once ensures the
        # current TZ-day's slots exist immediately (idempotent).
        self.preallocation_scheduler.run_once()
        self.execution.reconcile()  # REQ-014 restart recovery
        try:
            self.event_publisher.declare_topology(self.config.instance_key)
        except Exception:  # noqa: BLE001
            _log.warning("event publisher topology declare failed", exc_info=True)

    def start(self) -> None:
        workers = [
            ("trex-execution", self._execution_loop),
            ("trex-ledger", self._ledger_loop),
            ("trex-intent-rpc", self._intent_rpc_loop),
            ("trex-snapshot-rpc", self._snapshot_rpc_loop),
            ("trex-control-rpc", self._control_rpc_loop),
            ("trex-snapshot-scheduler", self._snapshot_scheduler_loop),
            ("trex-preallocation", self._preallocation_loop),
            ("trex-safety-sell", self._safety_sell_loop),
        ]
        for name, target in workers:
            t = threading.Thread(target=target, name=name, daemon=True)
            t.start()
            self._threads.append(t)

    def shutdown(self) -> None:
        self.stop.set()
        for t in self._threads:
            t.join(
                timeout=max(self.config.execution_interval_sec, self.config.ledger_interval_sec) * 3
            )
        for tr in self._transports:
            try:
                tr.close()
            except Exception:  # noqa: BLE001
                pass

    def run(self) -> None:  # pragma: no cover - process entrypoint
        self.setup()
        self.start()
        signal.signal(signal.SIGINT, lambda *_: self.stop.set())
        signal.signal(signal.SIGTERM, lambda *_: self.stop.set())
        try:
            while not self.stop.is_set():
                self.stop.wait(1.0)
        finally:
            self.shutdown()


def main() -> None:  # pragma: no cover - process entrypoint
    import argparse

    logging.basicConfig(level=logging.INFO)
    parser = argparse.ArgumentParser(prog="trex-core-daemon")
    parser.add_argument(
        "--config", required=True,
        help="per-instance TOML config path (instance_key, portfolio_hub_token). "
             "Required — no default (REQ-048 / dcy-009 OQ-2).",
    )
    args = parser.parse_args()
    config = DaemonConfig.from_config(args.config)
    Daemon(config).run()


if __name__ == "__main__":  # pragma: no cover
    main()