From 823b775bf51c85f84e4a79ec44792197ee0f68ad Mon Sep 17 00:00:00 2001
From: Georgy Moshkin <gmoshkin@picodata.io>
Date: Wed, 19 Oct 2022 15:34:32 +0300
Subject: [PATCH] refactor: add a macro for defining PeerChange

---
 src/traft/mod.rs      | 65 ++++++++++++++++++++++++++++++-------------
 src/traft/topology.rs | 33 ++++++++++------------
 2 files changed, 59 insertions(+), 39 deletions(-)

diff --git a/src/traft/mod.rs b/src/traft/mod.rs
index 56dadca5a7..8501db575d 100644
--- a/src/traft/mod.rs
+++ b/src/traft/mod.rs
@@ -893,11 +893,51 @@ pub struct UpdatePeerRequest {
     pub changes: Vec<PeerChange>,
 }
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
-pub enum PeerChange {
-    CurrentGrade(CurrentGrade),
-    TargetGrade(TargetGrade),
-    FailureDomain(FailureDomain),
+macro_rules! define_peer_change {
+    (
+        $(#[$meta:meta])*
+        pub enum $enum:ident {
+            $(
+                #[setter = $setter:ident, field = $field:ident]
+                $variant:ident($type:ty),
+            )*
+        }
+    ) => {
+        $(#[$meta])*
+        pub enum $enum {
+            $( $variant($type), )*
+        }
+
+        impl $enum {
+            pub fn apply(self, peer: &mut Peer) {
+                match self {
+                    $( Self::$variant(value) => peer.$field = value, )*
+                }
+            }
+        }
+
+        impl UpdatePeerRequest {
+            $(
+                #[inline]
+                pub fn $setter(mut self, value: $type) -> Self {
+                    self.changes.push($enum::$variant(value));
+                    self
+                }
+            )*
+        }
+    }
+}
+
+define_peer_change! {
+    #[derive(Clone, Debug, Serialize, Deserialize)]
+    pub enum PeerChange {
+        #[setter = with_current_grade, field = current_grade]
+        CurrentGrade(CurrentGrade),
+        #[setter = with_target_grade, field = target_grade]
+        TargetGrade(TargetGrade),
+        #[setter = with_failure_domain, field = failure_domain]
+        FailureDomain(FailureDomain),
+    }
 }
 
 impl Encode for UpdatePeerRequest {}
@@ -910,21 +950,6 @@ impl UpdatePeerRequest {
             changes: vec![],
         }
     }
-    #[inline]
-    pub fn with_current_grade(mut self, current_grade: CurrentGrade) -> Self {
-        self.changes.push(PeerChange::CurrentGrade(current_grade));
-        self
-    }
-    #[inline]
-    pub fn with_target_grade(mut self, target_grade: TargetGrade) -> Self {
-        self.changes.push(PeerChange::TargetGrade(target_grade));
-        self
-    }
-    #[inline]
-    pub fn with_failure_domain(mut self, failure_domain: FailureDomain) -> Self {
-        self.changes.push(PeerChange::FailureDomain(failure_domain));
-        self
-    }
 }
 
 ///////////////////////////////////////////////////////////////////////////////
diff --git a/src/traft/topology.rs b/src/traft/topology.rs
index dc53cbec92..69677fdbde 100644
--- a/src/traft/topology.rs
+++ b/src/traft/topology.rs
@@ -193,31 +193,26 @@ impl Topology {
     pub fn update_peer(&mut self, req: UpdatePeerRequest) -> Result<Peer, String> {
         let this = self as *const Self;
 
-        let mut peer = self
+        let peer = self
             .instance_map
             .get_mut(&req.instance_id)
             .ok_or_else(|| format!("unknown instance {}", req.instance_id))?;
 
+        if peer.current_grade == CurrentGrade::Expelled {
+            return Err(format!(
+                "cannot update expelled peer \"{}\"",
+                peer.instance_id
+            ));
+        }
+
         for change in req.changes {
-            match change {
-                PeerChange::CurrentGrade(grade) => {
-                    if peer.current_grade != CurrentGrade::Expelled {
-                        peer.current_grade = grade;
-                    }
-                }
-                PeerChange::TargetGrade(target_grade) => {
-                    if peer.target_grade != TargetGrade::Expelled {
-                        peer.target_grade = target_grade;
-                    }
-                }
-                PeerChange::FailureDomain(fd) => {
-                    // SAFETY: this is safe, because rust doesn't complain if you inline
-                    // the function
-                    unsafe { &*this }.check_required_failure_domain(&fd)?;
-                    self.failure_domain_names.extend(fd.names().cloned());
-                    peer.failure_domain = fd;
-                }
+            if let PeerChange::FailureDomain(fd) = &change {
+                // SAFETY: this is safe, because rust doesn't complain if you inline
+                // the function
+                unsafe { &*this }.check_required_failure_domain(fd)?;
+                self.failure_domain_names.extend(fd.names().cloned());
             }
+            change.apply(peer);
         }
 
         Ok(peer.clone())
-- 
GitLab