From 0bb6940c1aa68dad29243a4b4ab8f29e73dcc8fc Mon Sep 17 00:00:00 2001
From: Stephen Akridge <stephen@solana.com>
Date: Tue, 12 Feb 2019 17:53:48 -0800
Subject: [PATCH] Faster exit for storage_stage client

Shorten the timeout and check for exit on every iteration
of fetching a last id.
---
 src/client.rs        |  6 ++++++
 src/storage_stage.rs | 18 +++++++++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 015647d58..65d7358a2 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -1,7 +1,13 @@
 use crate::cluster_info::{NodeInfo, FULLNODE_PORT_RANGE};
 use crate::thin_client::ThinClient;
+use std::time::Duration;
 
 pub fn mk_client(r: &NodeInfo) -> ThinClient {
     let (_, transactions_socket) = solana_netutil::bind_in_range(FULLNODE_PORT_RANGE).unwrap();
     ThinClient::new(r.rpc, r.tpu, transactions_socket)
 }
+
+pub fn mk_client_with_timeout(r: &NodeInfo, timeout: Duration) -> ThinClient {
+    let (_, transactions_socket) = solana_netutil::bind_in_range(FULLNODE_PORT_RANGE).unwrap();
+    ThinClient::new_with_timeout(r.rpc, r.tpu, transactions_socket, timeout)
+}
diff --git a/src/storage_stage.rs b/src/storage_stage.rs
index 8c3b49b69..c27be78e7 100644
--- a/src/storage_stage.rs
+++ b/src/storage_stage.rs
@@ -5,7 +5,7 @@
 use crate::blocktree::Blocktree;
 #[cfg(all(feature = "chacha", feature = "cuda"))]
 use crate::chacha_cuda::chacha_cbc_encrypt_file_many_keys;
-use crate::client::mk_client;
+use crate::client::mk_client_with_timeout;
 use crate::cluster_info::ClusterInfo;
 use crate::entry::EntryReceiver;
 use crate::result::{Error, Result};
@@ -227,7 +227,7 @@ impl StorageStage {
         account_to_create: Option<Pubkey>,
     ) -> io::Result<()> {
         if let Some(leader_info) = cluster_info.read().unwrap().leader_data() {
-            let mut client = mk_client(leader_info);
+            let mut client = mk_client_with_timeout(leader_info, Duration::from_secs(5));
 
             if let Some(account) = account_to_create {
                 if client.get_account_userdata(&account).is_ok() {
@@ -235,7 +235,19 @@ impl StorageStage {
                 }
             }
 
-            if let Some(last_id) = client.try_get_last_id(10) {
+            let mut last_id = None;
+            for _ in 0..10 {
+                if let Some(new_last_id) = client.try_get_last_id(1) {
+                    last_id = Some(new_last_id);
+                    break;
+                }
+
+                if exit.load(Ordering::Relaxed) {
+                    Err(io::Error::new(io::ErrorKind::Other, "exit signaled"))?;
+                }
+            }
+
+            if let Some(last_id) = last_id {
                 tx.sign(&[keypair.as_ref()], last_id);
 
                 if exit.load(Ordering::Relaxed) {