From 214e5daac9af1909be6eb29c13a5e689e28e9a8a Mon Sep 17 00:00:00 2001
From: Hoverbear <andrew@hoverbear.org>
Date: Wed, 8 Aug 2018 11:36:09 -0700
Subject: [PATCH] Make ProgressSet not panic.

Moves failable functions to returning `Result<(), Error>` and adds
useful error types for determining the problem.

Adds an `panic` to all calling code to preserve existing behavior.
---
 src/errors.rs            | 13 ++++++++++++-
 src/progress.rs          | 41 +++++++++++++++++++++++++---------------
 src/raft.rs              | 10 +++++-----
 tests/cases/test_raft.rs |  4 ++--
 4 files changed, 45 insertions(+), 23 deletions(-)

diff --git a/src/errors.rs b/src/errors.rs
index 0b5935c69..8fe1b24aa 100644
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -55,7 +55,18 @@ quick_error! {
             description(err.description())
             display("protobuf error {:?}", err)
         }
-
+        /// The node exists, but should not.
+        Exists(id: u64, set: &'static str) {
+            display("The node {} aleady exists in the {} set.", id, set)
+        }
+        /// The node does not exist, but should.
+        NotExists(id: u64, set: &'static str) {
+            display("The node {} is not in the {} set.", id, set)
+        }
+        /// The node is attempting to be demoted, which isn't allowed.
+        Demotion(id: u64) {
+            display("The node {} is a voter already, and cannot become a learner.", id)
+        }
     }
 }
 
diff --git a/src/progress.rs b/src/progress.rs
index 1ed39b000..a48b89473 100644
--- a/src/progress.rs
+++ b/src/progress.rs
@@ -28,6 +28,7 @@
 use fxhash::FxHashMap;
 use std::cmp;
 use std::collections::hash_map::HashMap;
+use errors::Error;
 
 /// The state of the progress.
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -118,13 +119,15 @@ impl ProgressSet {
     /// # Panics
     ///
     /// Panics if the node already has been added.
-    pub fn insert_voter(&mut self, id: u64, pr: Progress) {
-        if self.learners.contains_key(&id) {
-            panic!("insert voter {} but already in learners", id);
+    pub fn insert_voter(&mut self, id: u64, pr: Progress) -> Result<(), Error> {
+        if self.voters.contains_key(&id) {
+            Err(Error::Exists(id, "voters"))?
         }
-        if self.voters.insert(id, pr).is_some() {
-            panic!("insert voter {} twice", id);
+        if self.learners.contains_key(&id) {
+            Err(Error::Exists(id, "learners"))?;
         }
+        self.voters.insert(id, pr);
+        Ok(())
     }
 
     /// Adds a learner to the cluster
@@ -132,13 +135,15 @@ impl ProgressSet {
     /// # Panics
     ///
     /// Panics if the node already has been added.
-    pub fn insert_learner(&mut self, id: u64, pr: Progress) {
+    pub fn insert_learner(&mut self, id: u64, pr: Progress) -> Result<(), Error> {
         if self.voters.contains_key(&id) {
-            panic!("insert learner {} but already in voters", id);
+            Err(Error::Demotion(id))?
         }
-        if self.learners.insert(id, pr).is_some() {
-            panic!("insert learner {} twice", id);
+        if self.learners.contains_key(&id) {
+            Err(Error::Exists(id, "learners"))?
         }
+        self.learners.insert(id, pr);
+        Ok(())
     }
 
     /// Removes the peer from the set of voters or learners.
@@ -154,13 +159,19 @@ impl ProgressSet {
     /// # Panics
     ///
     /// Panics if the node doesn't exist.
-    pub fn promote_learner(&mut self, id: u64) {
-        if let Some(mut pr) = self.learners.remove(&id) {
-            pr.is_learner = false;
-            self.voters.insert(id, pr);
-            return;
+    pub fn promote_learner(&mut self, id: u64) -> Result<(), Error> {
+        if self.voters.contains_key(&id) {
+            Err(Error::Exists(id, "voters"))?;
         }
-        panic!("promote not exists learner: {}", id);
+        if !self.learners.contains_key(&id) {
+            Err(Error::NotExists(id, "learners"))?
+        }
+
+        self.learners.remove(&id).map(|mut learner| {
+            learner.is_learner = false;
+            self.voters.insert(id, learner);
+        });
+        Ok(())
     }
 }
 
diff --git a/src/raft.rs b/src/raft.rs
index 8b86d9cab..64213b94e 100644
--- a/src/raft.rs
+++ b/src/raft.rs
@@ -263,12 +263,12 @@ impl<T: Storage> Raft<T> {
         };
         for p in peers {
             let pr = new_progress(1, r.max_inflight);
-            r.mut_prs().insert_voter(*p, pr);
+            r.mut_prs().insert_voter(*p, pr).unwrap();
         }
         for p in learners {
             let mut pr = new_progress(1, r.max_inflight);
             pr.is_learner = true;
-            r.mut_prs().insert_learner(*p, pr);
+            r.mut_prs().insert_learner(*p, pr).unwrap();
             if *p == r.id {
                 r.is_learner = true;
             }
@@ -1854,7 +1854,7 @@ impl<T: Storage> Raft<T> {
                 // Ignore redundant add learner.
                 return;
             }
-            self.mut_prs().promote_learner(id);
+            self.mut_prs().promote_learner(id).unwrap();
             if id == self.id {
                 self.is_learner = false;
             }
@@ -1905,9 +1905,9 @@ impl<T: Storage> Raft<T> {
         p.matched = matched;
         p.is_learner = is_learner;
         if is_learner {
-            self.mut_prs().insert_learner(id, p);
+            self.mut_prs().insert_learner(id, p).unwrap();
         } else {
-            self.mut_prs().insert_voter(id, p);
+            self.mut_prs().insert_voter(id, p).unwrap();
         }
     }
 
diff --git a/tests/cases/test_raft.rs b/tests/cases/test_raft.rs
index 5b385c548..f4a6e7eaa 100644
--- a/tests/cases/test_raft.rs
+++ b/tests/cases/test_raft.rs
@@ -209,12 +209,12 @@ impl Interface {
                         is_learner: true,
                         ..Default::default()
                     };
-                    self.mut_prs().insert_learner(*id, progress);
+                    self.mut_prs().insert_learner(*id, progress).unwrap();
                 } else {
                     let progress = Progress {
                         ..Default::default()
                     };
-                    self.mut_prs().insert_voter(*id, progress);
+                    self.mut_prs().insert_voter(*id, progress).unwrap();
                 }
             }
             let term = self.term;