chore: add more documentation to the quic transport code and make some errors infallible (#2515)

Author: rumenov

rs/p2p/quic_transport/src/connection_handle.rs

@@ -1,6 +1,6 @@
 //! The module implements the RPC abstraction over an established QUIC connection.
 //!
-use anyhow::{anyhow, Context};
+use anyhow::Context;
 use bytes::Bytes;
 use http::{Method, Request, Response, Version};
 use ic_protobuf::transport::v1 as pb;
@@ -42,39 +42,34 @@ impl Drop for SendStreamDropGuard {
 #[derive(Clone, Debug)]
 pub struct ConnectionHandle {
     pub connection: Connection,
-    metrics: QuicTransportMetrics,
-    conn_id: ConnId,
+    pub metrics: QuicTransportMetrics,
+    pub conn_id: ConnId,
 }
 
 impl ConnectionHandle {
-    pub fn new(connection: Connection, metrics: QuicTransportMetrics, conn_id: ConnId) -> Self {
-        Self {
-            connection,
-            metrics,
-            conn_id,
-        }
-    }
-
-    pub fn conn_id(&self) -> ConnId {
-        self.conn_id
-    }
-
-    /// Performs an RPC operation on the already established connection.
+    /// Executes an RPC operation over an already-established connection.
+    ///
+    /// This method leverages the QUIC transport layer, which continuously monitors the connection’s health
+    /// and automatically attempts reconnection as necessary. As a result, any errors returned by this method
+    /// should be considered transient (retryable).
+    ///
+    /// In this P2P architecture, where there is a designated dialer and receiver, connection management
+    /// is delegated solely to the transport layer. This differs from typical client-server architectures,
+    /// where connections can be managed directly by the caller.
     ///
-    /// Since the QUIC transport layer continuously monitors connection health and automatically reconnects as needed,
-    /// all errors returned by this method can be safely treated as transient (retryable).
-    /// Therefore, we return an anyhow::Error for error handling.
+    /// Note: This method provides the same cancellation safety guarantees as the `quinn::Connection` methods.
     pub async fn rpc(&self, request: Request<Bytes>) -> Result<Response<Bytes>, anyhow::Error> {
         let _timer = self
             .metrics
             .connection_handle_duration_seconds
             .with_label_values(&[request.uri().path()])
             .start_timer();
-        self.metrics
+
+        let bytes_sent_counter = self
+            .metrics
             .connection_handle_bytes_sent_total
-            .with_label_values(&[request.uri().path()])
-            .inc_by(request.body().len() as u64);
-        let in_counter = self
+            .with_label_values(&[request.uri().path()]);
+        let bytes_received_counter = self
             .metrics
             .connection_handle_bytes_received_total
             .with_label_values(&[request.uri().path()]);
@@ -95,6 +90,7 @@ impl ConnectionHandle {
             .unwrap_or_default();
         let _ = send_stream.set_priority(priority.into());
 
+        bytes_sent_counter.inc_by(request.body().len() as u64);
         write_request(send_stream, request).await.inspect_err(|_| {
             self.metrics
                 .connection_handle_errors_total
@@ -123,7 +119,7 @@ impl ConnectionHandle {
                 .inc();
         })?;
 
-        in_counter.inc_by(response.body().len() as u64);
+        bytes_received_counter.inc_by(response.body().len() as u64);
         Ok(response)
     }
 }
@@ -137,16 +133,10 @@ async fn read_response(mut recv_stream: RecvStream) -> Result<Response<Bytes>, a
     let response_proto = pb::HttpResponse::decode(raw_msg.as_slice())
         .with_context(|| "Failed to decode response header.")?;
 
-    let status: u16 = match response_proto.status_code.try_into() {
-        Ok(status) => status,
-        Err(e) => {
-            return Err(anyhow!(
-                "Received invalid status code {} {}",
-                response_proto.status_code,
-                e
-            ))
-        }
-    };
+    let status: u16 = response_proto
+        .status_code
+        .try_into()
+        .with_context(|| "Failed to decode status code.")?;
 
     let mut response = Response::builder().status(status).version(Version::HTTP_3);
     for h in response_proto.headers {
@@ -188,7 +178,7 @@ async fn write_request(
             Method::CONNECT => pb::HttpMethod::Connect.into(),
             Method::PATCH => pb::HttpMethod::Patch.into(),
             Method::TRACE => pb::HttpMethod::Trace.into(),
-            _ => return Err(anyhow!("invalid method")),
+            _ => pb::HttpMethod::Unspecified.into(),
         },
         body: body.into(),
     };
@@ -197,8 +187,7 @@ async fn write_request(
     send_stream
         .write_all(&request_bytes)
         .await
-        .with_context(|| "Failed to write request to stream.")?;
-    Ok(())
+        .with_context(|| "Failed to write request to stream.")
 }
 
 // tests

rs/p2p/quic_transport/src/connection_manager.rs

@@ -523,12 +523,14 @@ impl ConnectionManager {
                 self.conn_id_counter.inc_assign();
                 let conn_id = self.conn_id_counter;
 
-                let connection_handle =
-                    ConnectionHandle::new(connection, self.metrics.clone(), conn_id);
-                let req_handler_connection_handle = connection_handle.clone();
+                let connection_handle = ConnectionHandle {
+                    connection,
+                    metrics: self.metrics.clone(),
+                    conn_id,
+                };
 
                 // dropping the old connection will result in closing it
-                if let Some(old_conn) = peer_map_mut.insert(peer_id, connection_handle) {
+                if let Some(old_conn) = peer_map_mut.insert(peer_id, connection_handle.clone()) {
                     old_conn
                         .connection
                         .close(VarInt::from_u32(0), b"using newer connection");
@@ -549,9 +551,7 @@ impl ConnectionManager {
                     run_stream_acceptor(
                         self.log.clone(),
                         peer_id,
-                        req_handler_connection_handle.conn_id(),
-                        req_handler_connection_handle.connection,
-                        self.metrics.clone(),
+                        connection_handle,
                         self.router.clone(),
                     ),
                     &self.rt,

rs/p2p/quic_transport/src/lib.rs

@@ -203,7 +203,7 @@ impl Transport for QuicTransport {
             .read()
             .unwrap()
             .iter()
-            .map(|(n, c)| (*n, c.conn_id()))
+            .map(|(n, c)| (*n, c.conn_id))
             .collect()
     }
 }

rs/p2p/quic_transport/src/request_handler.rs

@@ -12,22 +12,20 @@
 //!
 use std::time::Duration;
 
-use anyhow::{anyhow, Context};
-use axum::{
-    body::Body,
-    http::{Method, Request, Response, Version},
-    Router,
-};
+use anyhow::Context;
+use axum::{body::Body, Router};
 use bytes::Bytes;
+use http::{Method, Request, Response, Version};
 use ic_base_types::NodeId;
 use ic_logger::{info, ReplicaLogger};
 use ic_protobuf::transport::v1 as pb;
 use prost::Message;
-use quinn::{Connection, RecvStream, SendStream};
+use quinn::{RecvStream, SendStream};
 use tower::ServiceExt;
 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,
@@ -37,12 +35,10 @@ use crate::{
 
 const QUIC_METRIC_SCRAPE_INTERVAL: Duration = Duration::from_secs(5);
 
-pub(crate) async fn run_stream_acceptor(
+pub async fn run_stream_acceptor(
     log: ReplicaLogger,
     peer_id: NodeId,
-    conn_id: ConnId,
-    connection: Connection,
-    metrics: QuicTransportMetrics,
+    conn_handle: ConnectionHandle,
     router: Router,
 ) {
     let mut inflight_requests = tokio::task::JoinSet::new();
@@ -55,18 +51,18 @@ pub(crate) async fn run_stream_acceptor(
     loop {
         tokio::select! {
              _ = quic_metrics_scrape.tick() => {
-                metrics.collect_quic_connection_stats(&connection, &peer_id);
+                conn_handle.metrics.collect_quic_connection_stats(&conn_handle.connection, &peer_id);
             }
-            bi = connection.accept_bi() => {
+            bi = conn_handle.connection.accept_bi() => {
                 match bi {
                     Ok((bi_tx, bi_rx)) => {
                         inflight_requests.spawn(
-                            metrics.request_task_monitor.instrument(
+                            conn_handle.metrics.request_task_monitor.instrument(
                                 handle_bi_stream(
                                     log.clone(),
                                     peer_id,
-                                    conn_id,
-                                    metrics.clone(),
+                                    conn_handle.conn_id,
+                                    conn_handle.metrics.clone(),
                                     router.clone(),
                                     bi_tx,
                                     bi_rx
@@ -76,7 +72,7 @@ pub(crate) async fn run_stream_acceptor(
                     }
                     Err(e) => {
                         info!(log, "Error accepting bi stream {}", e.to_string());
-                        metrics
+                        conn_handle.metrics
                             .request_handle_errors_total
                             .with_label_values(&[
                                 STREAM_TYPE_BIDI,
@@ -87,8 +83,8 @@ pub(crate) async fn run_stream_acceptor(
                     }
                 }
             },
-            _ = connection.accept_uni() => {},
-            _ = connection.read_datagram() => {},
+            _ = conn_handle.connection.accept_uni() => {},
+            _ = conn_handle.connection.read_datagram() => {},
             Some(completed_request) = inflight_requests.join_next() => {
                 if let Err(err) = completed_request {
                     // Cancelling tasks is ok. Panicking tasks are not.
@@ -181,33 +177,34 @@ async fn read_request(mut recv_stream: RecvStream) -> Result<Request<Body>, anyh
     let request_proto = pb::HttpRequest::decode(raw_msg.as_slice())
         .with_context(|| "Failed to decode http request.")?;
 
-    let mut request = Request::builder()
-        .method(match pb::HttpMethod::try_from(request_proto.method) {
-            Ok(pb::HttpMethod::Get) => Method::GET,
-            Ok(pb::HttpMethod::Post) => Method::POST,
-            Ok(pb::HttpMethod::Put) => Method::PUT,
-            Ok(pb::HttpMethod::Delete) => Method::DELETE,
-            Ok(pb::HttpMethod::Head) => Method::HEAD,
-            Ok(pb::HttpMethod::Options) => Method::OPTIONS,
-            Ok(pb::HttpMethod::Connect) => Method::CONNECT,
-            Ok(pb::HttpMethod::Patch) => Method::PATCH,
-            Ok(pb::HttpMethod::Trace) => Method::TRACE,
-            Ok(pb::HttpMethod::Unspecified) => {
-                return Err(anyhow!("received http method unspecified."));
-            }
-            Err(e) => {
-                return Err(anyhow!("received invalid method {}", e));
-            }
-        })
+    let pb_http_method = pb::HttpMethod::try_from(request_proto.method)
+        .with_context(|| "Failed to decode http method.")?;
+    let http_method = match pb_http_method {
+        pb::HttpMethod::Get => Some(Method::GET),
+        pb::HttpMethod::Post => Some(Method::POST),
+        pb::HttpMethod::Put => Some(Method::PUT),
+        pb::HttpMethod::Delete => Some(Method::DELETE),
+        pb::HttpMethod::Head => Some(Method::HEAD),
+        pb::HttpMethod::Options => Some(Method::OPTIONS),
+        pb::HttpMethod::Connect => Some(Method::CONNECT),
+        pb::HttpMethod::Patch => Some(Method::PATCH),
+        pb::HttpMethod::Trace => Some(Method::TRACE),
+        pb::HttpMethod::Unspecified => None,
+    };
+    let mut request_builder = Request::builder();
+    if let Some(http_method) = http_method {
+        request_builder = request_builder.method(http_method);
+    }
+    request_builder = request_builder
         .version(Version::HTTP_3)
         .uri(request_proto.uri);
     for h in request_proto.headers {
         let pb::HttpHeader { key, value } = h;
-        request = request.header(key, value);
+        request_builder = request_builder.header(key, value);
     }
     // This consumes the body without requiring allocation or cloning the whole content.
     let body_bytes = Bytes::from(request_proto.body);
-    request
+    request_builder
         .body(Body::from(body_bytes))
         .with_context(|| "Failed to build request.")
 }
@@ -241,6 +238,5 @@ async fn write_response(
     send_stream
         .write_all(&response_bytes)
         .await
-        .with_context(|| "Failed to write request to stream.")?;
-    Ok(())
+        .with_context(|| "Failed to write request to stream.")
 }