rabbit_auto/
comms.rs

1//! Helper to create a connection and keep it connected.
2
3use async_rs::traits::{Executor, Reactor};
4use async_rs::Runtime;
5use futures::lock::Mutex;
6use futures::Future;
7use lapin::types::{AMQPValue, ShortUInt};
8use lapin::{Channel, Connection, ConnectionProperties, ConnectionState};
9use std::collections::HashMap;
10use std::pin::Pin;
11use std::sync::{Arc, Weak};
12// use crate::comms::comms_data::Data;
13use crate::config::Config;
14use crate::exchanges::DeclareExchange;
15
16pub const MAX_CHANNELS: usize = 16;
17use tokio::sync::mpsc::{channel, Receiver, Sender};
18
19/// Result of the creator function
20pub type CreatorResult<T> = Pin<Box<dyn Future<Output = anyhow::Result<T>> + Send>>;
21/// Creator function
22pub type Creator<T> =
23    Pin<Box<dyn Fn(Arc<Channel>, Option<Arc<DeclareExchange>>) -> CreatorResult<T> + Send + Sync>>;
24
25pub type ChannelSender = Arc<Sender<Weak<Channel>>>;
26
27// mod comms_data;
28
29/// RabbitMQ connection singleton
30pub struct Comms;
31/// Internal size of the channel buffer
32const COMMS_MSG_SIZE: usize = 32;
33
34pub(crate) enum CommsMsg {
35    DeclareExchange {
36        exchange: String,
37        declare: DeclareExchange,
38    },
39    RequestChannel(ChannelSender),
40}
41
42impl Comms {
43    /// Declare exchange
44    pub async fn declare_exchange(exchange: String, declare: DeclareExchange) {
45        let sender = Self::channel_comms().0;
46        log::debug!("Sending {exchange} declaration to the rabbit loop");
47        if let Err(err) = sender
48            .send(CommsMsg::DeclareExchange { exchange, declare })
49            .await
50        {
51            log::error!("Declare exchange internal comms is broken: {err}");
52            std::process::exit(-1);
53        }
54    }
55
56    /// Gets the singleton global data
57    /// The Receiver is taken away by the configure method when is started.
58    fn channel_comms() -> (
59        &'static Arc<Sender<CommsMsg>>,
60        &'static Arc<Mutex<Option<Receiver<CommsMsg>>>>,
61    ) {
62        lazy_static::lazy_static! {
63            static ref COMMS: (Arc<Sender<CommsMsg>>, Arc<Mutex<Option<Receiver<CommsMsg>>>>) = {
64                let (tx, rx) = channel(COMMS_MSG_SIZE);
65                (Arc::new(tx), Arc::new(Mutex::new(Some(rx))))
66            };
67        }
68        (&COMMS.0, &COMMS.1)
69    }
70
71    /// Gets the singleton channel sender. This is already cloned and ready to use
72    pub(crate) fn get_channel_comms() -> Arc<Sender<CommsMsg>> {
73        Self::channel_comms().0.clone()
74    }
75
76    /// Creates a new channel to transfer Weak<Channel> object
77    pub(crate) fn create_channel_channel() -> (Arc<Sender<Weak<Channel>>>, Receiver<Weak<Channel>>)
78    {
79        let (tx, rx) = channel(2);
80        (Arc::new(tx), rx)
81    }
82
83    /// The connection needs to be configured before the first using
84    /// This is no longer async method and this will start the rabbit auto connection manager.
85    pub fn configure(config: Config) -> anyhow::Result<()> {
86        let rt = Runtime::tokio()?;
87        let spawner = rt.clone();
88        spawner.spawn(Box::pin(async move {
89            // let (mut channel_requester, mut declare_exchange_requester)  = {
90            let mut msg_receiver = {
91                log::debug!("Taking the channel requester from the global area");
92                let singleton = Self::channel_comms().1.clone();
93                let mut channel_receiver = singleton.lock().await;
94                if let Some(cr) = channel_receiver.take() {
95                    cr
96                } else {
97                    log::error!("Fatal error. No channel receiver found. Is the Comms::configure called the second time?");
98                    std::process::exit(1);
99                }
100            };
101            log::debug!("Channel requester taken, we can start the rabbit loop");
102            let mut connection: Option<Connection> = None;
103            let mut connection_attempt = 0;
104            let mut declare_exchanges = HashMap::new();
105
106            let mut channels = channels_ring::Channels::new();
107            let mut is_connecting = false;
108            loop {
109                let mut channel_sender = Vec::with_capacity(COMMS_MSG_SIZE);
110                let mut msgs = Vec::with_capacity(COMMS_MSG_SIZE);
111                let _ = msg_receiver.recv_many(&mut msgs, COMMS_MSG_SIZE).await;
112                for msg in msgs.into_iter() {
113                    match msg {
114                        CommsMsg::DeclareExchange { exchange, declare } => {
115                            log::trace!("Received a new exchange declaration for {exchange}");
116                            declare_exchanges.insert(exchange, declare);
117                        }
118                        CommsMsg::RequestChannel(cs) => {
119                            log::trace!("There is a request for a channel");
120                            channel_sender.push(cs);
121                        }
122                    }
123                }
124                if channel_sender.is_empty() {
125                    log::debug!("No channel requesters yet, going back waiting");
126                    continue;
127                }
128                // try to establish a connection until we have one and can send a channel back to the requester
129                loop {
130                    // connect
131                    if let Some(conn) = connection.take() {
132                        // validate the connection
133                        log::debug!("Validating an existing connection");
134                        'connection: loop {
135                            match conn.status().state() {
136                                ConnectionState::Initial | ConnectionState::Connecting | ConnectionState::Reconnecting => {
137                                    log::debug!("Connecting is in progress");
138                                    // lets check again in a while
139                                    rt.sleep(std::time::Duration::from_millis(100)).await;
140                                }
141                                ConnectionState::Connected => {
142                                    // the connection is valid
143                                    if is_connecting {
144                                        log::info!("Connection to RabbitMQ established");
145                                        is_connecting = false;
146                                    } else {
147                                        log::debug!("RabbitMQ is already connected");
148                                    }
149                                    if !declare_exchanges.is_empty() {
150                                        log::debug!("There are {} exchanges to declare", declare_exchanges.len());
151                                        match channels.create_channel(&conn).await {
152                                            Ok(channel) => {
153                                                // run all exchange declarations
154                                                while !declare_exchanges.is_empty() {
155                                                    let mut done = None;
156                                                    for (exchange, declare) in declare_exchanges.iter() {
157                                                        log::debug!("Running exchange declaration for '{exchange}'");
158                                                        if let Err(err) = declare(channel.clone()).await {
159                                                            log::error!("Failed to declare an exchange: {err}");
160                                                            continue 'connection;
161                                                        }
162                                                        log::debug!("Exchange '{exchange}' declared");
163                                                        done = Some(exchange.clone());
164                                                        break;
165                                                    }
166                                                    if let Some(exchange) = done {
167                                                        declare_exchanges.remove(&exchange);
168                                                    }
169                                                }
170                                            }
171                                            Err(err) => {
172                                                // creating the channel failed, we need a new connection
173                                                log::error!("Connection to rabbit failed, reconnecting: {err}");
174                                            }
175                                        }
176                                    }
177                                    connection = Some(conn);
178                                    log::debug!("Connection initialised");
179                                    break;
180                                }
181                                ConnectionState::Closing |
182                                ConnectionState::Closed |
183                                ConnectionState::Error => {
184                                    log::warn!("Failed to connect to the rabbit server, will try to reconnect");
185                                    log::trace!("Sleeping for {} seconds", config.reconnect_delay.as_secs_f64());
186                                    rt.sleep(config.reconnect_delay).await;
187                                    break;
188                                }
189                            }
190                        }
191                    } else {
192                        log::debug!("There is no connection yet");
193                    }
194                    // check if the connection exists or we need a new one
195                    if let Some(conn) = connection.take() {
196                        // so the connection should be valid, we can finally get the channel and send it
197                        log::debug!("Connection exist, getting channels");
198                        match channels.create_channels(channel_sender.len(), &conn).await {
199                            Ok(channels) => {
200                                log::debug!("Channels are created, sending them to the individual requesters");
201                                for (cs, channel) in channel_sender.drain(..).zip(channels) {
202                                    let channel = Arc::downgrade(&channel);
203                                    if let Err(err) = cs.send(channel).await {
204                                        log::error!("Fatal error. Failed to send the channel back to the requester: {err}");
205                                        std::process::exit(1);
206                                    }
207                                }
208                                connection = Some(conn);
209                                break;
210                            }
211                            Err(err) => {
212                                // creating the channel failed, we need a new connection
213                                log::error!("Connection to rabbit failed, reconnecting: {err}");
214                            }
215                        }
216                    } else {
217                        // we need a new one
218                        log::debug!("Closing any existing channels if any");
219                        channels.try_close().await; // close any existing channels
220
221                        log::debug!("No connection, attempting to connect [{connection_attempt}]",);
222                        let mut properties = ConnectionProperties::default()
223                            .with_connection_name(config.name.clone().into());
224                        properties
225                            .client_properties
226                            .insert("channel_max".into(), AMQPValue::ShortUInt(MAX_CHANNELS as ShortUInt));
227                        let addr = config.address[connection_attempt % config.address.len()].as_ref();
228                        connection_attempt += 1;
229                        let con = Connection::connect(addr, properties).await;
230                        match con {
231                            Ok(con) => {
232                                // we have a new connection, we will wait until fully connected
233                                log::info!("Connecting to RabbitMQ");
234                                connection = Some(con);
235                            }
236                            Err(err) => {
237                                log::error!("Failed to connect: {err}");
238                                log::debug!("Sleeping for {} seconds", config.reconnect_delay.as_secs_f64());
239                                rt.sleep(config.reconnect_delay).await;
240                            }
241                        }
242                    }
243                }
244            }
245        }));
246        Ok(())
247    }
248}
249
250/// This will dispatch an operation on a channel. If the channel is not
251/// valid, it will request a new channel and when all is ok, it will dispatch the operation.
252#[async_trait::async_trait]
253pub(crate) trait RabbitDispatcher: Sync + Send {
254    type Object;
255    /// This can only fail if the rabbit connection is broken, otherwise it will get into a dead lock
256    async fn dispatcher(&self, channel: &Channel) -> anyhow::Result<Self::Object>;
257    async fn channel_setup(&self, _channel: &Channel) -> anyhow::Result<()> {
258        Ok(())
259    }
260    async fn start_dispatch(
261        &self,
262        mut channel: Option<Weak<Channel>>,
263        channel_sender: &ChannelSender,
264        channel_receiver: &mut Receiver<Weak<Channel>>,
265        channel_requester: &Arc<Sender<CommsMsg>>,
266    ) -> (Self::Object, Weak<Channel>) {
267        let mut setup_channel = false;
268        loop {
269            // ensure we have some channel
270            let channel = if let Some(channel) = channel.take() {
271                channel
272            } else {
273                // if the channel is None, we need to request a new channel from the Comms
274                if let Err(err) = channel_requester
275                    .send(CommsMsg::RequestChannel(channel_sender.clone()))
276                    .await
277                {
278                    log::error!("Fatal internal error. Failed to request a new channel: {err}");
279                    std::process::exit(1);
280                }
281                // now eventually the Channel should arrive over the channel_receiver
282                if let Some(channel) = channel_receiver.recv().await {
283                    log::trace!("Channel received");
284                    setup_channel = true;
285                    channel
286                } else {
287                    log::error!("Fatal internal error. Failed to receive a new channel");
288                    std::process::exit(1);
289                }
290            };
291            // ensure the channel is alive, if not, get a new one
292            let channel = match channel.upgrade() {
293                None => {
294                    log::trace!("No channel, we are going to wait for a new one");
295                    continue;
296                }
297                Some(channel) => channel,
298            };
299
300            if setup_channel {
301                log::trace!("Setting up the channel");
302                if let Err(err) = self.channel_setup(&channel).await {
303                    log::error!("Failed to setup the channel: {err}");
304                    continue;
305                }
306                log::trace!("Channel setup done");
307            }
308
309            log::trace!("Dispatching the operation");
310            match self.dispatcher(channel.as_ref()).await {
311                Ok(object) => {
312                    log::trace!("Operation is dispatched");
313                    return (object, Arc::downgrade(&channel));
314                }
315                Err(err) => {
316                    log::error!("Failed to failed to dispatch the operation: {}", err);
317                    continue;
318                }
319            }
320        }
321    }
322}
323
324mod channels_ring;
rabbit_auto/comms.rs

rabbit_auto/
comms.rs
