support replace node policy

Author: zhaohaiyuan

rust/src/hdfs/block_writer.rs

@@ -10,6 +10,7 @@ use crate::{
     hdfs::{
         connection::{DatanodeConnection, DatanodeReader, DatanodeWriter, Op, WritePacket},
         protocol::NamenodeProtocol,
+        replace_datanode::ReplaceDatanodeOnFailure,
     },
     proto::hdfs,
     HdfsError, Result,
@@ -300,6 +301,7 @@ pub(crate) struct ReplicatedBlockWriter {
 
     current_packet: WritePacket,
     pipeline: Option<Pipeline>,
+    replace_datanode: ReplaceDatanodeOnFailure,
 }
 
 impl ReplicatedBlockWriter {
@@ -340,6 +342,7 @@ impl ReplicatedBlockWriter {
             current_packet,
 
             pipeline: Some(pipeline),
+            replace_datanode: ReplaceDatanodeOnFailure::default(),
         };
 
         Ok(this)
@@ -363,10 +366,51 @@ impl ReplicatedBlockWriter {
         debug!("Recovering block writer");
 
         let mut new_block = self.block.clone();
-        for failed_node in failed_nodes {
-            new_block.locs.remove(failed_node);
-            new_block.storage_i_ds.remove(failed_node);
-            new_block.storage_types.remove(failed_node);
+
+        // Check if we should try to replace failed datanodes
+        let should_replace = self.replace_datanode.should_replace(
+            new_block.locs.len() as u32,
+            &new_block.locs,
+            new_block.b.num_bytes() > 0,
+            false,
+        );
+
+        if should_replace && self.replace_datanode.check_enabled() {
+            // Request additional datanodes to replace failed ones
+            let additional_nodes = self
+                .protocol
+                .get_additional_datanodes(&new_block.b, &new_block.locs, failed_nodes.len() as u32)
+                .await;
+
+            match additional_nodes {
+                Ok(new_nodes) => {
+                    // Replace failed nodes with new ones
+                    for (failed_idx, new_node) in failed_nodes.iter().zip(new_nodes.into_iter()) {
+                        if *failed_idx < new_block.locs.len() {
+                            new_block.locs[*failed_idx] = new_node;
+                        }
+                    }
+                }
+                Err(e) => {
+                    if !self.replace_datanode.is_best_effort() {
+                        return Err(e);
+                    }
+                    // In best effort mode, continue with remaining nodes
+                    warn!("Failed to get replacement nodes: {}", e);
+                    for failed_node in failed_nodes {
+                        new_block.locs.remove(failed_node);
+                        new_block.storage_i_ds.remove(failed_node);
+                        new_block.storage_types.remove(failed_node);
+                    }
+                }
+            }
+        } else {
+            // Original behavior - remove failed nodes
+            for failed_node in failed_nodes {
+                new_block.locs.remove(failed_node);
+                new_block.storage_i_ds.remove(failed_node);
+                new_block.storage_types.remove(failed_node);
+            }
         }
 
         let mut bytes_acked = new_block.b.num_bytes();

rust/src/hdfs/mod.rs

@@ -3,3 +3,4 @@ pub(crate) mod block_writer;
 pub(crate) mod connection;
 pub(crate) mod protocol;
 pub(crate) mod proxy;
+pub(crate) mod replace_datanode;

rust/src/hdfs/protocol.rs

@@ -444,6 +444,29 @@ impl NamenodeProtocol {
         };
         self.call("getAclStatus", message, false).await
     }
+
+    pub(crate) async fn get_additional_datanodes(
+        &self,
+        block: &hdfs::ExtendedBlockProto,
+        existing_nodes: &[hdfs::DatanodeInfoProto],
+        num_additional_nodes: u32,
+    ) -> Result<Vec<hdfs::DatanodeInfoProto>> {
+        let message = hdfs::GetAdditionalDatanodeRequestProto {
+            src: "".to_string(), // Not used by the namenode
+            blk: block.clone(),
+            existings: existing_nodes.to_vec(),
+            excludes: vec![], // No excluded nodes
+            num_additional_nodes,
+            client_name: self.client_name.clone(),
+            existing_storage_uuids: vec![], // Not used in this context
+            ..Default::default()
+        };
+
+        let response: hdfs::GetAdditionalDatanodeResponseProto =
+            self.call("getAdditionalDatanode", message, true).await?;
+
+        Ok(response.block.locs)
+    }
 }
 
 impl Drop for NamenodeProtocol {

rust/src/hdfs/replace_datanode.rs

@@ -0,0 +1,83 @@
+use crate::proto::hdfs::DatanodeInfoProto;
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum Policy {
+    /// The feature is disabled in the entire site
+    Disable,
+    /// Never add a new datanode
+    Never,
+    /// Default policy based on replication conditions
+    Default,
+    /// Always add a new datanode when an existing datanode is removed
+    Always,
+}
+
+pub struct ReplaceDatanodeOnFailure {
+    policy: Policy,
+    best_effort: bool,
+}
+
+impl ReplaceDatanodeOnFailure {
+    pub fn new(policy: Policy, best_effort: bool) -> Self {
+        Self {
+            policy,
+            best_effort,
+        }
+    }
+
+    /// Check if the feature is enabled
+    pub fn check_enabled(&self) -> bool {
+        self.policy != Policy::Disable
+    }
+
+    /// Best effort means that the client will try to replace the failed datanode
+    /// (provided that the policy is satisfied), however, it will continue the
+    /// write operation in case that the datanode replacement also fails.
+    pub fn is_best_effort(&self) -> bool {
+        self.best_effort
+    }
+
+    /// Does it need a replacement according to the policy?
+    pub fn should_replace(
+        &self,
+        replication: u32,
+        existing_datanodes: &[DatanodeInfoProto],
+        is_append: bool,
+        is_hflushed: bool,
+    ) -> bool {
+        let n = existing_datanodes.len();
+        if n == 0 || n >= replication as usize {
+            // Don't need to add datanode for any policy
+            return false;
+        }
+
+        match self.policy {
+            Policy::Disable | Policy::Never => false,
+            Policy::Always => true,
+            Policy::Default => {
+                // DEFAULT condition:
+                // Let r be the replication number.
+                // Let n be the number of existing datanodes.
+                // Add a new datanode only if r >= 3 and either
+                // (1) floor(r/2) >= n; or
+                // (2) r > n and the block is hflushed/appended.
+                if replication < 3 {
+                    false
+                } else if n <= (replication as usize / 2) {
+                    true
+                } else {
+                    is_append || is_hflushed
+                }
+            }
+        }
+    }
+}
+
+impl Default for ReplaceDatanodeOnFailure {
+    fn default() -> Self {
+        Self {
+            policy: Policy::Default,
+            best_effort: true,
+        }
+    }
+}