feat: add messaging transport rollout
This commit is contained in:
@@ -0,0 +1,380 @@
|
||||
"""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 (
|
||||
EVENTS_DLQ_EXCHANGE,
|
||||
EVENTS_DLQ_QUEUE,
|
||||
EVENTS_EXCHANGE,
|
||||
INTENT_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},
|
||||
)
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user