kis_bot/kis_trader/execution/order_worker.py

"""계좌 단일 주문 실행 Worker — 4전략 BUY/SELL intent 를 우선순위 FIFO 로 직렬 place."""
from __future__ import annotations

import itertools
import queue
import threading
import time
from typing import Any, Dict, Optional, TYPE_CHECKING

from ..utils.env import get_env_bool, get_env_int
from ..utils.logger import get_logger
from .order_intent import OrderIntent, resolve_order_priority

if TYPE_CHECKING:
    from ..strategies.base import BaseStrategy
    from .order_manager import OrderManager

logger = get_logger("kis_trader.order_worker")

_SENTINEL = object()


class AccountOrderWorker:
    """증권 계좌 1개에 맞춘 주문 대기줄 — 한 번에 place() 1건만."""

    def __init__(self, order_mgr: "OrderManager") -> None:
        self.order_mgr = order_mgr
        self._pq: "queue.PriorityQueue[tuple]" = queue.PriorityQueue()
        self._seq = itertools.count()
        self._strategies: Dict[str, Any] = {}
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._start_lock = threading.Lock()
        self.enqueued_total = 0
        self.processed_total = 0
        self.dropped_total = 0

    def register_strategy(self, strategy: "BaseStrategy") -> None:
        sid = str(getattr(strategy, "strategy_id", "") or "").strip()
        if sid:
            self._strategies[sid] = strategy

    def start(self) -> None:
        with self._start_lock:
            if self._thread is not None and self._thread.is_alive():
                return
            self._running = True
            th = threading.Thread(
                target=self._worker_loop,
                name="AccountOrderWorker",
                daemon=True,
            )
            self._thread = th
            th.start()
            logger.info("▶ [OrderWorker] 계좌 단일 주문 큐 기동")

    def stop(self) -> None:
        self._running = False
        try:
            self._pq.put_nowait((9999, next(self._seq), _SENTINEL))
        except Exception:
            pass
        th = self._thread
        if th is not None:
            th.join(timeout=5.0)

    def queue_depth(self) -> int:
        try:
            return int(self._pq.qsize())
        except Exception:
            return 0

    def enqueue(
        self,
        strategy: "BaseStrategy",
        side: str,
        signal: Dict[str, Any],
        *,
        source: str = "scan",
        priority: Optional[int] = None,
    ) -> bool:
        """intent 큐 적재. False=큐 만료 drop (inflight 는 호출자가 해제)."""
        if not self._running:
            return False
        max_q = max(1, int(get_env_int("ORDER_WORKER_MAX_QUEUE", 500) or 500))
        if self.queue_depth() >= max_q:
            self.dropped_total += 1
            code = str((signal or {}).get("code") or "")
            logger.warning(
                "⚠️ [OrderWorker] 큐 만료 drop [%s] %s %s depth>=%d",
                getattr(strategy, "strategy_id", "?"),
                side,
                code,
                max_q,
            )
            return False
        prio = (
            int(priority)
            if priority is not None
            else resolve_order_priority(side, signal or {}, source)
        )
        intent = OrderIntent(
            priority=prio,
            seq=next(self._seq),
            strategy_id=str(getattr(strategy, "strategy_id", "") or ""),
            side=str(side or "").upper(),
            signal=dict(signal or {}),
            source=str(source or "scan"),
        )
        self._pq.put((intent.priority, intent.seq, intent))
        self.enqueued_total += 1
        return True

    def _worker_loop(self) -> None:
        while True:
            try:
                _prio, _seq, item = self._pq.get(timeout=0.3)
            except queue.Empty:
                if not self._running:
                    break
                continue
            if item is _SENTINEL:
                break
            if not isinstance(item, OrderIntent):
                continue
            if not self._running:
                self._release_inflight_for_intent(item)
                continue
            self._process_intent(item)

    def _process_intent(self, intent: OrderIntent) -> None:
        strat = self._strategies.get(intent.strategy_id)
        code = str((intent.signal or {}).get("code") or "").strip()
        side = str(intent.side or "").upper()
        t0 = time.perf_counter()
        try:
            if strat is None:
                logger.warning(
                    "⚠️ [OrderWorker] 미등록 전략 %s — skip %s %s",
                    intent.strategy_id,
                    side,
                    code,
                )
                return
            if side == "SELL" and get_env_bool("REAL_BALANCE_VERIFY_BEFORE_SELL", True):
                try:
                    self.order_mgr.prefetch_broker_holdings(force=False)
                except Exception:
                    pass
            if side == "BUY":
                strat._submit_buy(intent.signal)
            elif side == "SELL":
                strat._submit_sell(intent.signal)
            else:
                logger.warning(
                    "⚠️ [OrderWorker] invalid side=%s [%s] %s",
                    side,
                    intent.strategy_id,
                    code,
                )
        except Exception as ex:
            logger.warning(
                "⚠️ [OrderWorker] place 예외 [%s] %s %s: %s",
                intent.strategy_id,
                side,
                code or "-",
                ex,
            )
        finally:
            self._release_inflight_for_intent(intent)
            self.processed_total += 1
            elapsed_ms = (time.perf_counter() - t0) * 1000.0
            if elapsed_ms >= 3000.0:
                logger.warning(
                    "⚠️ [OrderWorker] 처리 지연 [%s] %s %s src=%s wait+place=%.0fms",
                    intent.strategy_id,
                    side,
                    code or "-",
                    intent.source,
                    elapsed_ms,
                )

    def _release_inflight_for_intent(self, intent: OrderIntent) -> None:
        strat = self._strategies.get(intent.strategy_id)
        if strat is None:
            return
        code = str((intent.signal or {}).get("code") or "").strip()
        if not code:
            return
        side = str(intent.side or "").upper()
        if side == "SELL":
            release = getattr(strat, "_release_sell_inflight", None)
            if callable(release):
                release(code)
        elif side == "BUY":
            release = getattr(strat, "_release_buy_inflight", None)
            if callable(release):
                release(code)