Coverage for /Users/eugene/Development/robotnikmq/robotnikmq/subscriber.py: 56%
67 statements
« prev ^ index » next coverage.py v7.3.4, created at 2023-12-26 19:13 -0500
« prev ^ index » next coverage.py v7.3.4, created at 2023-12-26 19:13 -0500
1from collections import namedtuple
2from random import uniform
3from time import sleep
4from traceback import format_exc
5from typing import Optional, Callable, List, Generator
7from pika.exceptions import AMQPError
8from pika.exchange_type import ExchangeType
9from typeguard import typechecked
11from robotnikmq.config import RobotnikConfig
12from robotnikmq.core import Robotnik, Message
13from robotnikmq.error import MalformedMessage
15OnMessageCallback = Callable[[Message], None]
17ExchangeBinding = namedtuple("ExchangeBinding", ["exchange", "binding_key"])
20class Subscriber(Robotnik):
21 TIMEOUT_MAX = 30
22 TIMEOUT_MIN = 0.5
23 TIMEOUT_STEP = 3
24 TIMEOUT_JITTER = 2
26 @typechecked
27 def __init__(self, exchange_bindings: Optional[List[ExchangeBinding]] = None,
28 config: Optional[RobotnikConfig] = None,):
29 super().__init__(config=config)
30 self.exchange_bindings = exchange_bindings or []
32 @typechecked
33 def _bind(self, exchange_binding: ExchangeBinding) -> "Subscriber":
34 self.exchange_bindings.append(exchange_binding)
35 return self
37 @typechecked
38 def bind(self, exchange: str, binding_key: str = "#") -> "Subscriber":
39 return self._bind(ExchangeBinding(exchange, binding_key))
41 @typechecked
42 def _consume(
43 self, inactivity_timeout: Optional[float]
44 ) -> Generator[Optional[Message], None, None]:
45 with self.open_channel() as channel:
46 queue_name = (
47 channel.queue_declare(queue="", exclusive=True).method.queue or ""
48 )
49 for ex_b in self.exchange_bindings:
50 channel.exchange_declare(
51 exchange=ex_b.exchange,
52 exchange_type=ExchangeType.topic,
53 auto_delete=True,
54 )
55 channel.queue_bind(
56 exchange=ex_b.exchange,
57 queue=queue_name,
58 routing_key=ex_b.binding_key,
59 )
60 try:
61 for method, ___, body in channel.consume(
62 queue=queue_name, # pragma: no cover
63 auto_ack=False,
64 inactivity_timeout=inactivity_timeout,
65 ):
66 if method and body:
67 channel.basic_ack(delivery_tag=method.delivery_tag or 0)
68 try:
69 yield Message.of_json(body.decode())
70 except MalformedMessage:
71 self.log.debug(format_exc())
72 else:
73 yield None
74 finally:
75 try:
76 channel.cancel()
77 self.close_channel(channel)
78 except AssertionError as exc:
79 self.log.warning(f"Unable to close channel: {exc}")
81 @typechecked
82 @staticmethod
83 def _jitter(step: float, jitter: float) -> float:
84 return uniform(step - jitter, step + jitter)
86 @typechecked
87 def _backoff_with_jitter(self, current_timeout: float,
88 timeout_min: Optional[float] = None,
89 timeout_step: Optional[float] = None,
90 timeout_jitter: Optional[float] = None,
91 timeout_max: Optional[float] = None) -> float:
92 timeout_min = timeout_min or self.TIMEOUT_MIN
93 timeout_step = timeout_step or self.TIMEOUT_STEP
94 timeout_jitter = timeout_jitter or self.TIMEOUT_JITTER
95 timeout_max = timeout_max or self.TIMEOUT_MAX
96 self.log.warning(f"Backing off for {current_timeout} seconds before reconnecting...")
97 sleep(current_timeout)
98 self.log.warning("Reconnecting")
99 return min(current_timeout + self._jitter(timeout_step, timeout_jitter), timeout_max)
101 @typechecked
102 def consume(
103 self, inactivity_timeout: Optional[float] = None
104 ) -> Generator[Optional[Message], None, None]:
105 timeout = Subscriber.TIMEOUT_MIN
106 while 42:
107 try:
108 for msg in self._consume(inactivity_timeout):
109 yield msg
110 timeout = Subscriber.TIMEOUT_MIN
111 except (AMQPError, OSError) as exc:
112 self.log.warning(f"Subscriber issue encountered: {exc}")
113 timeout = self._backoff_with_jitter(timeout)