fix(connection): Give pong its own channel

The pong replies to the broker used to share the outbound channel with
all other connection traffic. When the channel is full under high load,
pongs were being discarded due to try_send on a bounded channel that was
full. Not only does this flood the log with errors, but if the broker
does not receive the pong in time, it will kill the connection and the
cycle will repeat.

This commit gives pong its own dedicated bounded(1) channel so that it
cannot be crowded out by other outbound traffic. The sink writer drains
the pong channel ahead of the main channel via select_biased!, ensuring
pong responses are flushed to the socket as soon as possible.

Fixes: #408

Author: darinspivey

src/connection.rs

@@ -93,7 +93,7 @@ impl crate::authentication::Authentication for Authentication {
 
 pub(crate) struct Receiver<S: Stream<Item = Result<Message, ConnectionError>>> {
     inbound: Pin<Box<S>>,
-    outbound: async_channel::Sender<Message>,
+    pong_tx: async_channel::Sender<Message>,
     error: SharedError,
     pending_requests: BTreeMap<RequestKey, oneshot::Sender<Message>>,
     consumers: BTreeMap<u64, mpsc::UnboundedSender<Message>>,
@@ -108,15 +108,15 @@ impl<S: Stream<Item = Result<Message, ConnectionError>>> Receiver<S> {
     #[cfg_attr(feature = "telemetry", tracing::instrument(skip_all))]
     pub fn new(
         inbound: S,
-        outbound: async_channel::Sender<Message>,
+        pong_tx: async_channel::Sender<Message>,
         error: SharedError,
         registrations: mpsc::UnboundedReceiver<Register>,
         shutdown: oneshot::Receiver<()>,
         auth_challenge: mpsc::UnboundedSender<()>,
     ) -> Receiver<S> {
         Receiver {
             inbound: Box::pin(inbound),
-            outbound,
+            pong_tx,
             error,
             pending_requests: BTreeMap::new(),
             received_messages: BTreeMap::new(),
@@ -182,8 +182,8 @@ impl<S: Stream<Item = Result<Message, ConnectionError>>> Future for Receiver<S>
                         command: BaseCommand { ping: Some(_), .. },
                         ..
                     } => {
-                        if let Err(e) = self.outbound.try_send(messages::pong()) {
-                            error!("failed to send pong: {}", e);
+                        if self.pong_tx.try_send(messages::pong()).is_err() {
+                            error!("failed to send pong: pong already pending, sink may be stalled");
                         }
                     }
                     Message {
@@ -1242,6 +1242,7 @@ impl<Exe: Executor> Connection<Exe> {
 
         let (mut sink, stream) = stream.split();
         let (tx, rx) = async_channel::bounded(outbound_channel_size);
+        let (pong_tx, pong_rx) = async_channel::bounded(1);
         let (registrations_tx, registrations_rx) = mpsc::unbounded();
         let error = SharedError::new();
         let (receiver_shutdown_tx, receiver_shutdown_rx) = oneshot::channel();
@@ -1251,7 +1252,7 @@ impl<Exe: Executor> Connection<Exe> {
             .spawn(Box::pin(
                 Receiver::new(
                     stream,
-                    tx.clone(),
+                    pong_tx,
                     error.clone(),
                     registrations_rx,
                     receiver_shutdown_rx,
@@ -1267,8 +1268,19 @@ impl<Exe: Executor> Connection<Exe> {
 
         let err = error.clone();
         let res = executor.spawn(Box::pin(async move {
-            while let Ok(msg) = rx.recv().await {
-                // println!("real sent msg: {:?}", msg);
+            loop {
+                // Drain pong responses ahead of regular outbound messages so that
+                // broker keepalive pings are answered promptly even under high load.
+                let msg = futures::select_biased! {
+                    msg = pong_rx.recv().fuse() => match msg {
+                        Ok(msg) => msg,
+                        Err(_) => break,
+                    },
+                    msg = rx.recv().fuse() => match msg {
+                        Ok(msg) => msg,
+                        Err(_) => break,
+                    },
+                };
                 if let Err(e) = sink.send(msg).await {
                     err.set(e);
                     break;
@@ -1852,7 +1864,7 @@ mod tests {
     ))]
     async fn receiver_auth_challenge_test() {
         let (message_tx, message_rx) = mpsc::unbounded();
-        let (tx, _) = async_channel::bounded(10);
+        let (pong_tx, _pong_rx) = async_channel::bounded(1);
         let (_registrations_tx, registrations_rx) = mpsc::unbounded();
         let error = SharedError::new();
         let (_receiver_shutdown_tx, receiver_shutdown_rx) = oneshot::channel();
@@ -1871,7 +1883,7 @@ mod tests {
 
         tokio::spawn(Box::pin(Receiver::new(
             message_rx,
-            tx,
+            pong_tx,
             error.clone(),
             registrations_rx,
             receiver_shutdown_rx,