Add error handling and retries for using cached connections

Author: Kimahriman

crates/hdfs-native/src/hdfs/block_reader.rs

@@ -6,12 +6,15 @@ use futures::{
     stream::{self, BoxStream},
     Stream, StreamExt,
 };
-use log::debug;
+use log::{debug, warn};
 
 use crate::{
     ec::EcSchema,
     hdfs::connection::{DatanodeConnection, Op, DATANODE_CACHE},
-    proto::{common, hdfs},
+    proto::{
+        common,
+        hdfs::{self, BlockOpResponseProto},
+    },
     HdfsError, Result,
 };
 
@@ -32,6 +35,55 @@ pub(crate) fn get_block_stream(
     }
 }
 
+/// Connects to a DataNode to do a read, attempting to used cached connections.
+async fn connect_and_send(
+    url: &str,
+    block: &hdfs::ExtendedBlockProto,
+    token: common::TokenProto,
+    offset: u64,
+    len: u64,
+) -> Result<(DatanodeConnection, BlockOpResponseProto)> {
+    let mut remaining_attempts = 2;
+    while remaining_attempts > 0 {
+        if let Some(mut conn) = DATANODE_CACHE.get(url) {
+            let message = hdfs::OpReadBlockProto {
+                header: conn.build_header(block, Some(token.clone())),
+                offset,
+                len,
+                send_checksums: Some(true),
+                ..Default::default()
+            };
+            debug!("Block read op request {:?}", &message);
+            match conn.send(Op::ReadBlock, &message).await {
+                Ok(response) => {
+                    debug!("Block read op response {:?}", response);
+                    return Ok((conn, response));
+                }
+                Err(e) => {
+                    warn!("Failed to use cached connection: {:?}", e);
+                }
+            }
+        } else {
+            break;
+        }
+        remaining_attempts -= 1;
+    }
+    let mut conn = DatanodeConnection::connect(url).await?;
+
+    let message = hdfs::OpReadBlockProto {
+        header: conn.build_header(block, Some(token)),
+        offset,
+        len,
+        send_checksums: Some(true),
+        ..Default::default()
+    };
+
+    debug!("Block read op request {:?}", &message);
+    let response = conn.send(Op::ReadBlock, &message).await?;
+    debug!("Block read op response {:?}", response);
+    Ok((conn, response))
+}
+
 struct ReplicatedBlockStream {
     block: hdfs::LocatedBlockProto,
     offset: usize,
@@ -63,26 +115,18 @@ impl ReplicatedBlockStream {
                 ));
             }
         }
-        let datanode = &self.block.locs[self.current_replica].id;
 
+        let datanode = &self.block.locs[self.current_replica].id;
         let datanode_url = format!("{}:{}", datanode.ip_addr, datanode.xfer_port);
-        let mut connection = if let Some(conn) = DATANODE_CACHE.get(&datanode_url) {
-            conn
-        } else {
-            DatanodeConnection::connect(&datanode_url).await?
-        };
 
-        let message = hdfs::OpReadBlockProto {
-            header: connection.build_header(&self.block.b, Some(self.block.block_token.clone())),
-            offset: self.offset as u64,
-            len: self.len as u64,
-            send_checksums: Some(true),
-            ..Default::default()
-        };
-
-        debug!("Block read op request {:?}", &message);
-        let response = connection.send(Op::ReadBlock, &message).await?;
-        debug!("Block read op response {:?}", response);
+        let (connection, response) = connect_and_send(
+            &datanode_url,
+            &self.block.b,
+            self.block.block_token.clone(),
+            self.offset as u64,
+            self.len as u64,
+        )
+        .await?;
 
         if response.status() != hdfs::Status::Success {
             return Err(HdfsError::DataTransferError(response.message().to_string()));
@@ -346,23 +390,14 @@ impl StripedBlockStream {
         }
 
         let datanode_url = format!("{}:{}", datanode.ip_addr, datanode.xfer_port);
-        let mut connection = if let Some(conn) = DATANODE_CACHE.get(&datanode_url) {
-            conn
-        } else {
-            DatanodeConnection::connect(&datanode_url).await?
-        };
-
-        let message = hdfs::OpReadBlockProto {
-            header: connection.build_header(block, Some(token.clone())),
-            offset: offset as u64,
-            len: len as u64,
-            send_checksums: Some(true),
-            ..Default::default()
-        };
-
-        debug!("Block read op request {:?}", &message);
-        let response = connection.send(Op::ReadBlock, &message).await?;
-        debug!("Block read op response {:?}", response);
+        let (mut connection, response) = connect_and_send(
+            &datanode_url,
+            block,
+            token.clone(),
+            self.offset as u64,
+            self.len as u64,
+        )
+        .await?;
 
         if response.status() != hdfs::Status::Success {
             return Err(HdfsError::DataTransferError(response.message().to_string()));