rabbit_auto/
consumer.rs

1//! A stream wrapper for rabbitmq consumer. This stream never fails and will consume until stopped being used.
2use crate::exchanges::DeclareExchange;
3use anyhow::Result;
4use core::pin::Pin;
5use futures::{
6    future::Future,
7    stream::Stream,
8    task::{Context, Poll},
9};
10use lapin::{message::Delivery, Channel, Consumer};
11use std::sync::{Arc, Weak};
12use tokio::sync::mpsc::{Receiver, Sender};
13
14use super::comms::*;
15
16/// Returns a future which creates the consumer from the provided channel.
17type ConsumerCreator = Box<dyn RabbitDispatcher<Object = Consumer>>;
18
19pub type CreatorResult<T> = Pin<Box<dyn Future<Output = Result<T>> + Send>>;
20pub type Creator<T> =
21    Pin<Box<dyn Fn(Arc<Channel>, Option<Arc<DeclareExchange>>) -> CreatorResult<T> + Send + Sync>>;
22
23pub type ChannelReceiver = Receiver<Weak<Channel>>;
24
25type NextFuture = Pin<Box<dyn Future<Output = (Delivery, Data)> + Send>>;
26
27struct Data {
28    /// RabbitMQ consumer
29    consumer: Consumer,
30    /// RabbitMQ channel, this has to be keep here for having the consumer alive, otherwise there will
31    /// be non channel alive for the consumer, and the consumer would be dropped.
32    channel: Weak<Channel>,
33    /// Creator of the consumer
34    creator: ConsumerCreator,
35    /// This will be send to the Comms when the channel is not ok
36    channel_sender: ChannelSender,
37    /// Here will be a new channel delivered after requesting a new one
38    channel_receiver: Receiver<Weak<Channel>>,
39    /// The channel sender is requested over this
40    channel_requester: Arc<Sender<CommsMsg>>,
41}
42
43enum State {
44    /// Waiting to start looking for the next item
45    Idle(Data),
46    /// Looking for the next item
47    Next {
48        /// Future to get the next item
49        next: NextFuture,
50    },
51}
52
53/// Consumer wrapper handles errors in the connection. If the rabbitmq is disconnected, instead of
54/// finishing the stream, the wrapper will try to reconnect and recreate the connection and continue
55/// consuming like nothing happened. But if the connection was never established at least once, the stream
56/// end right away!
57pub struct ConsumerWrapper {
58    state: Option<State>,
59}
60
61impl ConsumerWrapper {
62    /// Create a new consumer by providing a creator function. This function might be called many times,
63    /// as often as we loose connection to the rabbitmq.
64    // pub async fn new(creator: ConsumerCreator) -> Result<Self> {
65    //     let (creator, (channel, consumer)) = Self::connect(creator).await?;
66    //     log::debug!("Consumer wrapper created");
67    //     Ok(Self { state: Some(State::Idle { consumer, channel, creator }) })
68    // }
69    pub(crate) async fn new(creator: ConsumerCreator) -> Self {
70        log::trace!("Getting channel requester");
71        let channel_requester = Comms::get_channel_comms();
72        let (channel_sender, mut channel_receiver) = Comms::create_channel_channel();
73        log::trace!("Creating the consumer using the creator");
74        let (consumer, channel) = creator
75            .start_dispatch(
76                None,
77                &channel_sender,
78                &mut channel_receiver,
79                &channel_requester,
80            )
81            .await;
82        log::trace!("Consumer wrapper created");
83        Self {
84            state: Some(State::Idle(Data {
85                consumer,
86                channel,
87                creator,
88                channel_sender,
89                channel_receiver,
90                channel_requester,
91            })),
92        }
93    }
94
95    /// Gets the next item from the consumer. If the consumer is broken, then a new consumer is automatically created
96    async fn next_item(mut data: Data) -> (Delivery, Data) {
97        loop {
98            use futures::stream::StreamExt;
99            log::trace!("Polling consumer");
100            match data.consumer.next().await {
101                Some(Ok(delivery)) => {
102                    log::trace!("Got delivery");
103                    return (delivery, data);
104                }
105                Some(Err(err)) => {
106                    log::error!("Failed to consume a message: {}", err);
107                }
108                None => {
109                    log::error!("Consumer has finished for some reason!");
110                }
111            }
112            log::warn!("Consumer is broken, waiting for a new connection");
113            let (consumer, channel) = data
114                .creator
115                .start_dispatch(
116                    Some(data.channel.clone()),
117                    &data.channel_sender,
118                    &mut data.channel_receiver,
119                    &data.channel_requester,
120                )
121                .await;
122            data.consumer = consumer;
123            data.channel = channel;
124        }
125    }
126}
127
128impl Stream for ConsumerWrapper {
129    type Item = Delivery;
130    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
131        log::trace!("Poll next");
132        let this = Pin::into_inner(self);
133
134        loop {
135            match this.state.take() {
136                Some(State::Idle(data)) => {
137                    this.state = Some(State::Next {
138                        next: Box::pin(Self::next_item(data)),
139                    });
140                }
141                Some(State::Next { mut next }) => {
142                    let action = next.as_mut();
143                    return match Future::poll(action, cx) {
144                        Poll::Pending => {
145                            this.state = Some(State::Next { next });
146                            log::trace!("Pending");
147                            Poll::Pending
148                        }
149                        Poll::Ready((delivery, data)) => {
150                            this.state = Some(State::Idle(data));
151                            log::trace!("Ready");
152                            Poll::Ready(Some(delivery))
153                        }
154                    };
155                }
156                None => unreachable!(),
157            }
158        }
159    }
160}
rabbit_auto/consumer.rs

rabbit_auto/
consumer.rs
