394 lines
15 KiB
Python
394 lines
15 KiB
Python
"""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 (
|
|
CONTROL_REQ_QUEUE,
|
|
EVENTS_DLQ_EXCHANGE,
|
|
EVENTS_DLQ_QUEUE,
|
|
EVENTS_EXCHANGE,
|
|
INTENT_REQ_QUEUE,
|
|
PAUSE_REQ_QUEUE,
|
|
SNAPSHOT_REQ_QUEUE,
|
|
STRATEGY_QUEUE_ARGS,
|
|
)
|
|
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) -> None:
|
|
self._ensure_connected()
|
|
ch = self._channel
|
|
assert ch is not None
|
|
# 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)
|
|
# RPC queues (durable for intent + snapshot)
|
|
ch.queue_declare(
|
|
INTENT_REQ_QUEUE,
|
|
durable=True,
|
|
arguments={"x-dead-letter-exchange": EVENTS_DLQ_EXCHANGE},
|
|
)
|
|
ch.queue_declare(
|
|
SNAPSHOT_REQ_QUEUE,
|
|
durable=True,
|
|
arguments={"x-dead-letter-exchange": EVENTS_DLQ_EXCHANGE},
|
|
)
|
|
ch.queue_declare(
|
|
PAUSE_REQ_QUEUE,
|
|
durable=True,
|
|
arguments={"x-dead-letter-exchange": EVENTS_DLQ_EXCHANGE},
|
|
)
|
|
# BAT-012 / REQ-043: operator control channel (action dispatch).
|
|
ch.queue_declare(
|
|
CONTROL_REQ_QUEUE,
|
|
durable=True,
|
|
arguments={"x-dead-letter-exchange": EVENTS_DLQ_EXCHANGE},
|
|
)
|
|
|
|
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)
|