From bfbc737b6c14bc9acf9a870194365a3a68764ad8 Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Tue, 1 Dec 2020 13:24:13 -0800
Subject: [PATCH] network: don't cancel heartbeat requests

The cancellation implementation changes made to the connection state machine
mean that if a response oneshot is dropped, the connection will avoid
cancelling the request.  So the heartbeat task does have to wait on the response.
---
 zebra-network/src/peer/handshake.rs | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/zebra-network/src/peer/handshake.rs b/zebra-network/src/peer/handshake.rs
index ed80a610..fe2f2b76 100644
--- a/zebra-network/src/peer/handshake.rs
+++ b/zebra-network/src/peer/handshake.rs
@@ -460,21 +460,35 @@ where
                         let shutdown_rx_ref = Pin::new(&mut shutdown_rx);
                         match future::select(interval_stream.next(), shutdown_rx_ref).await {
                             Either::Left(_) => {
-                                // We don't wait on a response because heartbeats are checked
-                                // internally to the connection logic, we just need a separate
-                                // task (this one) to generate them.
-                                let (request_tx, _) = oneshot::channel();
+                                let (tx, rx) = oneshot::channel();
+                                let request = Request::Ping(Nonce::default());
+                                tracing::trace!(?request, "queueing heartbeat request");
                                 if server_tx
                                     .send(ClientRequest {
-                                        request: Request::Ping(Nonce::default()),
-                                        tx: request_tx,
+                                        request,
+                                        tx,
                                         span: tracing::Span::current(),
                                     })
                                     .await
                                     .is_err()
                                 {
+                                    tracing::trace!(
+                                        "error sending heartbeat request, shutting down"
+                                    );
                                     return;
                                 }
+                                // Heartbeats are checked internally to the
+                                // connection logic, but we need to wait on the
+                                // response to avoid canceling the request.
+                                match rx.await {
+                                    Ok(_) => tracing::trace!("got heartbeat response"),
+                                    Err(_) => {
+                                        tracing::trace!(
+                                            "error awaiting heartbeat response, shutting down"
+                                        );
+                                        return;
+                                    }
+                                }
                             }
                             Either::Right(_) => return, // got shutdown signal
                         }