trex-core/trex_core/adapters/messaging/rabbitmq.py

"""RabbitMQ (AMQP) transport adapter — pika BlockingConnection (REQ-022/023/024).

pika is NOT thread-safe across channels/connections, so every consumer or
publisher uses its own connection/channel. The adapter owns one connection
for *its* thread; the daemon creates multiple adapter instances (one per
consumer/publisher worker thread).

Design highlights:
  - Reconnect backoff (1s -> 2s -> 4s -> 8s -> 16s -> 30s cap), implemented
    in the consume_*_forever loops. (pika has no built-in auto-reconnect.)
  - Publisher confirms (`channel.confirm_delivery`) + mandatory=True so an
    unroutable message NACKs via UnroutableError → DLQ.
  - Manual ack only AFTER the handler returns its response. If the handler
    raises, the message is NACKed with requeue=False so it goes to the DLQ
    instead of poison-looping (REQ-023 AC-6).
  - intent RPC consumer uses prefetch=1 (REQ-023 AC-3) for FIFO single-writer
    alignment.
"""
from __future__ import annotations

import json
import logging
import threading
import time
from typing import Optional

import pika
import pika.exceptions

from trex_core.adapters.messaging.topology import (
    BROADCAST_EXCHANGE,
    DEFAULT_INSTANCE_KEY,
    DIRECT_EXCHANGE,
    EVENTS_DLQ_EXCHANGE,
    EVENTS_DLQ_QUEUE,
    EVENTS_EXCHANGE,
    INTENT_REQ_QUEUE,
    PAUSE_REQ_QUEUE,
    STRATEGY_QUEUE_ARGS,
    control_queue,
    direct_routing_key,
    intent_queue,
    snapshot_queue,
)
from trex_core.application.messaging import (
    EventHandler,
    RpcHandler,
    RpcRequest,
    RpcTimeoutError,
    TransportHealth,
    TransportUnavailableError,
)

_log = logging.getLogger(__name__)

_BACKOFF_SEQ = (1.0, 2.0, 4.0, 8.0, 16.0, 30.0)


class RabbitMQTransport:
    def __init__(
        self,
        url: str,
        *,
        timeout: float = 5.0,
        intent_prefetch: int = 1,
        snapshot_prefetch: int = 8,
    ) -> None:
        self._url = url
        self._timeout = timeout
        self._intent_prefetch = intent_prefetch
        self._snapshot_prefetch = snapshot_prefetch
        self._lock = threading.Lock()
        self._conn: Optional[pika.BlockingConnection] = None
        self._channel: Optional[pika.adapters.blocking_connection.BlockingChannel] = None
        self._health = TransportHealth.DISCONNECTED

    # ------------------------------------------------------------ lifecycle
    def _connect(self) -> None:
        params = pika.URLParameters(self._url)
        params.socket_timeout = self._timeout
        params.blocked_connection_timeout = self._timeout
        params.heartbeat = 30
        with self._lock:
            self._health = TransportHealth.RECONNECTING
        conn = pika.BlockingConnection(params)
        ch = conn.channel()
        ch.confirm_delivery()  # publisher confirms (REQ-023 AC-6)
        with self._lock:
            self._conn = conn
            self._channel = ch
            self._health = TransportHealth.CONNECTED

    def _ensure_connected(self) -> None:
        with self._lock:
            ok = (
                self._conn is not None
                and self._conn.is_open
                and self._channel is not None
                and self._channel.is_open
            )
        if not ok:
            self._connect()

    def close(self) -> None:
        with self._lock:
            ch, conn = self._channel, self._conn
            self._channel = None
            self._conn = None
            self._health = TransportHealth.DISCONNECTED
        try:
            if ch is not None and ch.is_open:
                ch.close()
        except Exception:  # noqa: BLE001
            pass
        try:
            if conn is not None and conn.is_open:
                conn.close()
        except Exception:  # noqa: BLE001
            pass

    def health(self) -> TransportHealth:
        with self._lock:
            return self._health

    # ------------------------------------------------------------ topology
    def declare_topology(self, instance_key: str = DEFAULT_INSTANCE_KEY) -> None:
        self._ensure_connected()
        ch = self._channel
        assert ch is not None
        dlx = {"x-dead-letter-exchange": EVENTS_DLQ_EXCHANGE}
        # Events topic + DLQ
        ch.exchange_declare(EVENTS_EXCHANGE, exchange_type="topic", durable=True)
        ch.exchange_declare(EVENTS_DLQ_EXCHANGE, exchange_type="fanout", durable=True)
        ch.queue_declare(EVENTS_DLQ_QUEUE, durable=True)
        ch.queue_bind(EVENTS_DLQ_QUEUE, EVENTS_DLQ_EXCHANGE)
        # BAT-013 / REQ-050: multi-instance control exchanges.
        ch.exchange_declare(DIRECT_EXCHANGE, exchange_type="topic", durable=True)
        ch.exchange_declare(BROADCAST_EXCHANGE, exchange_type="fanout", durable=True)
        # Per-instance RPC queues (default → literal names, REQ-049 backward-compat).
        iq = intent_queue(instance_key)
        sq = snapshot_queue(instance_key)
        cq = control_queue(instance_key)
        for q, channel_type in ((iq, "intent"), (sq, "snapshot"), (cq, "control")):
            ch.queue_declare(q, durable=True, arguments=dlx)
            # Directed binding: trex.direct → {instance_key}.{channel_type}
            ch.queue_bind(q, DIRECT_EXCHANGE,
                          routing_key=direct_routing_key(instance_key, channel_type))
        # Control queue is also fanout-bound so global broadcasts hit every instance.
        ch.queue_bind(cq, BROADCAST_EXCHANGE)
        # Legacy pause queue — declared for the default instance only (REQ-032 alias).
        if instance_key == DEFAULT_INSTANCE_KEY:
            ch.queue_declare(PAUSE_REQ_QUEUE, durable=True, arguments=dlx)

    def declare_strategy_queue(self, queue: str, binding_pattern: str) -> None:
        """Per-strategy durable queue (REQ-023 AC-4) bound to events exchange."""
        self._ensure_connected()
        ch = self._channel
        assert ch is not None
        ch.queue_declare(queue, durable=True, arguments=dict(STRATEGY_QUEUE_ARGS))
        ch.queue_bind(queue, EVENTS_EXCHANGE, routing_key=binding_pattern)

    # ------------------------------------------------------------ publish
    def publish_event(
        self,
        routing_key: str,
        payload: dict,
        *,
        mandatory: bool = True,
    ) -> None:
        body = json.dumps(payload).encode("utf-8")
        try:
            self._ensure_connected()
            ch = self._channel
            assert ch is not None
            ch.basic_publish(
                exchange=EVENTS_EXCHANGE,
                routing_key=routing_key,
                body=body,
                properties=pika.BasicProperties(
                    content_type="application/json",
                    delivery_mode=2,  # persistent
                ),
                mandatory=mandatory,
            )
        except pika.exceptions.UnroutableError:
            # No matching queue subscribers → push to DLQ collector (REQ-023 AC-6).
            self._publish_dlq(routing_key, payload, reason="unroutable")
        except (pika.exceptions.AMQPError, OSError) as exc:
            self._mark_disconnected()
            raise TransportUnavailableError(
                f"publish_event failed: {exc.__class__.__name__}: {exc}"
            ) from exc

    def _publish_dlq(self, routing_key: str, payload: dict, *, reason: str) -> None:
        try:
            self._ensure_connected()
            ch = self._channel
            assert ch is not None
            dlq_body = json.dumps(
                {"reason": reason, "original_routing_key": routing_key, "payload": payload}
            ).encode("utf-8")
            ch.basic_publish(
                exchange=EVENTS_DLQ_EXCHANGE,
                routing_key="",
                body=dlq_body,
                properties=pika.BasicProperties(
                    content_type="application/json", delivery_mode=2
                ),
            )
        except Exception:  # noqa: BLE001 - DLQ best-effort
            _log.warning("DLQ publish failed", exc_info=True)

    # ------------------------------------------------------------ RPC client
    def call_rpc(
        self,
        queue: str,
        payload: dict,
        *,
        correlation_id: str,
        timeout: float,
    ) -> dict:
        try:
            self._ensure_connected()
        except (pika.exceptions.AMQPError, OSError) as exc:
            self._mark_disconnected()
            raise TransportUnavailableError(f"call_rpc connect failed: {exc}") from exc

        # We open a dedicated short-lived channel for this RPC to keep the
        # caller-thread reply isolated (pika BlockingConnection requires
        # consume from same thread that owns the channel).
        params = pika.URLParameters(self._url)
        params.socket_timeout = self._timeout
        try:
            conn = pika.BlockingConnection(params)
        except (pika.exceptions.AMQPError, OSError) as exc:
            self._mark_disconnected()
            raise TransportUnavailableError(f"call_rpc connect failed: {exc}") from exc

        try:
            ch = conn.channel()
            reply_q = ch.queue_declare(queue="", exclusive=True, auto_delete=True).method.queue

            holder: dict[str, dict] = {}

            def on_reply(_ch, method, props, body):
                if props.correlation_id == correlation_id:
                    holder["payload"] = json.loads(body)
                    _ch.basic_ack(delivery_tag=method.delivery_tag)
                    _ch.stop_consuming()

            ch.basic_consume(queue=reply_q, on_message_callback=on_reply, auto_ack=False)
            ch.basic_publish(
                exchange="",
                routing_key=queue,
                body=json.dumps(payload).encode("utf-8"),
                properties=pika.BasicProperties(
                    content_type="application/json",
                    delivery_mode=2,
                    correlation_id=correlation_id,
                    reply_to=reply_q,
                ),
                mandatory=True,
            )
            deadline = time.monotonic() + timeout
            while "payload" not in holder and time.monotonic() < deadline:
                conn.process_data_events(time_limit=min(0.2, deadline - time.monotonic()))
            if "payload" not in holder:
                raise RpcTimeoutError(
                    f"RPC to {queue} (correlation_id={correlation_id}) timed out after {timeout}s"
                )
            return holder["payload"]
        finally:
            try:
                conn.close()
            except Exception:  # noqa: BLE001
                pass

    # ------------------------------------------------------------ RPC server
    def consume_rpc_forever(
        self,
        queue: str,
        handler: RpcHandler,
        stop: threading.Event,
    ) -> None:
        prefetch = self._intent_prefetch if queue == INTENT_REQ_QUEUE else self._snapshot_prefetch
        backoff_idx = 0
        while not stop.is_set():
            try:
                self._ensure_connected()
                ch = self._channel
                assert ch is not None
                ch.basic_qos(prefetch_count=prefetch)
                # Declare in case adapter restarted before daemon-level declare.
                ch.queue_declare(
                    queue,
                    durable=True,
                    arguments={"x-dead-letter-exchange": EVENTS_DLQ_EXCHANGE},
                )

                def on_message(_ch, method, props, body):
                    try:
                        payload = json.loads(body) if body else {}
                        req = RpcRequest(
                            payload=payload,
                            correlation_id=props.correlation_id or "",
                            reply_to=props.reply_to,
                            properties={"headers": props.headers or {}},
                        )
                        resp = handler(req)
                        if props.reply_to:
                            _ch.basic_publish(
                                exchange="",
                                routing_key=props.reply_to,
                                body=json.dumps(resp.payload).encode("utf-8"),
                                properties=pika.BasicProperties(
                                    content_type="application/json",
                                    correlation_id=props.correlation_id,
                                ),
                            )
                        _ch.basic_ack(delivery_tag=method.delivery_tag)
                    except Exception:  # noqa: BLE001
                        # Handler failure: NACK without requeue -> DLQ (REQ-023 AC-6)
                        _log.exception("RPC handler raised; routing to DLQ")
                        try:
                            _ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)
                        except Exception:  # noqa: BLE001
                            pass

                ch.basic_consume(queue=queue, on_message_callback=on_message, auto_ack=False)
                backoff_idx = 0
                while not stop.is_set() and ch.is_open:
                    self._conn.process_data_events(time_limit=0.2)
                try:
                    ch.stop_consuming()
                except Exception:  # noqa: BLE001
                    pass
            except (pika.exceptions.AMQPError, OSError) as exc:
                _log.warning("RPC consumer reconnecting: %s", exc)
                self._mark_disconnected()
                self._sleep_backoff(stop, backoff_idx)
                backoff_idx = min(backoff_idx + 1, len(_BACKOFF_SEQ) - 1)

    # ------------------------------------------------------------ Events
    def consume_events_forever(
        self,
        queue: str,
        binding_pattern: str,
        handler: EventHandler,
        stop: threading.Event,
    ) -> None:
        backoff_idx = 0
        while not stop.is_set():
            try:
                self._ensure_connected()
                self.declare_strategy_queue(queue, binding_pattern)
                ch = self._channel
                assert ch is not None

                def on_message(_ch, method, props, body):
                    try:
                        payload = json.loads(body) if body else {}
                        handler(method.routing_key, payload)
                        _ch.basic_ack(delivery_tag=method.delivery_tag)
                    except Exception:  # noqa: BLE001
                        _log.exception("event handler raised; routing to DLQ")
                        try:
                            _ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)
                        except Exception:  # noqa: BLE001
                            pass

                ch.basic_consume(queue=queue, on_message_callback=on_message, auto_ack=False)
                backoff_idx = 0
                while not stop.is_set() and ch.is_open:
                    self._conn.process_data_events(time_limit=0.2)
                try:
                    ch.stop_consuming()
                except Exception:  # noqa: BLE001
                    pass
            except (pika.exceptions.AMQPError, OSError) as exc:
                _log.warning("event consumer reconnecting: %s", exc)
                self._mark_disconnected()
                self._sleep_backoff(stop, backoff_idx)
                backoff_idx = min(backoff_idx + 1, len(_BACKOFF_SEQ) - 1)

    # ------------------------------------------------------------ helpers
    def _mark_disconnected(self) -> None:
        with self._lock:
            self._health = TransportHealth.DISCONNECTED
            self._conn = None
            self._channel = None

    def _sleep_backoff(self, stop: threading.Event, idx: int) -> None:
        delay = _BACKOFF_SEQ[min(idx, len(_BACKOFF_SEQ) - 1)]
        stop.wait(delay)