Change: trait IntoNodes adds two new method has_nodes() and node_ids()

drmingdrmer · drmingdrmer · commit f591726a1a9e · 2023-02-19T11:29:24.000+08:00
`trait IntoNodes` converts types `T` such as `Vec` or `BTreeSet` into
`BTreeMap&lt;NID, Node&gt;`.

This patch changes the functionality of the `IntoNodes` trait to provide
two new methods `has_nodes()` and `node_ids()`, in addition to the existing
`into_nodes()` method. The `has_nodes()` method returns true if the type `T`
contains any `Node` objects, and `node_ids()` returns a `Vec` of the `NodeId`
objects associated with the `Node` objects in `T`.

Refactor:

The patch also refactors the `Membership::next_safe()` method to return an
`Err(LearnerNotFound)` if it attempts to build a `Membership` object
containing a `voter_id` that does not correspond to any `Node`.
diff --git a/openraft/src/membership/membership.rs b/openraft/src/membership/membership.rs
@@ -3,6 +3,7 @@ use std::collections::BTreeSet;
 
 use maplit::btreemap;
 
+use crate::error::LearnerNotFound;
 use crate::membership::NodeRole;
 use crate::node::Node;
 use crate::quorum::AsJoint;
@@ -18,6 +19,8 @@ where
     N: Node,
     NID: NodeId,
 {
+    fn has_nodes(&self) -> bool;
+    fn node_ids(&self) -> Vec<NID>;
     fn into_nodes(self) -> BTreeMap<NID, N>;
 }
 
@@ -26,6 +29,14 @@ where
     N: Node,
     NID: NodeId,
 {
+    fn has_nodes(&self) -> bool {
+        false
+    }
+
+    fn node_ids(&self) -> Vec<NID> {
+        vec![]
+    }
+
     fn into_nodes(self) -> BTreeMap<NID, N> {
         btreemap! {}
     }
@@ -36,6 +47,14 @@ where
     N: Node,
     NID: NodeId,
 {
+    fn has_nodes(&self) -> bool {
+        false
+    }
+
+    fn node_ids(&self) -> Vec<NID> {
+        self.iter().copied().collect()
+    }
+
     fn into_nodes(self) -> BTreeMap<NID, N> {
         self.into_iter().map(|node_id| (node_id, N::default())).collect()
     }
@@ -46,6 +65,19 @@ where
     N: Node,
     NID: NodeId,
 {
+    fn has_nodes(&self) -> bool {
+        true
+    }
+
+    fn node_ids(&self) -> Vec<NID> {
+        match self {
+            None => {
+                vec![]
+            }
+            Some(bs) => bs.iter().copied().collect(),
+        }
+    }
+
     fn into_nodes(self) -> BTreeMap<NID, N> {
         match self {
             None => BTreeMap::new(),
@@ -59,6 +91,14 @@ where
     N: Node,
     NID: NodeId,
 {
+    fn has_nodes(&self) -> bool {
+        true
+    }
+
+    fn node_ids(&self) -> Vec<NID> {
+        self.keys().copied().collect()
+    }
+
     fn into_nodes(self) -> BTreeMap<NID, N> {
         self
     }
@@ -304,11 +344,28 @@ where
     ///     curr = next;
     /// }
     /// ```
-    pub(crate) fn next_safe<T>(&self, goal: T, removed_to_learner: bool) -> Self
+    pub(crate) fn next_safe<T>(&self, goal: T, removed_to_learner: bool) -> Result<Self, LearnerNotFound<NID>>
     where T: IntoNodes<NID, N> {
-        let goal = goal.into_nodes();
+        let goal = if goal.has_nodes() {
+            goal.into_nodes()
+        } else {
+            // If `goal` does not contains Node, inherit them from current config.
+
+            let mut voters_map = BTreeMap::new();
+
+            // There has to be corresponding `Node` for every voter_id
+            for node_id in goal.node_ids().iter() {
+                let n = self.get_node(node_id);
+                if let Some(n) = n {
+                    voters_map.insert(*node_id, n.clone());
+                } else {
+                    return Err(LearnerNotFound { node_id: *node_id });
+                }
+            }
+            voters_map
+        };
 
-        let goal_ids = goal.keys().cloned().collect::<BTreeSet<_>>();
+        let goal_ids = goal.keys().copied().collect::<BTreeSet<_>>();
 
         let config = Joint::from(self.configs.clone()).find_coherent(goal_ids).children().clone();
 
@@ -323,7 +380,7 @@ where
             }
         };
 
-        Membership::with_nodes(config, nodes)
+        Ok(Membership::with_nodes(config, nodes))
     }
 
     /// Build a QuorumSet from current joint config
diff --git a/openraft/src/membership/membership_state.rs b/openraft/src/membership/membership_state.rs
@@ -4,7 +4,6 @@ use std::sync::Arc;
 use crate::error::ChangeMembershipError;
 use crate::error::EmptyMembership;
 use crate::error::InProgress;
-use crate::error::LearnerNotFound;
 use crate::less_equal;
 use crate::node::Node;
 use crate::validate::Validate;
@@ -92,21 +91,6 @@ where
         let effective = self.effective();
         let committed = self.committed();
 
-        let last = effective.membership.get_joint_config().last().unwrap();
-        let new_voter_ids = changes.apply_to(last);
-
-        // Ensure cluster will have at least one voter.
-        if new_voter_ids.is_empty() {
-            return Err(EmptyMembership {}.into());
-        }
-
-        // There has to be corresponding `Node` for every voter_id
-        for node_id in new_voter_ids.iter() {
-            if !effective.contains(node_id) {
-                return Err(LearnerNotFound { node_id: *node_id }.into());
-            }
-        }
-
         if committed.log_id == effective.log_id {
             // Ok: last membership(effective) is committed
         } else {
@@ -117,7 +101,15 @@ where
             .into());
         }
 
-        let new_membership = effective.membership.next_safe(new_voter_ids, removed_to_learner);
+        let last = effective.membership.get_joint_config().last().unwrap();
+        let new_voter_ids = changes.apply_to(last);
+
+        // Ensure cluster will have at least one voter.
+        if new_voter_ids.is_empty() {
+            return Err(EmptyMembership {}.into());
+        }
+
+        let new_membership = effective.membership.next_safe(new_voter_ids, removed_to_learner)?;
 
         tracing::debug!(?new_membership, "new membership config");
         Ok(new_membership)
diff --git a/openraft/src/membership/membership_test.rs b/openraft/src/membership/membership_test.rs
@@ -5,6 +5,8 @@ use std::fmt::Formatter;
 use maplit::btreemap;
 use maplit::btreeset;
 
+use crate::error::LearnerNotFound;
+use crate::membership::IntoNodes;
 use crate::Membership;
 use crate::MessageSummary;
 
@@ -206,28 +208,56 @@ fn test_membership_with_nodes_panic() {
 
 #[test]
 fn test_membership_next_safe() -> anyhow::Result<()> {
+    let nodes = || btreeset! {1,2,3,4,5,6,7,8,9}.into_nodes();
     let c1 = || btreeset! {1,2,3};
     let c2 = || btreeset! {3,4,5};
     let c3 = || btreeset! {7,8,9};
 
-    let m1 = Membership::<u64, ()>::new(vec![c1()], None);
-    let m2 = Membership::<u64, ()>::new(vec![c2()], None);
-    let m12 = Membership::<u64, ()>::new(vec![c1(), c2()], None);
-    let m23 = Membership::<u64, ()>::new(vec![c2(), c3()], None);
+    #[allow(clippy::redundant_closure)]
+    let new_mem = |voter_ids, ns| Membership::<u64, ()>::new(voter_ids, ns);
 
-    assert_eq!(m1, m1.next_safe(c1(), false));
-    assert_eq!(m12, m1.next_safe(c2(), false));
-    assert_eq!(m1, m12.next_safe(c1(), false));
-    assert_eq!(m2, m12.next_safe(c2(), false));
-    assert_eq!(m23, m12.next_safe(c3(), false));
+    let m1 = new_mem(vec![c1()], nodes());
+    let m12 = new_mem(vec![c1(), c2()], nodes());
+
+    assert_eq!(m1, m1.next_safe(c1(), false)?);
+    assert_eq!(m12, m1.next_safe(c2(), false)?);
+    assert_eq!(
+        new_mem(vec![c1()], btreeset! {1,2,3,6,7,8,9}.into_nodes()),
+        m12.next_safe(c1(), false)?
+    );
+    assert_eq!(
+        new_mem(vec![c2()], btreeset! {3,4,5,6,7,8,9}.into_nodes()),
+        m12.next_safe(c2(), false)?
+    );
+    assert_eq!(
+        new_mem(vec![c2(), c3()], btreeset! {3,4,5,6,7,8,9}.into_nodes()),
+        m12.next_safe(c3(), false)?
+    );
 
     // Turn removed members to learners
 
     let old_learners = || btreeset! {1, 2};
     let learners = || btreeset! {1, 2, 3, 4, 5};
     let m23_with_learners_old = Membership::<u64, ()>::new(vec![c2(), c3()], Some(old_learners()));
     let m23_with_learners_new = Membership::<u64, ()>::new(vec![c3()], Some(learners()));
-    assert_eq!(m23_with_learners_new, m23_with_learners_old.next_safe(c3(), true));
+    assert_eq!(m23_with_learners_new, m23_with_learners_old.next_safe(c3(), true)?);
+
+    Ok(())
+}
+
+#[test]
+fn test_membership_next_safe_learner_not_found() -> anyhow::Result<()> {
+    let c1 = || btreeset! {1,2,3};
+    let c2 = || btreeset! {3,4,5};
+    let c3 = || btreeset! {7};
+
+    #[allow(clippy::redundant_closure)]
+    let new_mem = |voter_ids, ns| Membership::<u64, ()>::new(voter_ids, ns);
+
+    let m12 = new_mem(vec![c1(), c2()], None);
+
+    let res = m12.next_safe(c3(), false);
+    assert_eq!(Err(LearnerNotFound { node_id: 7 }), res);
 
     Ok(())
 }
@@ -246,15 +276,15 @@ fn test_membership_next_safe_with_nodes() -> anyhow::Result<()> {
 
     // joint [{2}, {1,2}]
 
-    let res = initial.next_safe(btreeset! {1,2}, false);
+    let res = initial.next_safe(btreeset! {1,2}, false)?;
     assert_eq!(
         btreemap! {1=>node("1"), 2=>node("2")},
         res.nodes().map(|(nid, n)| (*nid, n.clone())).collect::<BTreeMap<_, _>>()
     );
 
     // Removed to learner
 
-    let res = initial.next_safe(btreeset! {1}, true);
+    let res = initial.next_safe(btreeset! {1}, true)?;
     assert_eq!(
         btreemap! {1=>node("1"), 2=>node("2")},
         res.nodes().map(|(nid, n)| (*nid, n.clone())).collect::<BTreeMap<_, _>>()
