From 9681c4d468a883b7b16a5475ad7c825c55a4b140 Mon Sep 17 00:00:00 2001
From: Sagar Dhawan <sagar@solana.com>
Date: Wed, 6 Feb 2019 15:31:41 -0800
Subject: [PATCH] Fix resource hogging when waiting for role transition

---
 src/fullnode.rs | 33 ++++++++++++++++++---------------
 1 file changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/fullnode.rs b/src/fullnode.rs
index bed88c256c..0d6aaa0b41 100644
--- a/src/fullnode.rs
+++ b/src/fullnode.rs
@@ -403,28 +403,25 @@ impl Fullnode {
     }
 
     fn handle_role_transition(&mut self) -> Option<(FullnodeReturnType, u64)> {
-        loop {
-            if self.exit.load(Ordering::Relaxed) {
-                return None;
+        if self.node_services.tpu.is_leader() {
+            let should_be_forwarder = self.role_notifiers.1.recv();
+            match should_be_forwarder {
+                Ok(TpuReturnType::LeaderRotation(tick_height)) => {
+                    Some((self.leader_to_validator(tick_height), tick_height + 1))
+                }
+                Err(_) => None,
             }
-            let should_be_forwarder = self.role_notifiers.1.try_recv();
-            let should_be_leader = self.role_notifiers.0.try_recv();
+        } else {
+            let should_be_leader = self.role_notifiers.0.recv();
             match should_be_leader {
                 Ok(TvuReturnType::LeaderRotation(tick_height, entry_height, last_entry_id)) => {
                     self.validator_to_leader(tick_height, entry_height, last_entry_id);
-                    return Some((
+                    Some((
                         FullnodeReturnType::ValidatorToLeaderRotation,
                         tick_height + 1,
-                    ));
+                    ))
                 }
-                _ => match should_be_forwarder {
-                    Ok(TpuReturnType::LeaderRotation(tick_height)) => {
-                        return Some((self.leader_to_validator(tick_height), tick_height + 1))
-                    }
-                    _ => {
-                        continue;
-                    }
-                },
+                Err(_) => None,
             }
         }
     }
@@ -438,6 +435,12 @@ impl Fullnode {
         let (sender, receiver) = channel();
         let exit = self.exit.clone();
         spawn(move || loop {
+            if self.exit.load(Ordering::Relaxed) {
+                debug!("node shutdown requested");
+                self.close().expect("Unable to close node");
+                sender.send(true).expect("Unable to signal exit");
+                break;
+            }
             let status = self.handle_role_transition();
             match status {
                 None => {