fix: reset send streams the futures get aborted and improve error reporting on the receive side of the quic transport (#2566)

rumenov · web-flow · commit e500f911f8b0 · 2024-11-13T10:47:33.000Z
diff --git a/rs/p2p/quic_transport/src/connection_handle.rs b/rs/p2p/quic_transport/src/connection_handle.rs
@@ -2,47 +2,22 @@
 //!
 use std::sync::atomic::{AtomicU64, Ordering};
 
-use anyhow::Context;
 use bytes::Bytes;
 use http::{Method, Request, Response, Version};
 use ic_protobuf::transport::v1 as pb;
 use prost::Message;
-use quinn::{Connection, ReadToEndError, SendStream, StoppedError, VarInt};
+use quinn::Connection;
 
 use crate::{
     metrics::{
-        observe_conn_error, observe_read_error, observe_write_error, QuicTransportMetrics,
-        INFALIBBLE,
+        observe_conn_error, observe_read_to_end_error, observe_stopped_error, observe_write_error,
+        QuicTransportMetrics, INFALIBBLE,
     },
-    ConnId, MessagePriority, MAX_MESSAGE_SIZE_BYTES,
+    ConnId, MessagePriority, ResetStreamOnDrop, MAX_MESSAGE_SIZE_BYTES,
 };
 
-/// QUIC error code for stream cancellation. See
-/// https://datatracker.ietf.org/doc/html/draft-ietf-quic-transport-03#section-12.3.
-const QUIC_STREAM_CANCELLED: VarInt = VarInt::from_u32(6);
-
 static CONN_ID_SEQ: AtomicU64 = AtomicU64::new(1);
 
-/// Drop guard to send a [`SendStream::reset`] frame on drop. QUINN sends a [`SendStream::finish`] frame by default when dropping a [`SendStream`],
-/// which can lead to the peer receiving the stream thinking a complete message was sent. This guard is used to send a reset frame instead, to signal
-/// that the transmission of the message was cancelled.
-struct SendStreamDropGuard {
-    send_stream: SendStream,
-}
-
-impl SendStreamDropGuard {
-    fn new(send_stream: SendStream) -> Self {
-        Self { send_stream }
-    }
-}
-
-impl Drop for SendStreamDropGuard {
-    fn drop(&mut self) {
-        // fails silently if the stream is already closed.
-        let _ = self.send_stream.reset(QUIC_STREAM_CANCELLED);
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct ConnectionHandle {
     conn: Connection,
@@ -77,7 +52,7 @@ impl ConnectionHandle {
     /// is delegated solely to the transport layer. This differs from typical client-server architectures,
     /// where connections can be managed directly by the caller.
     ///
-    /// Note: This method provides the same cancellation safety guarantees as the `quinn::Connection` methods.
+    /// Note: The method is cancel-safe.
     pub async fn rpc(&self, request: Request<Bytes>) -> Result<Response<Bytes>, anyhow::Error> {
         let _timer = self
             .metrics
@@ -98,7 +73,7 @@ impl ConnectionHandle {
             observe_conn_error(err, "open_bi", &self.metrics.connection_handle_errors_total);
         })?;
 
-        let mut send_stream_guard = SendStreamDropGuard::new(send_stream);
+        let mut send_stream_guard = ResetStreamOnDrop::new(send_stream);
         let send_stream = &mut send_stream_guard.send_stream;
 
         let priority = request
@@ -130,36 +105,18 @@ impl ConnectionHandle {
                 .inc();
         })?;
 
-        send_stream.stopped().await.inspect_err(|err| match err {
-            StoppedError::ConnectionLost(conn_err) => {
-                observe_conn_error(
-                    conn_err,
-                    "stopped",
-                    &self.metrics.connection_handle_errors_total,
-                );
-            }
-            StoppedError::ZeroRttRejected => {
-                self.metrics
-                    .connection_handle_errors_total
-                    .with_label_values(&["stopped", INFALIBBLE])
-                    .inc();
-            }
+        send_stream.stopped().await.inspect_err(|err| {
+            observe_stopped_error(err, "stopped", &self.metrics.connection_handle_errors_total)
         })?;
-
         let response_bytes = recv_stream
             .read_to_end(MAX_MESSAGE_SIZE_BYTES)
             .await
-            .inspect_err(|err| match err {
-                ReadToEndError::TooLong => self
-                    .metrics
-                    .connection_handle_errors_total
-                    .with_label_values(&["read_to_end", INFALIBBLE])
-                    .inc(),
-                ReadToEndError::Read(read_err) => observe_read_error(
-                    read_err,
+            .inspect_err(|err| {
+                observe_read_to_end_error(
+                    err,
                     "read_to_end",
                     &self.metrics.connection_handle_errors_total,
-                ),
+                )
             })?;
 
         let response = to_response(response_bytes)?;
@@ -169,14 +126,10 @@ impl ConnectionHandle {
     }
 }
 
+// The function returns infallible error.
 fn to_response(response_bytes: Vec<u8>) -> Result<Response<Bytes>, anyhow::Error> {
-    let response_proto = pb::HttpResponse::decode(response_bytes.as_slice())
-        .with_context(|| "Failed to decode response header.")?;
-
-    let status: u16 = response_proto
-        .status_code
-        .try_into()
-        .with_context(|| "Failed to decode status code.")?;
+    let response_proto = pb::HttpResponse::decode(response_bytes.as_slice())?;
+    let status: u16 = response_proto.status_code.try_into()?;
 
     let mut response = Response::builder().status(status).version(Version::HTTP_3);
     for h in response_proto.headers {
@@ -185,9 +138,7 @@ fn to_response(response_bytes: Vec<u8>) -> Result<Response<Bytes>, anyhow::Error
     }
     // This consumes the body without requiring allocation or cloning the whole content.
     let body_bytes = Bytes::from(response_proto.body);
-    response
-        .body(body_bytes)
-        .with_context(|| "Failed to build response.")
+    Ok(response.body(body_bytes)?)
 }
 
 fn into_request_bytes(request: Request<Bytes>) -> Vec<u8> {
@@ -242,11 +193,11 @@ mod tests {
     use tokio::sync::Barrier;
     use turmoil::Builder;
 
-    use crate::connection_handle::SendStreamDropGuard;
+    use crate::ResetStreamOnDrop;
 
     const MAX_READ_SIZE: usize = 10_000;
 
-    /// Test that [`SendStreamDropGuard`] sends a reset frame on drop. Also tests that
+    /// Test that [`ResetStreamOnDrop`] sends a reset frame on drop. Also tests that
     /// the receiver will receive the message if the stream is finished and stopped,
     /// before dropping the guard.
     #[rstest]
@@ -341,7 +292,7 @@ mod tests {
                 .unwrap();
 
             let (send_stream, _recv_stream) = connection.open_bi().await.unwrap();
-            let mut drop_guard = SendStreamDropGuard::new(send_stream);
+            let mut drop_guard = ResetStreamOnDrop::new(send_stream);
             let send_stream = &mut drop_guard.send_stream;
             send_stream
                 .write_chunk(Bytes::from(&b"hello wo"[..]))
diff --git a/rs/p2p/quic_transport/src/lib.rs b/rs/p2p/quic_transport/src/lib.rs
@@ -49,7 +49,7 @@ use ic_interfaces_registry::RegistryClient;
 use ic_logger::{info, ReplicaLogger};
 use ic_metrics::MetricsRegistry;
 use phantom_newtype::AmountOf;
-use quinn::AsyncUdpSocket;
+use quinn::{AsyncUdpSocket, SendStream, VarInt};
 use tokio::sync::watch;
 use tokio::task::{JoinError, JoinHandle};
 use tokio_util::{sync::CancellationToken, task::task_tracker::TaskTracker};
@@ -180,6 +180,30 @@ impl QuicTransport {
     }
 }
 
+/// QUIC error code for stream cancellation. See
+/// https://datatracker.ietf.org/doc/html/draft-ietf-quic-transport-03#section-12.3.
+const QUIC_STREAM_CANCELLED: VarInt = VarInt::from_u32(6);
+
+/// Drop guard to send a [`SendStream::reset`] frame on drop. QUINN sends a [`SendStream::finish`] frame by default when dropping a [`SendStream`],
+/// which can lead to the peer receiving the stream thinking a complete message was sent. This guard is used to send a reset frame instead, to signal
+/// that the transmission of the message was cancelled.
+struct ResetStreamOnDrop {
+    send_stream: SendStream,
+}
+
+impl ResetStreamOnDrop {
+    fn new(send_stream: SendStream) -> Self {
+        Self { send_stream }
+    }
+}
+
+impl Drop for ResetStreamOnDrop {
+    fn drop(&mut self) {
+        // fails silently if the stream is already closed.
+        let _ = self.send_stream.reset(QUIC_STREAM_CANCELLED);
+    }
+}
+
 #[async_trait]
 impl Transport for QuicTransport {
     #[instrument(skip(self, request))]
diff --git a/rs/p2p/quic_transport/src/metrics.rs b/rs/p2p/quic_transport/src/metrics.rs
@@ -3,7 +3,7 @@ use ic_metrics::{
     buckets::decimal_buckets, tokio_metrics_collector::TokioTaskMetricsCollector, MetricsRegistry,
 };
 use prometheus::{GaugeVec, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec};
-use quinn::{Connection, ConnectionError, ReadError, WriteError};
+use quinn::{Connection, ConnectionError, ReadError, ReadToEndError, StoppedError, WriteError};
 use tokio_metrics::TaskMonitor;
 
 const CONNECTION_RESULT_LABEL: &str = "status";
@@ -15,13 +15,8 @@ const ERROR_TYPE_LABEL: &str = "error";
 const REQUEST_TYPE_LABEL: &str = "request";
 pub(crate) const CONNECTION_RESULT_SUCCESS_LABEL: &str = "success";
 pub(crate) const CONNECTION_RESULT_FAILED_LABEL: &str = "failed";
-pub(crate) const ERROR_TYPE_ACCEPT: &str = "accept";
 pub(crate) const ERROR_TYPE_APP: &str = "app";
-pub(crate) const ERROR_TYPE_FINISH: &str = "finish";
-pub(crate) const ERROR_TYPE_STOPPED: &str = "stopped";
-pub(crate) const ERROR_TYPE_READ: &str = "read";
 pub(crate) const INFALIBBLE: &str = "infallible";
-pub(crate) const ERROR_TYPE_WRITE: &str = "write";
 const ERROR_CLOSED_STREAM: &str = "closed_stream";
 const ERROR_RESET_STREAM: &str = "reset_stream";
 const ERROR_STOPPED_STREAM: &str = "stopped_stream";
@@ -247,3 +242,17 @@ pub fn observe_read_error(err: &ReadError, op: &str, counter: &IntCounterVec) {
         }
     }
 }
+
+pub fn observe_stopped_error(err: &StoppedError, op: &str, counter: &IntCounterVec) {
+    match err {
+        StoppedError::ConnectionLost(conn_err) => observe_conn_error(conn_err, op, counter),
+        StoppedError::ZeroRttRejected => counter.with_label_values(&[op, INFALIBBLE]).inc(),
+    }
+}
+
+pub fn observe_read_to_end_error(err: &ReadToEndError, op: &str, counter: &IntCounterVec) {
+    match err {
+        ReadToEndError::TooLong => counter.with_label_values(&[op, INFALIBBLE]).inc(),
+        ReadToEndError::Read(read_err) => observe_read_error(read_err, op, counter),
+    }
+}
diff --git a/rs/p2p/quic_transport/src/request_handler.rs b/rs/p2p/quic_transport/src/request_handler.rs
@@ -27,10 +27,10 @@ use tracing::instrument;
 use crate::{
     connection_handle::ConnectionHandle,
     metrics::{
-        QuicTransportMetrics, ERROR_TYPE_ACCEPT, ERROR_TYPE_APP, ERROR_TYPE_FINISH,
-        ERROR_TYPE_READ, ERROR_TYPE_STOPPED, ERROR_TYPE_WRITE, STREAM_TYPE_BIDI,
+        observe_conn_error, observe_read_to_end_error, observe_stopped_error, observe_write_error,
+        QuicTransportMetrics, ERROR_TYPE_APP, INFALIBBLE, STREAM_TYPE_BIDI,
     },
-    ConnId, MAX_MESSAGE_SIZE_BYTES,
+    ConnId, ResetStreamOnDrop, MAX_MESSAGE_SIZE_BYTES,
 };
 
 const QUIC_METRIC_SCRAPE_INTERVAL: Duration = Duration::from_secs(5);
@@ -70,15 +70,9 @@ pub async fn run_stream_acceptor(
                             )
                         );
                     }
-                    Err(e) => {
-                        info!(log, "Error accepting bi stream {}", e.to_string());
-                        metrics
-                            .request_handle_errors_total
-                            .with_label_values(&[
-                                STREAM_TYPE_BIDI,
-                                ERROR_TYPE_ACCEPT,
-                            ])
-                            .inc();
+                    Err(err) => {
+                        info!(log, "Error accepting bi stream {}", err.to_string());
+                        observe_conn_error(&err, "accept_bi", &metrics.request_handle_errors_total);
                         break;
                     }
                 }
@@ -105,26 +99,30 @@ pub async fn run_stream_acceptor(
     inflight_requests.shutdown().await;
 }
 
-#[instrument(skip(metrics, router, bi_tx, bi_rx))]
+#[instrument(skip(metrics, router, send_stream, recv_stream))]
+/// Note: The method is cancel-safe.
 async fn handle_bi_stream(
     peer_id: NodeId,
     conn_id: ConnId,
     metrics: QuicTransportMetrics,
     router: Router,
-    mut bi_tx: SendStream,
-    bi_rx: RecvStream,
+    send_stream: SendStream,
+    mut recv_stream: RecvStream,
 ) -> Result<(), anyhow::Error> {
-    let mut request = read_request(bi_rx).await.inspect_err(|_| {
-        metrics
-            .request_handle_errors_total
-            .with_label_values(&[STREAM_TYPE_BIDI, ERROR_TYPE_READ])
-            .inc();
-    })?;
+    let mut send_stream_guard = ResetStreamOnDrop::new(send_stream);
+    let send_stream = &mut send_stream_guard.send_stream;
+    let request_bytes = recv_stream
+        .read_to_end(MAX_MESSAGE_SIZE_BYTES)
+        .await
+        .inspect_err(|err| {
+            observe_read_to_end_error(err, "read_to_end", &metrics.request_handle_errors_total)
+        })?;
+    let mut request = to_request(request_bytes)?;
     request.extensions_mut().insert::<NodeId>(peer_id);
     request.extensions_mut().insert::<ConnId>(conn_id);
 
     let svc = router.oneshot(request);
-    let stopped = bi_tx.stopped();
+    let stopped = send_stream.stopped();
     let response = tokio::select! {
         response = svc => response.expect("Infallible"),
         stopped_res = stopped => {
@@ -143,40 +141,29 @@ async fn handle_bi_stream(
     // We can ignore the errors because if both peers follow the protocol an errors will only occur
     // if the other peer has closed the connection. In this case `accept_bi` in the peer event
     // loop will close this connection.
-    write_response(&mut bi_tx, response)
+    let response_bytes = to_response_bytes(response).await?;
+    send_stream
+        .write_all(&response_bytes)
         .await
-        .inspect_err(|_| {
-            metrics
-                .request_handle_errors_total
-                .with_label_values(&[STREAM_TYPE_BIDI, ERROR_TYPE_WRITE])
-                .inc();
+        .inspect_err(|err| {
+            observe_write_error(err, "write_all", &metrics.request_handle_errors_total);
         })?;
-    bi_tx.finish().inspect_err(|_| {
+    send_stream.finish().inspect_err(|_| {
         metrics
             .request_handle_errors_total
-            .with_label_values(&[STREAM_TYPE_BIDI, ERROR_TYPE_FINISH])
+            .with_label_values(&["finish", INFALIBBLE])
             .inc();
     })?;
-    bi_tx.stopped().await.inspect_err(|_| {
-        metrics
-            .request_handle_errors_total
-            .with_label_values(&[STREAM_TYPE_BIDI, ERROR_TYPE_STOPPED])
-            .inc();
+    send_stream.stopped().await.inspect_err(|err| {
+        observe_stopped_error(err, "stopped", &metrics.request_handle_errors_total);
     })?;
     Ok(())
 }
 
-async fn read_request(mut recv_stream: RecvStream) -> Result<Request<Body>, anyhow::Error> {
-    let raw_msg = recv_stream
-        .read_to_end(MAX_MESSAGE_SIZE_BYTES)
-        .await
-        .with_context(|| "Failed to read request from the stream.")?;
-
-    let request_proto = pb::HttpRequest::decode(raw_msg.as_slice())
-        .with_context(|| "Failed to decode http request.")?;
-
-    let pb_http_method = pb::HttpMethod::try_from(request_proto.method)
-        .with_context(|| "Failed to decode http method.")?;
+// The function returns infallible error.
+fn to_request(request_bytes: Vec<u8>) -> Result<Request<Body>, anyhow::Error> {
+    let request_proto = pb::HttpRequest::decode(request_bytes.as_slice())?;
+    let pb_http_method = pb::HttpMethod::try_from(request_proto.method)?;
     let http_method = match pb_http_method {
         pb::HttpMethod::Get => Some(Method::GET),
         pb::HttpMethod::Post => Some(Method::POST),
@@ -207,10 +194,7 @@ async fn read_request(mut recv_stream: RecvStream) -> Result<Request<Body>, anyh
         .with_context(|| "Failed to build request.")
 }
 
-async fn write_response(
-    send_stream: &mut SendStream,
-    response: Response<Body>,
-) -> Result<(), anyhow::Error> {
+async fn to_response_bytes(response: Response<Body>) -> Result<Vec<u8>, anyhow::Error> {
     let (parts, body) = response.into_parts();
     // Check for axum error in body
     // TODO: Think about this. What is the error that can happen here?
@@ -231,10 +215,5 @@ async fn write_response(
             .collect(),
         body: body.into(),
     };
-
-    let response_bytes = response_proto.encode_to_vec();
-    send_stream
-        .write_all(&response_bytes)
-        .await
-        .with_context(|| "Failed to write response to stream.")
+    Ok(response_proto.encode_to_vec())
 }
