Change: EffectiveMembership::get_node() should return an Option

`EffectiveMembership::get_node()` should return an `Option<&Node>` instead of a `&Node`. Otherwise it panic if the node is not found.
databendlabs · Nov 2, 2022 · e8ec9c5 · e8ec9c5
1 parent 61d47bc
commit e8ec9c5
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/openraft/src/core/raft_core.rs b/openraft/src/core/raft_core.rs
@@ -326,7 +326,8 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
             };
 
             let my_id = self.id;
-            let target_node = self.engine.state.membership_state.effective.get_node(&target).clone();
+            // Safe unwrap(): target is in membership
+            let target_node = self.engine.state.membership_state.effective.get_node(&target).unwrap().clone();
             let mut client = match self.network.new_client(target, &target_node).await {
                 Ok(n) => n,
                 Err(e) => {
@@ -891,7 +892,12 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
     }
 
     pub(crate) fn get_leader_node(&self, leader_id: Option<C::NodeId>) -> Option<C::Node> {
-        leader_id.map(|id| self.engine.state.membership_state.effective.get_node(&id).clone())
+        let leader_id = match leader_id {
+            None => return None,
+            Some(x) => x,
+        };
+
+        self.engine.state.membership_state.effective.get_node(&leader_id).cloned()
     }
 
     #[tracing::instrument(level = "debug", skip_all)]
@@ -979,7 +985,9 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
         target: C::NodeId,
         matched: Option<LogId<C::NodeId>>,
     ) -> Result<ReplicationStream<C::NodeId>, N::ConnectionError> {
-        let target_node = self.engine.state.membership_state.effective.get_node(&target);
+        // Safe unwrap(): target must be in membership
+        let target_node = self.engine.state.membership_state.effective.get_node(&target).unwrap();
+
         let membership_log_id = self.engine.state.membership_state.effective.log_id;
         let network = self.network.new_client(target, target_node).await?;
 
@@ -1125,7 +1133,9 @@ impl<C: RaftTypeConfig, N: RaftNetworkFactory<C>, S: RaftStorage<C>> RaftCore<C,
             }
 
             let req = vote_req.clone();
-            let target_node = self.engine.state.membership_state.effective.get_node(&target).clone();
+
+            // Safe unwrap(): target must be in membership
+            let target_node = self.engine.state.membership_state.effective.get_node(&target).unwrap().clone();
             let mut client = match self.network.new_client(target, &target_node).await {
                 Ok(n) => n,
                 Err(err) => {

diff --git a/openraft/src/membership/effective_membership.rs b/openraft/src/membership/effective_membership.rs
@@ -151,7 +151,7 @@ where
     }
 
     /// Get a the node(either voter or learner) by node id.
-    pub fn get_node(&self, node_id: &NID) -> &N {
+    pub fn get_node(&self, node_id: &NID) -> Option<&N> {
         self.membership.get_node(node_id)
     }
 

diff --git a/openraft/src/membership/membership.rs b/openraft/src/membership/membership.rs
@@ -1,4 +1,3 @@
-use core::option::Option;
 use std::collections::BTreeMap;
 use std::collections::BTreeSet;
 
@@ -117,7 +116,7 @@ where
                 }
                 res.push(format!("{}", node_id));
 
-                let n = self.get_node(node_id);
+                let n = self.get_node(node_id).unwrap();
                 res.push(format!(":{{{:?}}}", n));
             }
             res.push("}".to_string());
@@ -135,7 +134,7 @@ where
 
             res.push(format!("{}", learner_id));
 
-            let n = self.get_node(learner_id);
+            let n = self.get_node(learner_id).unwrap();
             res.push(format!(":{{{:?}}}", n));
         }
         res.push("]".to_string());
@@ -262,8 +261,8 @@ where
     }
 
     /// Get a the node(either voter or learner) by node id.
-    pub(crate) fn get_node(&self, node_id: &NID) -> &N {
-        &self.nodes[node_id]
+    pub(crate) fn get_node(&self, node_id: &NID) -> Option<&N> {
+        self.nodes.get(node_id)
     }
 
     /// Returns an Iterator of all nodes(voters and learners).